#!/usr/bin/env python

import time
starttime = int(time.time())

import asyncio
import websockets
import websockets.exceptions
import os
import re
import argparse
import sys
from functools import partial
sys.path.insert(0, '.')
import json
import random
import string
import socket

parser = argparse.ArgumentParser(description='Server')
parser.add_argument('--stream', help='Stream mode')
parser.add_argument('--port', help='Port number')
parser.add_argument('--model', help='Model name')
parser.add_argument('--home', help='Home directory')
parser.add_argument('--norand', help='No random key')
parser.add_argument('--platform', help='The CLI tool name, claude, opencode, codex')

args = parser.parse_args()

#host = "127.0.0.1"
host = "0.0.0.0"
port = 11998
if args.port:
    port = int(args.port);

if args.norand:
    rand_str = ""
else:
    rand_str = ''.join(random.choices(string.ascii_letters + string.digits, k=8))



needpipe = 0
stream = 1
if args.stream:
    stream = int(args.stream)
    print(f"Stream mode set to {stream}")
print("Setup DB ...")
cli_platform = ""
if args.platform:
    cli_platform = args.platform
    para = {}
    para['sel_vend'] = 1
    para['sel_model'] = 1
    para['stream'] = 1
    print(f"CLI Platform is {cli_platform}")
else:
    print("ERROR! --platform doesn't present\n")
    exit(1)
    #import nand
    #para = nand.setupDB(domain, hasLLM, stream, needpipe, local)
    #nand.nprint("Ready to receive from chat client")

model = ""
if args.model:
    model = args.model
    print(f"LLM model is {model}")


endtime = int(time.time())
delta = endtime - starttime
print(f"Total {delta} seconds to load")
async def handle_client(websocket, cli_platform=""):
  
    addr = websocket.remote_address #writer.get_extra_info("peername")
    print(f"Connection from : {addr}")
    #idata = await reader.read(4096)
    ddata = "" #idata.decode('utf8')
    
    idata = "" #b''
    while 1:
        #ddata and not re.search("___ENDQQ___", ddata):
        #    print(f"Received data without ___ENDQQ___ {ddata}")
        try:
            idata += await websocket.recv() #reader.read(8896)
        except websockets.exceptions.ConnectionClosedOK:
            print("Client closed unexcepted")
            return
        ddata = idata  #.decode('utf8')
        #ddata += sdata
        if re.search("___ENDQQ___", ddata):
            ddata = re.sub("___ENDQQ___", "", ddata)
            break
    print(f"Reveived {ddata}")
    if re.search(r'^(\w\w\w\w\w\w\w\w)', ddata):
        recv_key = re.search(r'^(\w\w\w\w\w\w\w\w)', ddata).group(1)
        if recv_key != rand_str:
            print(f"Rand key mismatch,  {recv_key} != {rand_str} close connection")
            return "KEY ERROR"
        ddata = re.sub(r'^(\w\w\w\w\w\w\w\w)', '', ddata)
    else:
        print(f"No key found in the frame")
        return "KEY NOT FOUND"

    if cli_platform:
        response = await process_query_simp(ddata, para, websocket, cli_platform)



async def process_query_simp(iniquery, para, websocket, cli_platform):
    print(f"iniquery={iniquery}")
    qlen =  len(iniquery)
    print(f"The query has length = {qlen} characters")
    if re.search(r"^\s*LLM\w+: ", iniquery):
         para['sel_vend'] = int(re.search(r"^\s*LLM_(\d+)_\d+_\d+: ", iniquery).group(1))
         para['sel_model'] = int(re.search(r"^\s*LLM_\d+_(\d+)_\d+: ", iniquery).group(1))
         para['stream'] = int(re.search(r"^\s*LLM_\d+_\d+_(\d+): ", iniquery).group(1))
         print(f"Selected vendor {para['sel_vend']} model {para['sel_model']} stream={para['stream']}")
    query = re.sub(r"^\s*LLM\w+: ", "", iniquery)
   
    fullm = ""
    messages = json.loads(query)
    sys_it = messages.pop(0)
    sys_info = sys_it['content']
    print(f"================= System ===================\n{sys_info}")
    fullm += sys_info
    for onem in messages:
        role = onem['role']
        print(f"================= {role} ====================")
        fullm += f"[{role}]: "
        cname = 'content'
        if 'parts' in onem:
            cname = 'parts'
        for item in onem[cname]:
            cont = item['text']
            print(f"{cont}")
            fullm += f"{cont}"
    if cli_platform == "claude":
        cmds = ["claude", '-p', '--dangerously-skip-permissions'] #, '--model', 'claude-sonnet-4-6']
    elif cli_platform == "codex":
        cmds = ["codex", 'exec']
    elif cli_platform == "opencode":
        cmds = ["opencode", 'run', '--dangerously-skip-permissions'] #, '--model', 'zhipuai-coding-plan/glm-5.2']
    if model:
        cmds.append('--model')
        cmds.append(model)
    env = os.environ.copy()
    proc = await asyncio.create_subprocess_exec(
        *cmds,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env
    )
    try:
        stdout_bytes,stderr_bytes = await asyncio.wait_for(
            proc.communicate(input=fullm.encode()), timeout=2000)
    except asyncio.TimeoutError:
        proc.kill()
        await proc.communicate()
        raise TimeoutError("subprocess timed out")
    stdout = stdout_bytes.decode()
    stderr = stderr_bytes.decode()
    if cli_platform == "opencode":
        stdout = re.sub(r'^\> build .*', '', stdout, re.MULTILINE) 
    print(f"CLI: {stdout}\n{stderr}")
    await websocket.send(f"{stdout}\n{stderr}" + "\n___ENDAA___\n")
    print(f"Sent at once")
    return stdout

        
    

async def main():
    max_retries = 10
    base_port = port
    for attempt in range(max_retries):
        try:
            current_port = base_port + attempt
            print(f"Trying to start server on {host}:{current_port}")
            async with websockets.serve(partial(handle_client, cli_platform=cli_platform), host, current_port, ping_interval=None):
                # Save the current port to a file after successful bind
               
                homedir = os.path.expanduser("~/.gofai")
                config_dir = homedir
                os.makedirs(config_dir, exist_ok=True)
                config_path = os.path.join(config_dir, "gofai.json")
               
                hostname = socket.gethostname()
                local_ip = socket.gethostbyname(hostname)
                print(f"Server hostname {hostname} local IP address: {local_ip}")
    
                config_data = {
                    "HOST": local_ip,
                    "PORT": current_port,
                    "RANDKEY": rand_str
                }
                with open(config_path, "w") as f:
                    json.dump(config_data, f, indent=2)
                    print(f"Server hostname {hostname} local IP address: {local_ip} randkey {rand_str} saved to {config_path}")
                try:
                    await asyncio.Future()
                except KeyboardInterrupt:
                    print("Ctrl-C pressed. Shutting down server.")
                    exit(0)
            break
        except OSError as e:
            if getattr(e, 'errno', None) == 10048 or 'address already in use' in str(e).lower():
                print(f"Port {current_port} in use, trying next port...")
                continue
            else:
                print(f"Failed to start server: {e}")
                raise
    else:
        print(f"Could not bind to any port in range {base_port} to {base_port + max_retries - 1}")
        exit(1)

asyncio.run(main())

