atomcode_proxy.py

#!/usr/bin/env python3
"""
AtomCode Reverse Proxy — OpenAI-compatible proxy via AtomCode Daemon

Architecture:
  ┌──────────────┐     OpenAI API calls     ┌──────────────────┐    POST /chat    ┌────────────┐
  │  Any AI Tool  │ ───────────────────────▶ │  Proxy Server    │ ──────────────▶ │  Daemon    │
  │  (Cline/Codex)│ ◀─────────────────────── │  (:8080/v1)      │ ◀────────────── │  :13456    │
  └──────────────┘                          └──────────────────┘                 └────────────┘

Usage:
  1. Start the proxy (auto-manages daemon):
     python3 atomcode_proxy.py start

  2. Or manually:
     ATOMCODE_TELEMETRY=0 atomcode daemon
     python3 atomcode_proxy.py serve --port 8080

  3. Configure your AI tool:
     API Base URL: http://localhost:8080/v1
     API Key: anything (not validated)
     Model: deepseek-v4-flash
"""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import json
import logging
import os
import pathlib
import platform
import re
import shutil
import signal
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, AsyncGenerator, Dict, List, Optional

import yaml
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

try:
    import aiohttp
except ImportError:
    aiohttp = None  # type: ignore

# ─── Logging ──────────────────────────────────────────────────────────────────

logger = logging.getLogger("atomcode-proxy")

# ─── Configuration ─────────────────────────────────────────────────────────────

DEFAULT_CONFIG = {
    "server": {"host": "0.0.0.0", "port": 8080},
    "daemon": {
        "url": os.environ.get("ATOMCODE_DAEMON_URL", "http://localhost:13456"),
        "binary_path": None,  # auto-detect
        "auto_manage": False,  # proxy manages daemon lifecycle
        "telemetry": False,  # ATOMCODE_TELEMETRY=0
    },
    "logging": {"level": "INFO", "file": None},
    "sessions": {"max_age_minutes": 30},
    "setup": {
        "auto_login": True,  # attempt OAuth login if not logged in
        "auto_claim": True,  # attempt CodingPlan claim if needed
    },
}

CONFIG_SEARCH_PATHS = [
    pathlib.Path("proxy.yaml"),
    pathlib.Path("~/.atomcode/proxy.yaml").expanduser(),
    pathlib.Path("/etc/atomcode/proxy.yaml"),
]


def find_atomcode_binary() -> str:
    candidates = [
        shutil.which("atomcode"),
        "/usr/local/bin/atomcode",
        os.path.expanduser("~/.local/bin/atomcode"),
    ]
    for c in candidates:
        if c:
            return c
    raise FileNotFoundError(
        "atomcode binary not found. Install it with:\n"
        "  curl -fsSL https://raw.atomgit.com/atomgit_atomcode/atomcode/raw/main/scripts/install.sh | sh"
    )


def load_config() -> dict:
    config = dict(DEFAULT_CONFIG)

    # Load from first existing search path
    for p in CONFIG_SEARCH_PATHS:
        resolved = p.expanduser().resolve() if hasattr(p, "expanduser") else pathlib.Path(p).expanduser().resolve()
        if resolved.exists():
            with open(resolved) as f:
                loaded = yaml.safe_load(f) or {}
                _deep_merge(config, loaded)
            logger.info("Loaded config from %s", resolved)
            break

    # Auto-detect binary path
    if config["daemon"]["binary_path"] is None:
        try:
            config["daemon"]["binary_path"] = find_atomcode_binary()
        except FileNotFoundError:
            config["daemon"]["binary_path"] = "atomcode"

    # Env overrides
    config["daemon"]["url"] = os.environ.get("ATOMCODE_DAEMON_URL", config["daemon"]["url"])
    if "ATOMCODE_TELEMETRY" in os.environ:
        config["daemon"]["telemetry"] = os.environ["ATOMCODE_TELEMETRY"] == "0"

    return config


def _deep_merge(base: dict, overlay: dict) -> None:
    """Recursively merge overlay into base dict."""
    for k, v in overlay.items():
        if k in base and isinstance(base[k], dict) and isinstance(v, dict):
            _deep_merge(base[k], v)
        else:
            base[k] = v


@dataclass
class AppState:
    config: dict = field(default_factory=load_config)
    daemon_proc: Optional[subprocess.Popen] = None
    daemon_running: bool = False
    daemon_fail_count: int = 0
    start_time: float = field(default_factory=time.time)
    sessions: Dict[str, str] = field(default_factory=dict)  # conv_hash → daemon_session_id
    session_times: Dict[str, float] = field(default_factory=dict)  # conv_hash → last used timestamp


# ─── Pydantic Models (OpenAI-Compatible) ──────────────────────────────────────


class Message(BaseModel):
    role: str = "user"
    content: Any = ""  # str or List[dict] (multi-part)


class ChatCompletionRequest(BaseModel):
    model: Optional[str] = None
    messages: List[Message] = Field(default_factory=list)
    stream: bool = False
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    top_p: Optional[float] = None
    frequency_penalty: Optional[float] = None
    presence_penalty: Optional[float] = None
    stop: Any = None
    user: Optional[str] = None


class ChatUsage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatChoice(BaseModel):
    index: int = 0
    message: dict = Field(default_factory=lambda: {"role": "assistant", "content": ""})
    finish_reason: Optional[str] = None


class ChatCompletionResponse(BaseModel):
    id: str = "chatcmpl-atomcode"
    object: str = "chat.completion"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str = "unknown"
    choices: List[ChatChoice] = Field(default_factory=list)
    usage: Optional[ChatUsage] = None


# ─── SSE Translation ──────────────────────────────────────────────────────────


def _translate_daemon_event(event: dict, model: str, tool_call_index: int) -> Optional[dict]:
    """
    Translate a daemon SSE event dict into an OpenAI SSE chunk dict.
    Returns None if the event should be skipped (no translation).
    Tracks tool_call_index to assign sequential tool call IDs.
    """
    etype = event.get("type")

    if etype == "text":
        return {
            "choices": [{"delta": {"content": event.get("content", "")}, "index": 0}],
        }

    elif etype == "reasoning":
        return {
            "choices": [{"delta": {"reasoning_content": event.get("content", "")}, "index": 0}],
        }

    elif etype == "tool_start":
        tid = event.get("id", f"call_{tool_call_index}")
        name = event.get("name", "")
        args_text = event.get("arguments", "{}")
        return {
            "choices": [{
                "delta": {
                    "tool_calls": [{
                        "index": 0,
                        "id": tid,
                        "type": "function",
                        "function": {"name": name, "arguments": args_text},
                    }]
                },
                "index": 0,
            }],
        }

    elif etype == "tool_output":
        # Intermediate progress — skip in OpenAI format
        return None

    elif etype == "tool_result":
        # Final tool result — skip in streaming (folded into next assistant response)
        return None

    elif etype == "tokens":
        pt = event.get("prompt", 0)
        ct = event.get("completion", 0)
        tt = event.get("total", 0)
        return {
            "choices": [],
            "usage": {"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": tt},
        }

    elif etype == "done":
        return {"type": "done", "session_id": event.get("session_id")}

    elif etype == "error":
        return {
            "choices": [{"delta": {}, "finish_reason": "error", "index": 0}],
        }

    elif etype == "stopped":
        return {"type": "done"}

    return None


def _build_openai_chunk(base: dict, model: str | None) -> dict:
    """Wrap a delta dict into a full OpenAI SSE chunk."""
    chunk = {
        "id": "chatcmpl-atomcode",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model or "unknown",
    }
    chunk.update(base)
    return chunk


# ─── Daemon Client ────────────────────────────────────────────────────────────


class DaemonClient:
    """Async HTTP client for the AtomCode daemon REST API."""

    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self._http_timeout = 120

    # ── Sync helpers (for simple request/response) ──

    def _sync_request(self, method: str, path: str, body: Any = None) -> tuple[int, Any]:
        url = f"{self.base_url}{path}"
        data = json.dumps(body).encode() if body is not None else None
        req = urllib.request.Request(url, data=data, method=method)
        req.add_header("Content-Type", "application/json")
        req.add_header("Accept", "application/json")
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                return resp.status, json.loads(resp.read())
        except urllib.error.HTTPError as e:
            body_text = e.read().decode()
            try:
                return e.code, json.loads(body_text)
            except json.JSONDecodeError:
                return e.code, {"error": body_text}
        except urllib.error.URLError as e:
            return 503, {"error": f"Daemon unreachable: {e.reason}"}

    def health(self) -> dict:
        _, data = self._sync_request("GET", "/health")
        return data

    def auth_status(self) -> dict:
        _, data = self._sync_request("GET", "/auth/status")
        return data

    def login_start(self) -> dict:
        _, data = self._sync_request("POST", "/auth/login/start", {"open_browser": False})
        return data

    def login_poll(self, login_id: str) -> dict:
        _, data = self._sync_request("POST", f"/auth/login/{login_id}/poll", {})
        return data

    def codingplan_setup(self) -> dict:
        _, data = self._sync_request("POST", "/codingplan/setup", {})
        return data

    def models_list(self) -> list:
        _, data = self._sync_request("GET", "/models")
        return data if isinstance(data, list) else []

    def providers_list(self) -> dict:
        _, data = self._sync_request("GET", "/providers")
        return data if isinstance(data, dict) else {}

    # ── Async SSE chat stream ──

    async def chat_stream(self, message: str, stream: bool = True,
                          provider: str | None = None, system: str | None = None,
                          session_id: str | None = None
                          ) -> AsyncGenerator[dict, None]:
        """
        POST /chat to daemon, yield parsed SSE event dicts asynchronously.
        Uses aiohttp if available, falls back to sync thread.
        Pass session_id to continue an existing conversation context.
        """
        body: dict = {"message": message, "stream": stream}
        if provider:
            body["provider"] = provider
        if system:
            body["system"] = system
        if session_id:
            body["session_id"] = session_id

        url = f"{self.base_url}/chat"
        headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}

        if aiohttp is not None:
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self._http_timeout)) as session:
                async with session.post(url, json=body, headers=headers) as resp:
                    async for raw_line in resp.content:
                        line = raw_line.decode("utf-8", errors="replace").strip()
                        if not line:
                            continue
                        if line.startswith("data: "):
                            data_str = line[6:]
                            if data_str == "[DONE]":
                                yield {"type": "done"}
                                return
                            try:
                                yield json.loads(data_str)
                            except json.JSONDecodeError:
                                pass
        else:
            # Fallback: sync request in thread pool
            loop = asyncio.get_event_loop()
            data = json.dumps(body).encode()
            req = urllib.request.Request(url, data=data, method="POST")
            for k, v in headers.items():
                req.add_header(k, v)

            def _read_sse():
                resp = urllib.request.urlopen(req, timeout=self._http_timeout)
                for raw_line in resp:
                    line = raw_line.decode("utf-8", errors="replace").strip()
                    yield line

            lines = await loop.run_in_executor(None, _read_sse)
            for line in lines:
                if not line:
                    continue
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str == "[DONE]":
                        yield {"type": "done"}
                        return
                    try:
                        yield json.loads(data_str)
                    except json.JSONDecodeError:
                        pass


# ─── Daemon Process Manager ──────────────────────────────────────────────────


class DaemonManager:
    """Manages the AtomCode daemon subprocess lifecycle."""

    def __init__(self, config: dict):
        self.config = config
        self._proc: Optional[subprocess.Popen] = None
        self._daemon_url = config["daemon"]["url"]
        self._binary = config["daemon"]["binary_path"]
        self._telemetry = config["daemon"].get("telemetry", True)

    @property
    def running(self) -> bool:
        if self._proc is None:
            return False
        poll = self._proc.poll()
        return poll is None

    def start(self) -> bool:
        if self.running:
            logger.info("Daemon already running (PID %d)", self._proc.pid)
            return True

        logger.info("Starting AtomCode daemon...")
        env = os.environ.copy()
        if not self._telemetry:
            env["ATOMCODE_TELEMETRY"] = "0"

        log_path = self.config["logging"].get("file")
        stdout = stderr = None
        if log_path:
            stdout = open(log_path, "a")
            stderr = subprocess.STDOUT

        try:
            self._proc = subprocess.Popen(
                [self._binary, "daemon"],
                env=env,
                stdout=stdout,
                stderr=stderr,
                stdin=subprocess.DEVNULL,
            )
            logger.info("Daemon started (PID %d)", self._proc.pid)

            # Wait for readiness
            for i in range(15):
                if self._check_health():
                    logger.info("Daemon ready after %ds", i + 1)
                    return True
                time.sleep(1)

            logger.error("Daemon failed to become ready within 15s")
            self.stop()
            return False

        except FileNotFoundError:
            logger.error("Daemon binary not found: %s", self._binary)
            return False

    def stop(self) -> bool:
        if self._proc is None:
            return True

        logger.info("Stopping daemon (PID %d)...", self._proc.pid)
        try:
            self._proc.terminate()
            self._proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            logger.warning("Daemon did not terminate, killing...")
            self._proc.kill()
            self._proc.wait(timeout=5)
        except Exception as e:
            logger.error("Error stopping daemon: %s", e)
            return False

        self._proc = None
        logger.info("Daemon stopped")
        return True

    def restart(self) -> bool:
        self.stop()
        time.sleep(1)
        return self.start()

    def _check_health(self) -> bool:
        try:
            data = json.loads(urllib.request.urlopen(
                urllib.request.Request(f"{self._daemon_url}/health"),
                timeout=3,
            ).read())
            return data.get("status") == "ok"
        except Exception:
            return False


# ─── Session Helper ──────────────────────────────────────────────────────────


def _conversation_key(messages: List[dict], system: str = "") -> str:
    """
    Generate a deterministic hash for a conversation prefix.
    Includes all messages except the last user message, plus system prompt.
    Used to identify the same conversation context across proxy restarts.
    """
    payload = system
    for m in messages[:-1]:
        role = m.get("role", "")
        content = str(m.get("content", ""))
        payload += f"|{role}:{content}"
    return hashlib.md5(payload.encode()).hexdigest()[:16]


def _format_messages_for_daemon(messages: List[dict], system_prompt: str) -> str:
    """
    Format the full messages history into a single message string
    for the daemon's /chat endpoint.
    """
    parts = []
    for m in messages:
        role = m.get("role", "user")
        content = m.get("content", "")
        if isinstance(content, list):
            # Multi-part content: extract text parts
            texts = [p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"]
            content = "\n".join(texts)

        if role == "system":
            continue  # system is handled separately
        if not content:
            continue

        role_label = {"user": "User", "assistant": "Assistant", "tool": "Tool"}.get(role, role.capitalize())
        parts.append(f"{role_label}: {content}")

    msg = "\n\n".join(parts) if parts else messages[-1].get("content", "")
    return msg


# ─── FastAPI App ──────────────────────────────────────────────────────────────


def create_app(config: dict) -> FastAPI:
    app = FastAPI(
        title="AtomCode Reverse Proxy",
        version="2.0.0",
        docs_url=None,  # disable docs in production
        redoc_url=None,
    )

    daemon_client = DaemonClient(config["daemon"]["url"])
    daemon_mgr = DaemonManager(config)
    app.state.config = config
    app.state.daemon_client = daemon_client
    app.state.daemon_mgr = daemon_mgr
    app.state.start_time = time.time()
    app.state.sessions: Dict[str, str] = {}
    app.state.session_times: Dict[str, float] = {}
    app.state.daemon_fail_count = 0

    # CORS
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # ── Events ──

    @app.on_event("startup")
    async def startup():
        # Start daemon if auto_manage is enabled
        if config["daemon"].get("auto_manage", False):
            daemon_mgr.start()

        # Start health monitor background task
        asyncio.create_task(_health_monitor())

    @app.on_event("shutdown")
    async def shutdown():
        if config["daemon"].get("auto_manage", False):
            daemon_mgr.stop()

    async def _health_monitor():
        """Periodically check daemon health and attempt recovery."""
        while True:
            await asyncio.sleep(15)
            try:
                h = daemon_client.health()
                if h.get("status") == "ok":
                    app.state.daemon_fail_count = 0
                    app.state.daemon_running = True
                else:
                    app.state.daemon_fail_count += 1
            except Exception:
                app.state.daemon_fail_count += 1
                app.state.daemon_running = False

            # Auto-recovery after 3 consecutive failures
            if app.state.daemon_fail_count >= 3 and config["daemon"].get("auto_manage", False):
                logger.warning("Daemon unreachable, attempting restart...")
                daemon_mgr.restart()
                app.state.daemon_fail_count = 0

    # ── Helpers ──

    async def _check_daemon_ready() -> bool:
        """Verify daemon is running and logged in. Returns True if OK."""
        try:
            h = daemon_client.health()
            if h.get("status") != "ok":
                raise HTTPException(status_code=502, detail={
                    "error": {"message": f"Daemon not ready: {h}", "code": 502}
                })
            app.state.daemon_running = True
            app.state.daemon_fail_count = 0

            auth = daemon_client.auth_status()
            if not auth.get("logged_in"):
                raise HTTPException(status_code=502, detail={
                    "error": {
                        "message": "Daemon running but not logged in. Run: python3 atomcode_proxy.py login",
                        "code": 502,
                    }
                })
            return True
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=502, detail={
                "error": {"message": f"Daemon unreachable at {config['daemon']['url']}: {e}", "code": 502}
            })

    # ── Endpoints ──

    @app.options("/v1/{path:path}")
    async def options_handler(path: str):
        return JSONResponse(status_code=204, headers={
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
            "Access-Control-Allow-Headers": "Content-Type, Authorization",
        })

    @app.get("/v1/health")
    async def health():
        try:
            h = daemon_client.health()
            auth = daemon_client.auth_status()
            daemon_ok = h.get("status") == "ok"
            logged_in = auth.get("logged_in", False)
        except Exception as e:
            daemon_ok = False
            logged_in = False
            h = {"error": str(e)}
            auth = {}

        return {
            "status": "ok",
            "uptime_seconds": int(time.time() - app.state.start_time),
            "daemon": {
                "connected": daemon_ok,
                "version": h.get("version", "?"),
                "service": h.get("service", "?"),
            },
            "logged_in": logged_in,
            "user": auth.get("user"),
        }

    @app.get("/v1/models")
    async def list_models():
        try:
            models = daemon_client.models_list()
        except Exception as e:
            raise HTTPException(status_code=502, detail={
                "error": {"message": f"Failed to fetch models: {e}", "code": 502}
            })

        return {
            "object": "list",
            "data": [
                {
                    "id": m.get("model", m.get("name", "unknown")),
                    "object": "model",
                    "created": int(time.time()),
                    "owned_by": "atomcode",
                }
                for m in models
            ],
        }

    @app.post("/v1/chat/completions")
    async def chat_completions(req: ChatCompletionRequest, request: Request):
        await _check_daemon_ready()

        messages = [m.model_dump() if hasattr(m, 'model_dump') else dict(m) for m in req.messages]
        model = req.model or ""
        stream = req.stream

        # 1. Extract system prompt
        system_prompt = ""
        remaining = []
        for m in messages:
            if m.get("role") == "system":
                system_prompt = str(m.get("content", ""))
            else:
                remaining.append(m)
        messages = remaining

        if not messages:
            raise HTTPException(status_code=400, detail={
                "error": {"message": "No user/assistant messages found", "code": 400}
            })

        # 2. Determine provider from model name
        provider = None
        if model:
            try:
                providers_data = daemon_client.providers_list()
                provider_list = providers_data.get("providers", []) if isinstance(providers_data, dict) else providers_data or []
                for p in provider_list:
                    if isinstance(p, dict) and p.get("model", "").lower() == model.lower():
                        provider = p.get("name")
                        break
            except Exception:
                pass

        # 3. Format messages into daemon message string
        daemon_message = _format_messages_for_daemon(messages, system_prompt)

        # 4. Session tracking — reuse daemon session_id for multi-turn context
        conv_key = _conversation_key(messages, system_prompt)
        daemon_session_id = app.state.sessions.get(conv_key)
        now = time.time()
        app.state.session_times[conv_key] = now

        # Clean expired sessions
        max_age = config["sessions"]["max_age_minutes"] * 60
        expired_keys = [k for k, v in app.state.session_times.items() if now - v > max_age]
        for k in expired_keys:
            app.state.sessions.pop(k, None)
            app.state.session_times.pop(k, None)

        logger.info("Chat: model=%s, stream=%s, provider=%s, msgs=%d, sys=%s, sid=%s",
                     model, stream, provider, len(messages), bool(system_prompt),
                     daemon_session_id or "(new)")

        # ── Non-streaming ──
        if not stream:
            full_text = ""
            usage_data = ChatUsage()
            tool_calls: list = []
            new_session_id = daemon_session_id

            try:
                async for event in daemon_client.chat_stream(
                    message=daemon_message, stream=True, provider=provider,
                    system=system_prompt, session_id=daemon_session_id
                ):
                    etype = event.get("type")
                    if etype == "text":
                        full_text += event.get("content", "")
                    elif etype == "reasoning":
                        pass  # skip reasoning in non-streaming
                    elif etype == "tool_start":
                        tool_calls.append({
                            "id": event.get("id", ""),
                            "type": "function",
                            "function": {
                                "name": event.get("name", ""),
                                "arguments": event.get("arguments", "{}"),
                            },
                        })
                    elif etype == "tokens":
                        usage_data = ChatUsage(
                            prompt_tokens=event.get("prompt", 0),
                            completion_tokens=event.get("completion", 0),
                            total_tokens=event.get("total", 0),
                        )
                    elif etype == "done":
                        new_session_id = event.get("session_id") or new_session_id
                        break
                    elif etype == "error":
                        raise HTTPException(status_code=502, detail={
                            "error": {"message": event.get("message", "Daemon error"), "code": 502}
                        })

            except HTTPException:
                raise
            except Exception as e:
                logger.error("Chat error: %s", e, exc_info=True)
                raise HTTPException(status_code=502, detail={
                    "error": {"message": f"Daemon error: {e}", "code": 502}
                })

            message: dict = {"role": "assistant", "content": full_text}
            if tool_calls:
                message["tool_calls"] = tool_calls
                finish_reason = "tool_calls"
            else:
                finish_reason = "stop"

            # Persist session_id for multi-turn continuity
            if new_session_id and new_session_id != daemon_session_id:
                app.state.sessions[conv_key] = new_session_id

            return ChatCompletionResponse(
                model=model or "unknown",
                choices=[ChatChoice(
                    message=message,
                    finish_reason=finish_reason,
                )],
                usage=usage_data if usage_data.total_tokens > 0 else None,
            )

        # ── Streaming ──
        else:
            async def event_stream():
                tool_idx = 0
                current_session_id = daemon_session_id
                try:
                    async for event in daemon_client.chat_stream(
                        message=daemon_message, stream=True, provider=provider,
                        system=system_prompt, session_id=daemon_session_id
                    ):
                        # Internal control events
                        if event.get("type") == "done":
                            current_session_id = event.get("session_id") or current_session_id
                            yield f"data: [DONE]\n\n"
                            # Persist session_id for multi-turn
                            if current_session_id and current_session_id != daemon_session_id:
                                app.state.sessions[conv_key] = current_session_id
                            return

                        translated = _translate_daemon_event(event, model or "unknown", tool_idx)
                        if translated is None:
                            continue

                        if translated.get("type") == "done":
                            yield f"data: [DONE]\n\n"
                            return

                        if "choices" in translated and translated["choices"]:
                            for c in translated["choices"]:
                                if "tool_calls" in c.get("delta", {}):
                                    tool_idx += 1

                        chunk = _build_openai_chunk(translated, model)
                        yield f"data: {json.dumps(chunk)}\n\n"

                    yield f"data: [DONE]\n\n"

                except HTTPException:
                    raise
                except Exception as e:
                    logger.error("Stream error: %s", e)
                    yield f"data: {json.dumps({'error': str(e)})}\n\n"

            return StreamingResponse(
                event_stream(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "X-Accel-Buffering": "no",
                    "Access-Control-Allow-Origin": "*",
                },
            )

    return app


# ─── CLI ──────────────────────────────────────────────────────────────────────


def load_config_for_cli() -> dict:
    try:
        return load_config()
    except Exception:
        return dict(DEFAULT_CONFIG)


def cmd_serve(args: argparse.Namespace):
    """Start the proxy server only (no daemon management)."""
    config = load_config_for_cli()
    if args.port:
        config["server"]["port"] = args.port
    if args.host:
        config["server"]["host"] = args.host

    # Setup logging
    log_level = getattr(logging, config["logging"].get("level", "INFO").upper(), logging.INFO)
    log_file = config["logging"].get("file")
    logging.basicConfig(level=log_level, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")

    host = config["server"]["host"]
    port = config["server"]["port"]

    logger.info("=" * 60)
    logger.info("AtomCode Reverse Proxy v2.0.0")
    logger.info("Listening on http://%s:%d", host, port)
    logger.info("OpenAI endpoint: http://localhost:%d/v1", port)
    logger.info("Daemon: %s", config["daemon"]["url"])
    logger.info("Daemon auto-manage: %s", config["daemon"].get("auto_manage", False))
    logger.info("=" * 60)

    app = create_app(config)
    import uvicorn
    uvicorn.run(app, host=host, port=port, log_level=config["logging"].get("level", "info").lower())


def cmd_start(args: argparse.Namespace):
    """Start daemon + proxy together."""
    config = load_config_for_cli()
    if args.port:
        config["server"]["port"] = args.port

    log_level = getattr(logging, config["logging"].get("level", "INFO").upper(), logging.INFO)
    logging.basicConfig(level=log_level, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")

    # Ensure auto_manage is enabled
    config["daemon"]["auto_manage"] = True

    host = config["server"]["host"]
    port = config["server"]["port"]

    logger.info("=" * 60)
    logger.info("AtomCode Reverse Proxy v2.0.0 — START")
    logger.info("=" * 60)

    # Start daemon
    mgr = DaemonManager(config)
    if not mgr.start():
        logger.error("Failed to start daemon. Aborting.")
        sys.exit(1)

    # Check login / attempt auto-login
    client = DaemonClient(config["daemon"]["url"])
    try:
        auth = client.auth_status()
        if not auth.get("logged_in"):
            logger.warning("Not logged in.")
            if config.get("setup", {}).get("auto_login", True):
                logger.info("Attempting auto-login...")
                _run_auto_login(client, config)
    except Exception as e:
        logger.warning("Cannot check auth: %s", e)

    # Start proxy (in same process via uvicorn)
    logger.info("Starting proxy on http://%s:%d", host, port)
    logger.info("Daemon: %s", config["daemon"]["url"])
    logger.info("=" * 60)

    app = create_app(config)

    import uvicorn
    try:
        uvicorn.run(app, host=host, port=port, log_level=config["logging"].get("level", "info").lower())
    except KeyboardInterrupt:
        logger.info("Shutting down...")
    finally:
        mgr.stop()


def cmd_stop(args: argparse.Namespace):
    """Stop daemon (and proxy if running via start)."""
    config = load_config_for_cli()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")

    mgr = DaemonManager(config)
    mgr.stop()
    logger.info("Stopped.")


def cmd_restart(args: argparse.Namespace):
    """Restart daemon + proxy."""
    cmd_stop(args)
    time.sleep(1)
    cmd_start(args)


def cmd_status(args: argparse.Namespace):
    """Show daemon and proxy status."""
    config = load_config_for_cli()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")

    daemon_url = config["daemon"]["url"]
    client = DaemonClient(daemon_url)

    logger.info("AtomCode Reverse Proxy — Status")
    logger.info("=" * 60)

    # Daemon health
    try:
        h = client.health()
        logger.info("Daemon:       ✅ %s (v%s)", h.get("status", "?"), h.get("version", "?"))
    except Exception as e:
        logger.info("Daemon:       ❌ unreachable (%s)", e)

    # Auth status
    try:
        auth = client.auth_status()
        if auth.get("logged_in"):
            user = auth.get("user", {})
            token = auth.get("token", {})
            logger.info("Logged in:    ✅ %s (%s)", user.get("name", ""), user.get("username", ""))
            expires = token.get("expires_in", 0)
            expires_h = expires // 3600
            logger.info("Token:        expires in %dh, has refresh: %s",
                         expires_h, token.get("has_refresh_token", False))
        else:
            logger.info("Logged in:    ❌")
    except Exception as e:
        logger.info("Logged in:    ❌ (%s)", e)

    # Models
    try:
        models = client.models_list()
        logger.info("Models:       %d available", len(models))
        for m in models:
            logger.info("  - %s", m.get("model", m.get("name", "?")))
    except Exception as e:
        logger.info("Models:       ❌ (%s)", e)

    # Proxy config
    logger.info("Config:")
    logger.info("  Daemon URL:  %s", daemon_url)
    logger.info("  Binary:      %s", config["daemon"].get("binary_path", "auto"))
    logger.info("  Auto-manage: %s", config["daemon"].get("auto_manage", False))
    logger.info("  Config file: %s", _find_config_file())
    logger.info("=" * 60)


def cmd_login(args: argparse.Namespace):
    """Run interactive OAuth login."""
    config = load_config_for_cli()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
    client = DaemonClient(config["daemon"]["url"])
    success = _run_auto_login(client, config)
    sys.exit(0 if success else 1)


def cmd_daemon(args: argparse.Namespace):
    """Manage daemon process."""
    config = load_config_for_cli()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
    mgr = DaemonManager(config)

    if args.daemon_action == "status":
        if mgr.running:
            logger.info("Daemon is running (PID %d)", mgr._proc.pid if mgr._proc else "?")
        else:
            # Check via HTTP as fallback (daemon may have been started externally)
            try:
                h = client.health()
                if h.get("status") == "ok":
                    logger.info("Daemon is running externally (v%s)", h.get("version", "?"))
                    logger.info("  (not managed by this proxy)")
                else:
                    logger.info("Daemon is not running")
            except Exception:
                logger.info("Daemon is not running")
    elif args.daemon_action == "start":
        mgr.start()
    elif args.daemon_action == "stop":
        mgr.stop()
    elif args.daemon_action == "restart":
        mgr.restart()
    else:
        logger.error("Unknown action: %s", args.daemon_action)


def cmd_setup(args: argparse.Namespace):
    """Login + CodingPlan claim."""
    config = load_config_for_cli()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
    client = DaemonClient(config["daemon"]["url"])
    success = _run_full_setup(client)
    sys.exit(0 if success else 1)


def _run_auto_login(client: DaemonClient, config: dict) -> bool:
    """Interactive OAuth login flow."""
    try:
        auth = client.auth_status()
        if auth.get("logged_in"):
            user = auth.get("user", {})
            logger.info("Already logged in as %s (%s)", user.get("name", ""), user.get("username", ""))
            return True
    except Exception:
        pass

    try:
        login_data = client.login_start()
    except Exception as e:
        logger.error("Cannot start login: %s", e)
        logger.error("Is the daemon running?")
        return False

    login_id = login_data.get("login_id")
    url = login_data.get("url")
    expires = login_data.get("expires_in_seconds", 600)

    if not url:
        logger.error("Failed to get login URL: %s", login_data)
        return False

    logger.info("")
    logger.info("=" * 60)
    logger.info("  Open this URL in your browser to authorize:")
    logger.info("  %s", url)
    logger.info("=" * 60)
    logger.info("")
    logger.info("Login session expires in %d seconds.", expires)

    deadline = time.time() + expires
    poll_interval = 2

    while time.time() < deadline:
        result = client.login_poll(login_id)
        if result.get("status") == "authorized":
            user = result.get("user", {})
            logger.info("")
            logger.info("✅ Login successful!")
            logger.info("   User: %s (%s)", user.get("name", ""), user.get("username", ""))
            return True
        if "error" in result:
            logger.error("Login error: %s", result["error"])
            return False
        sys.stdout.write(".")
        sys.stdout.flush()
        time.sleep(poll_interval)

    logger.error("Login timed out.")
    return False


def _run_full_setup(client: DaemonClient) -> bool:
    """Full setup: login + CodingPlan claim."""
    if not _run_auto_login(client, {}):
        return False

    try:
        logger.info("Claiming CodingPlan...")
        cp = client.codingplan_setup()
        if cp.get("success"):
            logger.info("✅ CodingPlan claimed successfully!")
            logger.info("   Default provider: %s", cp.get("default_provider", ""))
            for m in cp.get("providers", []):
                logger.info("   - %s  (%s)", m.get("name", ""), m.get("model", ""))
            return True
        else:
            logger.error("CodingPlan setup failed: %s", cp)
            return False
    except Exception as e:
        logger.error("CodingPlan error: %s", e)
        return False


def _find_config_file() -> str:
    for p in CONFIG_SEARCH_PATHS:
        resolved = p.expanduser().resolve() if hasattr(p, "expanduser") else pathlib.Path(p).expanduser().resolve()
        if resolved.exists():
            return str(resolved)
    return "(default config)"


def main():
    parser = argparse.ArgumentParser(
        description="AtomCode Reverse Proxy v2.0 — OpenAI-compatible API for CodingPlan",
    )
    parser.add_argument("--config", help="Path to config YAML file (optional, auto-detected)")

    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # serve
    serve_p = subparsers.add_parser("serve", help="Start proxy server only (no daemon management)")
    serve_p.add_argument("--port", type=int, default=None, help="Port (default: config or 8080)")
    serve_p.add_argument("--host", type=str, default=None, help="Host (default: config or 0.0.0.0)")

    # start
    start_p = subparsers.add_parser("start", help="Start daemon + proxy together")
    start_p.add_argument("--port", type=int, default=None, help="Proxy port")

    # stop / restart
    subparsers.add_parser("stop", help="Stop daemon")
    subparsers.add_parser("restart", help="Restart daemon + proxy")

    # status
    subparsers.add_parser("status", help="Show daemon + proxy status")

    # login
    subparsers.add_parser("login", help="Run OAuth login (requires running daemon)")

    # setup
    subparsers.add_parser("setup", help="Login + CodingPlan claim (requires running daemon)")

    # daemon
    daemon_p = subparsers.add_parser("daemon", help="Manage daemon process")
    daemon_p.add_argument("daemon_action", nargs="?", default="status",
                          choices=["status", "start", "stop", "restart"],
                          help="Daemon action (default: status)")

    # models
    subparsers.add_parser("models", help="List available models from daemon")

    args = parser.parse_args()

    if args.command == "serve":
        cmd_serve(args)
    elif args.command == "start":
        cmd_start(args)
    elif args.command == "stop":
        cmd_stop(args)
    elif args.command == "restart":
        cmd_restart(args)
    elif args.command == "status":
        cmd_status(args)
    elif args.command == "login":
        cmd_login(args)
    elif args.command == "setup":
        cmd_setup(args)
    elif args.command == "daemon":
        cmd_daemon(args)
    elif args.command == "models":
        config = load_config_for_cli()
        logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
        client = DaemonClient(config["daemon"]["url"])
        try:
            models = client.models_list()
            if models:
                logger.info("Available models (%d):", len(models))
                for m in models:
                    logger.info("  - %s  (provider: %s)", m.get("model", "?"), m.get("provider", "?"))
            else:
                logger.warning("No models returned.")
        except Exception as e:
            logger.error("Failed to fetch models: %s", e)
            sys.exit(1)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()