#!/usr/bin/env python3
import sys
import os
import signal
import json
import asyncio
import psutil
import urllib.request
import urllib.parse
from google.antigravity import Agent, LocalAgentConfig, CapabilitiesConfig
from google.antigravity.utils.interactive import run_interactive_loop


def kill_existing_instances():
    current_pid = os.getpid()
    current_script = os.path.basename(__file__) if "__file__" in globals() else "tahcia-antigravity"
    for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
        try:
            if proc.info['pid'] == current_pid:
                continue
            cmdline = proc.info['cmdline']
            if cmdline and any("tahcia-antigravity" in str(arg) for arg in cmdline):
                os.kill(proc.info['pid'], signal.SIGTERM)
        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
            pass

def send_antigravity_message(api_url, session_key, thread_hash, agent_hash, content):
    try:
        url = f"{api_url}/antigravity/{agent_hash}/message/{thread_hash}"
        headers = {
            "Authorization": f"Bearer {session_key}",
            "Content-Type": "application/json",
            "User-Agent": "TahciaCLI/1.0"
        }
        data = json.dumps({"content": content}).encode('utf-8')
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        with urllib.request.urlopen(req) as resp:
            return resp.read().decode('utf-8')
    except Exception as e:
        sys.stderr.write(f"Failed to send message: {e}\n")
        return None

def fetch_backend_config(api_key):
    if not api_key:
        return {}
    try:
        url = "https://api.tahcia.com/antigravity/config"
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": "TahciaCLI/1.0"
        }
        req = urllib.request.Request(url, headers=headers, method="GET")
        with urllib.request.urlopen(req, timeout=5) as resp:
            return json.loads(resp.read().decode('utf-8'))
    except Exception as e:
        sys.stderr.write(f"Notice: Could not load backend config: {e}\n")
        return {}

async def main():
    kill_existing_instances()

    api_key = ""
    port = "6999"
    workspace_arg = ""
    custom_system_prompt = ""
    python_exe = sys.executable
    is_interactive = False
    auto_approve_val = True
    confirm_run_command_val = False

    args = sys.argv[1:]
    i = 0
    positionals = []
    while i < len(args):
        arg = args[i]
        if arg in ('-k', '--apiKey', '-apiKey') and i + 1 < len(args):
            api_key = args[i + 1]
            i += 2
        elif arg in ('-port', '--port') and i + 1 < len(args):
            port = args[i + 1]
            i += 2
        elif arg in ('-p', '-paths', '--paths') and i + 1 < len(args):
            workspace_arg = args[i + 1]
            i += 2
        elif arg in ('-py', '--python', '-python') and i + 1 < len(args):
            python_exe = args[i + 1]
            i += 2
        elif arg in ('-i', '--interactive'):
            is_interactive = True
            i += 1
        else:
            positionals.append(arg)
            i += 1

    if not api_key and len(positionals) > 0:
        api_key = positionals[0]

    # Load configuration from backend if available
    backend_config = fetch_backend_config(api_key)
    if backend_config:
        if backend_config.get('port') and port == "6999":
            port = str(backend_config['port'])
        if backend_config.get('filepath'):
            fp = backend_config['filepath']
            if isinstance(fp, list):
                workspace_folders = fp
            elif isinstance(fp, str):
                workspace_arg = workspace_arg + "," + fp if workspace_arg else fp
        if 'auto_approve' in backend_config:
            auto_approve_val = bool(backend_config['auto_approve'])
        if 'confirm_run_command' in backend_config:
            confirm_run_command_val = bool(backend_config['confirm_run_command'])


    if len(positionals) > 1 and port == "6999" and positionals[1].isdigit():
        port = positionals[1]
    if len(positionals) > 2 and not workspace_arg:
        workspace_arg = positionals[2]
    if len(positionals) > 3 and not custom_system_prompt:
        custom_system_prompt = positionals[3]

    if not 'workspace_folders' in locals() or not workspace_folders:
        workspace_folders = []
    if workspace_arg:
        try:
            cli_folders = json.loads(workspace_arg)
        except Exception:
            cli_folders = [f.strip() for f in workspace_arg.split(',') if f.strip()]
        for f in cli_folders:
            if f not in workspace_folders:
                workspace_folders.append(f)

    print(f"✓ Starting Tahcia Antigravity Agent (port: {port})")
    if api_key:
        print(f"  API Key: {api_key[:6]}...")
    if workspace_folders:
        print(f"  Workspace Folders: {workspace_folders}")

    from google.antigravity.hooks import policy

    agent_policies = []
    if confirm_run_command_val:
        agent_policies.append(policy.confirm_run_command())
    else:
        agent_policies.append(policy.allow_all())

    config = LocalAgentConfig(
        system_instructions=custom_system_prompt,
        capabilities=CapabilitiesConfig(),
        policies=agent_policies,
        workspaces=workspace_folders
    )

    async with Agent(config) as agent:
        if is_interactive:
            print(f"✓ Tahcia Antigravity Agent initialized (Interactive Mode). Type 'exit' or 'quit' to end:")
            interactive_arg = getattr(agent, 'config', config)
            await run_interactive_loop(interactive_arg)
        else:
            print(f"✓ Tahcia Antigravity Agent daemon running on 0.0.0.0:{port}")
            async def run_task(prompt, agent_h, thread_h):
                print(f"→ Processing task for thread {thread_h}: {prompt[:60]}...")
                try:
                    res = await execute_antigravity_agent(agent, prompt)
                    output = str(res)
                except Exception as e:
                    output = f"Error executing task: {e}"
                print(f"← Replying back to Tahcia...")
                send_antigravity_message("https://api.tahcia.com", api_key, thread_h, agent_h, output)

            async def _handle_request(reader, writer):
                data = await reader.read(65536)
                req_str = data.decode('utf-8', errors='ignore')
                body = req_str.split("\r\n\r\n", 1)[1] if "\r\n\r\n" in req_str else ""
                try:
                    payload = json.loads(body) if body else {}
                except Exception:
                    payload = {}

                # Extract prompt from OpenAI messages format or simple json
                prompt = ""
                messages = payload.get('messages', [])
                if messages and isinstance(messages, list):
                    prompt_parts = []
                    for m in messages:
                        if isinstance(m, dict):
                            role = m.get('role', '')
                            content = m.get('content', '')
                            if isinstance(content, str) and content:
                                prompt_parts.append(f"[{role}]: {content}")
                            elif isinstance(content, list):
                                for c in content:
                                    if isinstance(c, dict) and c.get('type') == 'text':
                                        prompt_parts.append(f"[{role}]: {c.get('text')}")
                    prompt = "\n".join(prompt_parts)
                
                if not prompt:
                    prompt = payload.get('prompt', payload.get('content', payload.get('message', '')))

                thread_hash = payload.get('user', '').replace('thread_', '') or payload.get('thread_hash', payload.get('thash', ''))
                agent_hash = payload.get('agent_hash', payload.get('hash', ''))
                is_stream = payload.get('stream', False) or '/v1/chat/completions' in req_str

                if is_stream:
                    http_header = (
                        "HTTP/1.1 200 OK\r\n"
                        "Content-Type: text/event-stream\r\n"
                        "Cache-Control: no-cache\r\n"
                        "Connection: keep-alive\r\n"
                        "Access-Control-Allow-Origin: *\r\n\r\n"
                    ).encode('utf-8')
                    writer.write(http_header)
                    await writer.drain()

                    print(f"→ Processing streaming thread chat: {prompt[:60]}...")
                    async def keepalive():
                        try:
                            while True:
                                await asyncio.sleep(15)
                                ping_chunk = "data: " + json.dumps({"choices": [{"delta": {"content": ""}}]}) + "\n\n"
                                writer.write(ping_chunk.encode('utf-8'))
                                await writer.drain()
                        except asyncio.CancelledError:
                            pass
                        except Exception:
                            pass

                    ka_task = asyncio.create_task(keepalive())
                    try:
                        if hasattr(agent, 'chat'):
                            res = await agent.chat(prompt)
                            if hasattr(res, '__aiter__'):
                                async for chunk in res:
                                    sse_chunk = "data: " + json.dumps({"choices": [{"delta": {"content": str(chunk)}}]}) + "\n\n"
                                    writer.write(sse_chunk.encode('utf-8'))
                                    await writer.drain()
                            else:
                                output = str(res)
                                sse_chunk = "data: " + json.dumps({"choices": [{"delta": {"content": output}}]}) + "\n\n"
                                writer.write(sse_chunk.encode('utf-8'))
                                await writer.drain()
                        else:
                            res = await execute_antigravity_agent(agent, prompt)
                            output = str(res)
                            sse_chunk = "data: " + json.dumps({"choices": [{"delta": {"content": output}}]}) + "\n\n"
                            writer.write(sse_chunk.encode('utf-8'))
                            await writer.drain()
                    except Exception as e:
                        output = f"Error executing task: {e}"
                        sse_chunk = "data: " + json.dumps({"choices": [{"delta": {"content": output}}]}) + "\n\n"
                        writer.write(sse_chunk.encode('utf-8'))
                        await writer.drain()
                    finally:
                        ka_task.cancel()

                    sse_done = "data: [DONE]\n\n"
                    writer.write(sse_done.encode('utf-8'))
                    await writer.drain()
                    writer.close()
                    await writer.wait_closed()
                    print("← Streamed response to Thread Chat.")
                else:
                    resp_json = json.dumps({"status": "received", "agent": "antigravity"}).encode('utf-8')
                    http_header = (
                        f"HTTP/1.1 200 OK\r\n"
                        f"Content-Type: application/json\r\n"
                        f"Content-Length: {len(resp_json)}\r\n"
                        f"Access-Control-Allow-Origin: *\r\n"
                        f"Connection: close\r\n\r\n"
                    ).encode('utf-8')
                    writer.write(http_header + resp_json)
                    await writer.drain()
                    writer.close()
                    await writer.wait_closed()

                    if prompt and thread_hash:
                        asyncio.create_task(run_task(prompt, agent_hash, thread_hash))

            async def handle_request(reader, writer):
                try:
                    await _handle_request(reader, writer)
                except (ConnectionError, BrokenPipeError):
                    pass
                except Exception as e:
                    print(f"Error handling request: {e}")
                finally:
                    if not writer.is_closing():
                        try:
                            writer.close()
                            await writer.wait_closed()
                        except Exception:
                            pass

            server = await asyncio.start_server(handle_request, '0.0.0.0', int(port))
            async with server:
                await server.serve_forever()

            
async def format_agent_output(res):
    if isinstance(res, str):
        return res

    if hasattr(res, '__aiter__'):
        try:
            parts = []
            async for chunk in res:
                parts.append(str(chunk))
            if parts:
                return "".join(parts)
        except Exception as e:
            return f"Error streaming ChatResponse: {e}"

    for attr in ['text', 'content', 'message', 'output', 'response']:
        if hasattr(res, attr):
            val = getattr(res, attr)
            if val is not None:
                if callable(val):
                    try:
                        val = val()
                    except Exception as e:
                        pass
                if asyncio.iscoroutine(val) or hasattr(val, '__await__'):
                    try:
                        val = await val
                    except Exception as e:
                        return f"Error awaiting {attr}(): {e}"
                if isinstance(val, str) and val:
                    return val
                if isinstance(val, list) and val:
                    parts = []
                    for item in val:
                        if isinstance(item, dict):
                            parts.append(str(item.get('text', item)))
                        else:
                            parts.append(str(item))
                    return "".join(parts)
                if val and not callable(val):
                    return str(val)

    if hasattr(res, '__dict__'):
        d = getattr(res, '__dict__')
        for k in ['text', 'content', 'message', 'output', 'response']:
            if k in d and d[k]:
                val = d[k]
                if callable(val):
                    try:
                        val = val()
                    except Exception:
                        pass
                if asyncio.iscoroutine(val) or hasattr(val, '__await__'):
                    try:
                        val = await val
                    except Exception:
                        pass
                if val and not callable(val):
                    return str(val)
                    
    return str(res)

async def execute_antigravity_agent(agent, prompt):
    methods_to_try = [
        'run', 'generate', 'send_message', 'send', 'query',
        'execute', 'chat', 'run_task', 'generate_response',
        'run_turn', 'start_turn', 'stream'
    ]
    for method_name in methods_to_try:
        if hasattr(agent, method_name):
            func = getattr(agent, method_name)
            if callable(func):
                try:
                    res = func(prompt)
                    if asyncio.iscoroutine(res) or hasattr(res, '__await__'):
                        res = await res
                    return await format_agent_output(res)
                except Exception as e:
                    sys.stderr.write(f"Method {method_name} failed: {e}\n")

    for attr in ['session', 'runner', 'client', 'chat']:
        if hasattr(agent, attr):
            target = getattr(agent, attr)
            for m in ['send_message', 'send', 'query', 'run', 'generate']:
                if hasattr(target, m):
                    func = getattr(target, m)
                    if callable(func):
                        try:
                            res = func(prompt)
                            if asyncio.iscoroutine(res) or hasattr(res, '__await__'):
                                res = await res
                            return await format_agent_output(res)
                        except Exception:
                            pass

    for attr_name in dir(agent):
        if not attr_name.startswith('_'):
            val = getattr(agent, attr_name)
            if callable(val) and attr_name not in ('close', 'start', 'stop', 'cleanup'):
                try:
                    res = val(prompt)
                    if asyncio.iscoroutine(res) or hasattr(res, '__await__'):
                        res = await res
                    return await format_agent_output(res)
                except Exception:
                    pass

    raise RuntimeError(f"Could not find callable prompt execution method on Agent object: {dir(agent)}")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nExiting Tahcia Antigravity Agent.")
        sys.exit(0)