#!/usr/bin/env python3
"""
Lightweight local HTTP bridge for AGY (Antigravity CLI).
Exposes an OpenAI-compatible /v1/chat/completions endpoint on port 8088
so Docker containers can seamlessly send AI requests to the host's agy CLI.
"""

import sys
import os
import json
import shutil
import logging
import subprocess
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] agy_bridge: %(message)s")
logger = logging.getLogger("agy_bridge")

AGY_BIN = shutil.which("agy") or ("/home/mamad/.local/bin/agy" if os.path.exists("/home/mamad/.local/bin/agy") else "agy")
PORT = int(os.getenv("AGY_BRIDGE_PORT", "8088"))
HOST = "0.0.0.0"


class AGYBridgeHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path in ("/", "/health", "/v1/models"):
            resp = json.dumps({"status": "ok", "service": "agy_bridge", "agy_bin": AGY_BIN}).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp)))
            self.end_headers()
            self.wfile.write(resp)
        else:
            self.send_error(404, "Not Found")

    def do_POST(self):
        if not self.path.startswith("/v1/chat/completions") and self.path != "/chat/completions":
            self.send_error(404, "Not Found")
            return

        try:
            content_length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(content_length).decode("utf-8")
            data = json.loads(body)

            messages = data.get("messages", [])
            effort = data.get("reasoning_effort", "")

            system_instructions = []
            user_prompts = []
            temp_image_files = []

            for msg in messages:
                role = msg.get("role", "user")
                content = msg.get("content", "")
                if isinstance(content, list):
                    for part in content:
                        if isinstance(part, dict):
                            p_type = part.get("type", "")
                            if p_type == "text":
                                txt = part.get("text", "")
                                if role == "system":
                                    system_instructions.append(txt)
                                else:
                                    user_prompts.append(txt)
                            elif p_type == "image_url":
                                url_data = part.get("image_url", {}).get("url", "")
                                if url_data.startswith("data:image/"):
                                    import tempfile, base64
                                    header, b64_str = url_data.split(",", 1)
                                    ext = ".png" if "png" in header else ".jpg"
                                    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp_f:
                                        tmp_f.write(base64.b64decode(b64_str))
                                        temp_image_files.append(tmp_f.name)
                elif isinstance(content, str):
                    if role == "system":
                        system_instructions.append(content)
                    else:
                        user_prompts.append(content)

            full_prompt = ""
            if system_instructions:
                full_prompt += f"System Instruction:\n" + "\n".join(system_instructions) + "\n\n"
            full_prompt += "User Prompt:\n" + "\n".join(user_prompts)

            if temp_image_files:
                full_prompt += "\n\nAttached Images to analyze:\n" + "\n".join([f"- {f}" for f in temp_image_files])

            cmd = [AGY_BIN, "-p", full_prompt, "--output-format", "json"]
            if effort in ("low", "medium", "high"):
                cmd.extend(["--effort", effort])

            logger.info(f"Executing agy for prompt length: {len(full_prompt)} chars (images: {len(temp_image_files)})")
            try:
                proc = subprocess.run(
                    cmd,
                    capture_output=True,
                    text=True,
                    timeout=120,
                    env=os.environ.copy()
                )
            finally:
                for img_f in temp_image_files:
                    if os.path.exists(img_f):
                        try:
                            os.remove(img_f)
                        except Exception:
                            pass

            if proc.returncode != 0:
                logger.error(f"agy error (code {proc.returncode}): {proc.stderr}")
                self.send_error(500, f"agy execution error: {proc.stderr}")
                return

            try:
                res_data = json.loads(proc.stdout)
                content = res_data.get("response", "")
            except Exception:
                content = proc.stdout

            resp_payload = {
                "id": "agy-response",
                "object": "chat.completion",
                "created": 1234567890,
                "model": data.get("model", "antigravity"),
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": content
                        },
                        "finish_reason": "stop"
                    }
                ]
            }
            resp_bytes = json.dumps(resp_payload, ensure_ascii=False).encode("utf-8")

            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(resp_bytes)))
            self.end_headers()
            self.wfile.write(resp_bytes)
            logger.info("Successfully handled AGY chat completion request.")

        except Exception as e:
            logger.error(f"Failed to handle request: {e}", exc_info=True)
            self.send_error(500, str(e))

    def log_message(self, format, *args):
        # Override default noisy stderr logging
        pass


def main():
    logger.info(f"Starting AGY Bridge Server on http://{HOST}:{PORT} (AGY: {AGY_BIN})")
    server = ThreadingHTTPServer((HOST, PORT), AGYBridgeHandler)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down AGY Bridge Server...")
        server.shutdown()


if __name__ == "__main__":
    main()