"""
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",
    # Short aliases accepted by the CLI
    "sonnet",
    "opus",
    "haiku",
]

DEFAULT_MODEL = "claude-sonnet-4-6"

# Valid effort levels accepted by the CLI
_VALID_EFFORT_LEVELS = {"low", "medium", "high", "xhigh", "max"}

# ---------------------------------------------------------------------------
# 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: map temperature float to effort level string
# ---------------------------------------------------------------------------

def _temperature_to_effort(temperature: float) -> str:
    if temperature < 0.2:
        return "low"
    elif temperature < 0.4:
        return "medium"
    elif temperature < 0.7:
        return "high"
    elif temperature < 0.9:
        return "xhigh"
    else:
        return "max"


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

def _build_command(
    model: str,
    system_prompt,
    effort=None,
    temperature=None,
    fallback_model=None,
    system_prompt_mode="replace",
    max_turns=None,
    max_budget_usd=None,
    json_schema=None,
    exclude_dynamic_sections=False,
):
    import json as _json
    cmd = [
        CLAUDE_BIN,
        "--print",
        "--output-format", "stream-json",
        "--input-format", "stream-json",
        "--verbose",
        "--include-partial-messages",
        "--no-session-persistence",
        "--mcp-config", "{\"mcpServers\":{}}",
        "--strict-mcp-config",
        "--tools", "",
        "--model", model,
    ]

    resolved_effort = None
    if effort is not None and effort in _VALID_EFFORT_LEVELS:
        resolved_effort = effort
    elif temperature is not None:
        resolved_effort = _temperature_to_effort(temperature)

    if resolved_effort is not None:
        cmd += ["--effort", resolved_effort]

    if fallback_model is not None:
        cmd += ["--fallback-model", fallback_model]

    if system_prompt:
        if system_prompt_mode == "append":
            cmd += ["--append-system-prompt", system_prompt]
        else:
            cmd += ["--system-prompt", system_prompt]

    if max_turns is not None:
        cmd += ["--max-turns", str(max_turns)]

    if max_budget_usd is not None:
        cmd += ["--max-budget-usd", str(max_budget_usd)]

    if json_schema is not None:
        cmd += ["--json-schema", _json.dumps(json_schema)]

    if exclude_dynamic_sections:
        cmd += ["--exclude-dynamic-system-prompt-sections"]

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

async def _run_claude(
    model: str,
    system_prompt,
    turns: list,
    effort=None,
    temperature=None,
    fallback_model=None,
    system_prompt_mode="replace",
    max_turns=None,
    max_budget_usd=None,
    json_schema=None,
    exclude_dynamic_sections=False,
):
    """
    Yields dicts with keys:
      {"event": "delta", "text": "..."}
      {"event": "done", "usage": {...}, "cost_usd": float | None, "rate_limit": dict | None}
      {"event": "error", "status": int, "message": "..."}
    """
    cmd = _build_command(
        model, system_prompt,
        effort=effort, temperature=temperature, fallback_model=fallback_model,
        system_prompt_mode=system_prompt_mode, max_turns=max_turns,
        max_budget_usd=max_budget_usd, json_schema=json_schema,
        exclude_dynamic_sections=exclude_dynamic_sections,
    )
    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 = {}
        cost_usd: Optional[float] = None
        rate_limit_info: Optional[dict] = None

        # Read stdout line by line with overall timeout
        async def _read_lines():
            nonlocal last_text, usage, cost_usd, rate_limit_info
            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)
                        ),
                    }
                    # Extract cost from result event
                    raw_cost = event.get("total_cost_usd")
                    if raw_cost is not None:
                        try:
                            cost_usd = float(raw_cost)
                        except (TypeError, ValueError):
                            pass
                    yield {
                        "event": "done",
                        "usage": usage,
                        "cost_usd": cost_usd,
                        "rate_limit": rate_limit_info,
                    }

                elif etype == "rate_limit_event":
                    # Capture rate limit info for the done event
                    rate_limit_info = event.get("rate_limit_info", event)
                    log.debug("Rate limit event: %s", rate_limit_info)

                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,
    turns: list,
    effort=None,
    temperature=None,
    fallback_model=None,
    system_prompt_mode="replace",
    max_turns=None,
    max_budget_usd=None,
    json_schema=None,
    exclude_dynamic_sections=False,
):
    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,
            effort=effort, temperature=temperature, fallback_model=fallback_model,
            system_prompt_mode=system_prompt_mode, max_turns=max_turns,
            max_budget_usd=max_budget_usd, json_schema=json_schema,
            exclude_dynamic_sections=exclude_dynamic_sections,
        ):
            # 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"

                # Emit cost/rate-limit metadata chunk before [DONE]
                cost_usd = ev.get("cost_usd")
                rate_limit = ev.get("rate_limit")
                if cost_usd is not None or rate_limit is not None:
                    meta_chunk = {
                        "id": completion_id,
                        "object": "chat.completion.chunk",
                        "created": created,
                        "model": model,
                        "choices": [],
                        "usage": ev.get("usage", {}),
                    }
                    if cost_usd is not None:
                        meta_chunk["x_claude_cost_usd"] = cost_usd
                    if rate_limit is not None:
                        meta_chunk["x_claude_rate_limit"] = rate_limit
                    yield f"data: {json.dumps(meta_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))

    # Extract temperature (kept for effort mapping; otherwise ignored)
    temperature: Optional[float] = None
    if "temperature" in body:
        try:
            temperature = float(body["temperature"])
        except (TypeError, ValueError):
            pass

    # Extract effort from extra_body or top-level body
    extra_body = body.get("extra_body") or {}
    effort: Optional[str] = extra_body.get("effort") or body.get("effort")
    if effort is not None:
        effort = str(effort)

    # Extract fallback_model from extra_body
    fallback_model: Optional[str] = extra_body.get("fallback_model")
    if fallback_model is not None:
        fallback_model = str(fallback_model)

    # Extract new CLI-mapped parameters from extra_body
    # system_prompt_mode: "replace" (default, --system-prompt) or "append" (--append-system-prompt)
    system_prompt_mode: str = str(extra_body.get("system_prompt_mode", "replace"))
    if system_prompt_mode not in ("replace", "append"):
        system_prompt_mode = "replace"

    max_turns: Optional[int] = None
    raw_max_turns = extra_body.get("max_turns")
    if raw_max_turns is not None:
        try:
            max_turns = int(raw_max_turns)
        except (TypeError, ValueError):
            pass

    max_budget_usd: Optional[float] = None
    raw_budget = extra_body.get("max_budget_usd")
    if raw_budget is not None:
        try:
            max_budget_usd = float(raw_budget)
        except (TypeError, ValueError):
            pass

    # response_format.json_schema → --json-schema
    json_schema: Optional[dict] = None
    response_format = body.get("response_format") or {}
    if isinstance(response_format, dict):
        if response_format.get("type") == "json_schema":
            js = response_format.get("json_schema")
            if isinstance(js, dict):
                # OpenAI wraps it: {name, strict, schema} — pass the schema sub-dict to CLI
                json_schema = js.get("schema") or js
        elif response_format.get("type") == "json_object":
            # Instruct via system prompt append instead of --json-schema
            json_schema = {"type": "object"}
    # Also allow direct extra_body.json_schema override
    if extra_body.get("json_schema"):
        json_schema = extra_body["json_schema"]

    # exclude_dynamic_system_prompt_sections improves cache reuse in multi-user deployments
    exclude_dynamic_sections: bool = bool(extra_body.get("exclude_dynamic_system_prompt_sections", False))

    # Warn about other unsupported / silently-ignored params
    _IGNORED = {"top_p", "frequency_penalty", "presence_penalty",
                "logit_bias", "n", "best_of", "logprobs", "top_logprobs",
                "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,
            effort=effort, temperature=temperature, fallback_model=fallback_model,
            system_prompt_mode=system_prompt_mode, max_turns=max_turns,
            max_budget_usd=max_budget_usd, json_schema=json_schema,
            exclude_dynamic_sections=exclude_dynamic_sections,
        )
        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 = {}
    cost_usd: Optional[float] = None
    rate_limit_info: Optional[dict] = None
    error_ev: Optional[dict] = None

    async for ev in _run_claude(
        model, system_prompt, turns,
        effort=effort, temperature=temperature, fallback_model=fallback_model,
        system_prompt_mode=system_prompt_mode, max_turns=max_turns,
        max_budget_usd=max_budget_usd, json_schema=json_schema,
        exclude_dynamic_sections=exclude_dynamic_sections,
    ):
        if ev["event"] == "delta":
            full_text += ev["text"]
        elif ev["event"] == "done":
            usage = ev.get("usage", {})
            cost_usd = ev.get("cost_usd")
            rate_limit_info = ev.get("rate_limit")
        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}"
    response_body: dict = {
        "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},
    }
    if cost_usd is not None:
        response_body["x_claude_cost_usd"] = cost_usd
    if rate_limit_info is not None:
        response_body["x_claude_rate_limit"] = rate_limit_info
    return response_body