"""
claude-max-bridge: OpenAI-compatible API server that translates requests
to Claude Code CLI subprocess calls, enabling use of a Claude Max subscription
via LiteLLM or any OpenAI-compatible client.
"""

import asyncio
import json
import logging
import os
import time
import uuid
from typing import AsyncGenerator, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

CLAUDE_BIN = os.environ.get("CLAUDE_BIN", "/usr/local/bin/claude")
CLAUDE_HOME = os.environ.get("CLAUDE_HOME", "/root/.claude")
LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO")
REQUEST_TIMEOUT_S = float(os.environ.get("REQUEST_TIMEOUT_S", "120"))

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL.upper(), logging.INFO),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
log = logging.getLogger("claude-max-bridge")

# ---------------------------------------------------------------------------
# Supported models (pass-through to CLI --model flag)
# ---------------------------------------------------------------------------

SUPPORTED_MODELS = [
    "claude-opus-4-5",
    "claude-sonnet-4-6",
    "claude-haiku-4-5",
    "claude-opus-4-0",
    "claude-sonnet-3-7",
    "claude-haiku-3-5",
]

DEFAULT_MODEL = "claude-sonnet-4-6"

# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(title="claude-max-bridge", version="1.0.0")


# ---------------------------------------------------------------------------
# Helper: build subprocess env
# ---------------------------------------------------------------------------

def _subprocess_env() -> dict:
    env = os.environ.copy()
    env["HOME"] = "/root"
    env["XDG_CONFIG_HOME"] = os.path.join(CLAUDE_HOME, "..")  # ~/.claude lives at CLAUDE_HOME
    # Ensure the claude binary can find its config
    env["CLAUDE_CONFIG_DIR"] = CLAUDE_HOME
    # Prevent interactive prompts
    env["TERM"] = "dumb"
    env.pop("DISPLAY", None)
    return env


# ---------------------------------------------------------------------------
# Helper: parse OpenAI messages → (system_prompt, conversation_ndjson)
# ---------------------------------------------------------------------------

def _parse_messages(messages: list[dict]) -> tuple[Optional[str], list[dict]]:
    """
    Split OpenAI messages into:
    - system_prompt: first system message content (str | None)
    - turns: list of {type, message} dicts suitable for NDJSON stdin
    """
    system_prompt: Optional[str] = None
    turns: list[dict] = []

    for msg in messages:
        role = msg.get("role", "")
        content = msg.get("content", "")

        # Normalise content: list-of-parts → plain string
        if isinstance(content, list):
            text_parts = []
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    text_parts.append(part.get("text", ""))
                elif isinstance(part, str):
                    text_parts.append(part)
            content = "".join(text_parts)

        if role == "system":
            if system_prompt is None:
                system_prompt = content
            else:
                # Additional system messages: prepend to next user turn or skip
                log.debug("Ignoring additional system message: %s", content[:80])
        elif role == "user":
            turns.append({"type": "user", "message": {"role": "user", "content": content}})
        elif role == "assistant":
            turns.append({"type": "assistant", "message": {"role": "assistant", "content": content}})
        else:
            log.debug("Ignoring unknown role: %s", role)

    return system_prompt, turns


# ---------------------------------------------------------------------------
# Helper: build CLI command
# ---------------------------------------------------------------------------

def _build_command(model: str, system_prompt: Optional[str]) -> list[str]:
    cmd = [
        CLAUDE_BIN,
        "--print",
        "--output-format", "stream-json",
        "--input-format", "stream-json",
        "--include-partial-messages",
        "--no-session-persistence",
        "--tools", "",
        "--model", model,
    ]
    if system_prompt:
        cmd += ["--system-prompt", system_prompt]
    return cmd


# ---------------------------------------------------------------------------
# Core: run subprocess and stream parsed events
# ---------------------------------------------------------------------------

async def _run_claude(
    model: str,
    system_prompt: Optional[str],
    turns: list[dict],
) -> AsyncGenerator[dict, None]:
    """
    Yields dicts with keys:
      {"event": "delta", "text": "..."}
      {"event": "done", "usage": {...}}
      {"event": "error", "status": int, "message": "..."}
    """
    cmd = _build_command(model, system_prompt)
    log.debug("Spawning: %s", " ".join(cmd[:6]) + " ...")
    env = _subprocess_env()

    proc: Optional[asyncio.subprocess.Process] = None
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )

        # Write conversation turns as NDJSON then close stdin
        stdin_data = "\n".join(json.dumps(t) for t in turns) + "\n"
        proc.stdin.write(stdin_data.encode())
        await proc.stdin.drain()
        proc.stdin.close()

        last_text = ""
        usage: dict = {}

        # Read stdout line by line with overall timeout
        async def _read_lines():
            nonlocal last_text, usage
            assert proc is not None
            assert proc.stdout is not None
            async for raw_line in proc.stdout:
                line = raw_line.decode("utf-8", errors="replace").strip()
                if not line:
                    continue
                try:
                    event = json.loads(line)
                except json.JSONDecodeError:
                    log.warning("Non-JSON stdout line: %s", line[:200])
                    continue

                etype = event.get("type", "")

                if etype == "assistant":
                    msg = event.get("message", {})
                    content = msg.get("content", [])
                    # content is a list; find the first text block
                    current_text = ""
                    for block in content:
                        if isinstance(block, dict) and block.get("type") == "text":
                            current_text = block.get("text", "")
                            break
                    if current_text and current_text != last_text:
                        delta = current_text[len(last_text):]
                        last_text = current_text
                        if delta:
                            yield {"event": "delta", "text": delta}

                elif etype == "result":
                    raw_usage = event.get("usage", {})
                    usage = {
                        "prompt_tokens": raw_usage.get("input_tokens", 0),
                        "completion_tokens": raw_usage.get("output_tokens", 0),
                        "total_tokens": (
                            raw_usage.get("input_tokens", 0)
                            + raw_usage.get("output_tokens", 0)
                        ),
                    }
                    yield {"event": "done", "usage": usage}

                elif etype == "system":
                    subtype = event.get("subtype", "")
                    if subtype == "error":
                        err_msg = event.get("error", {}).get("message", "Unknown CLI error")
                        yield {"event": "error", "status": 500, "message": err_msg}
                        return

        try:
            async with asyncio.timeout(REQUEST_TIMEOUT_S):
                async for ev in _read_lines():
                    yield ev
        except asyncio.TimeoutError:
            log.error("Subprocess timed out after %.0fs", REQUEST_TIMEOUT_S)
            yield {"event": "error", "status": 504, "message": f"Claude CLI timed out after {REQUEST_TIMEOUT_S:.0f}s"}
            return

        # Wait for process to finish and collect stderr
        try:
            _, stderr_bytes = await asyncio.wait_for(
                proc.communicate(), timeout=10.0
            )
        except asyncio.TimeoutError:
            stderr_bytes = b""

        stderr_text = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else ""
        rc = proc.returncode

        if rc != 0:
            log.error("Claude CLI exited %d. stderr: %s", rc, stderr_text[:500])
            stderr_lower = stderr_text.lower()
            if "rate limit" in stderr_lower or (rc == 1 and "usage limit" in stderr_lower):
                yield {"event": "error", "status": 429, "message": stderr_text or "Rate limit exceeded"}
            elif "not logged in" in stderr_lower or "authentication" in stderr_lower:
                yield {"event": "error", "status": 401, "message": "Claude CLI authentication error"}
            else:
                yield {"event": "error", "status": 500, "message": f"Claude CLI error (exit {rc}): {stderr_text[:300]}"}

    except Exception as exc:
        log.exception("Unexpected error running Claude CLI")
        yield {"event": "error", "status": 500, "message": str(exc)}
    finally:
        if proc is not None and proc.returncode is None:
            log.debug("Killing lingering subprocess")
            try:
                proc.kill()
                await asyncio.wait_for(proc.wait(), timeout=5.0)
            except Exception:
                pass


# ---------------------------------------------------------------------------
# Streaming response builder
# ---------------------------------------------------------------------------

async def _stream_response(
    request: Request,
    model: str,
    system_prompt: Optional[str],
    turns: list[dict],
) -> AsyncGenerator[str, None]:
    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
    created = int(time.time())

    async def _gen():
        got_done = False
        async for ev in _run_claude(model, system_prompt, turns):
            # Check if client disconnected
            if await request.is_disconnected():
                log.info("Client disconnected, stopping stream")
                return

            if ev["event"] == "delta":
                chunk = {
                    "id": completion_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"content": ev["text"]},
                            "finish_reason": None,
                        }
                    ],
                }
                yield f"data: {json.dumps(chunk)}\n\n"

            elif ev["event"] == "done":
                got_done = True
                # Final chunk with finish_reason
                final_chunk = {
                    "id": completion_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {},
                            "finish_reason": "stop",
                        }
                    ],
                    "usage": ev.get("usage", {}),
                }
                yield f"data: {json.dumps(final_chunk)}\n\n"
                yield "data: [DONE]\n\n"

            elif ev["event"] == "error":
                # Emit an error chunk — client will see it in the stream
                err_chunk = {
                    "id": completion_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model,
                    "choices": [],
                    "error": {"type": "server_error", "message": ev["message"]},
                }
                yield f"data: {json.dumps(err_chunk)}\n\n"
                yield "data: [DONE]\n\n"
                return

        if not got_done:
            # Ensure stream is always terminated
            yield "data: [DONE]\n\n"

    return _gen()


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/health")
async def health():
    return {"status": "ok"}


@app.get("/v1/models")
async def list_models():
    now = int(time.time())
    return {
        "object": "list",
        "data": [
            {
                "id": m,
                "object": "model",
                "created": now,
                "owned_by": "anthropic",
            }
            for m in SUPPORTED_MODELS
        ],
    }


@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    try:
        body = await request.json()
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    # --- Parse request ---
    messages = body.get("messages", [])
    if not messages:
        raise HTTPException(status_code=400, detail="messages must not be empty")

    raw_model = body.get("model", DEFAULT_MODEL)
    # Accept model aliases / prefixes gracefully
    model = raw_model if raw_model in SUPPORTED_MODELS else DEFAULT_MODEL
    if model != raw_model:
        log.debug("Unknown model %r, falling back to %s", raw_model, model)

    stream = bool(body.get("stream", False))

    # Warn about unsupported / silently-ignored params
    _IGNORED = {"temperature", "top_p", "frequency_penalty", "presence_penalty",
                "logit_bias", "n", "best_of", "logprobs", "top_logprobs",
                "response_format", "seed", "stop", "user"}
    for param in _IGNORED:
        if param in body:
            log.debug("Ignoring unsupported parameter: %s", param)

    system_prompt, turns = _parse_messages(messages)

    if not turns:
        raise HTTPException(status_code=400, detail="No user/assistant turns found in messages")

    # --- Streaming ---
    if stream:
        gen = await _stream_response(request, model, system_prompt, turns)
        return StreamingResponse(
            gen,
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "X-Accel-Buffering": "no",
            },
        )

    # --- Non-streaming: buffer all deltas ---
    full_text = ""
    usage: dict = {}
    error_ev: Optional[dict] = None

    async for ev in _run_claude(model, system_prompt, turns):
        if ev["event"] == "delta":
            full_text += ev["text"]
        elif ev["event"] == "done":
            usage = ev.get("usage", {})
        elif ev["event"] == "error":
            error_ev = ev
            break

    if error_ev:
        status_code = error_ev.get("status", 500)
        err_type = {
            429: "rate_limit_error",
            401: "authentication_error",
        }.get(status_code, "server_error")
        return JSONResponse(
            status_code=status_code,
            content={"error": {"type": err_type, "message": error_ev["message"]}},
        )

    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
    return {
        "id": completion_id,
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": full_text},
                "finish_reason": "stop",
            }
        ],
        "usage": usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }