"""
claude-code-proxy: Ollama-compatible HTTP facade for `claude -p`.

Exposes a subset of the Ollama API on http://127.0.0.1:11435 and translates
each request into a `claude -p` subprocess invocation. This lets external
tools that already speak Ollama (Open WebUI, AnythingLLM, n8n nodes, etc.)
talk to Claude Code instead of a local Ollama instance.

Endpoints:
    GET  /                  health check
    GET  /api/version       Ollama version stub
    GET  /api/tags          list "models" (so clients can validate)
    POST /api/show          model details stub
    POST /api/generate      single-shot prompt -> response
    POST /api/chat          multi-message conversation -> response (vision-aware)
    POST /api/vision        multipart image diagnosis endpoint (Garden Buddy
                            Buddy Check contract: image/image2/image3 files +
                            prompt/system_prompt/backend/model fields)

Both /api/generate and /api/chat honour the `stream` flag in the request
body (Ollama default is True). When true, responses are emitted as
NDJSON chunks; when false, a single JSON object is returned.

Vision support: when /api/chat messages contain an `images` list (Ollama
vision format), each base64 blob is written to a temp file under
CLAUDE_PROXY_IMG_DIR and the file path is appended to the prompt so that
claude's Read tool can load the image as a vision input.

Environment variables:
    CLAUDE_BIN                  path to claude CLI (default: "claude")
    CLAUDE_PROXY_CONCURRENCY    max concurrent claude subprocesses (default: 3)
    CLAUDE_PROXY_MODEL          name advertised in /api/tags (default: "claude-code")
    CLAUDE_PROXY_TIMEOUT        per-request timeout in seconds (default: 300)
    CLAUDE_PROXY_IMG_DIR        temp dir for vision images (default: /tmp/claude_proxy_imgs)
    CLAUDE_CODE_OAUTH_TOKEN     long-lived auth token; used by the claude
                                subprocess AND as Bearer auth for direct
                                /v1/messages calls on the /api/vision route
                                (anthropic-beta: oauth-2025-04-20)
    CLAUDE_PROXY_VISION_TOKEN   shared secret for /api/vision Bearer auth
                                (matches gardenbuddy.au gbc_proxy_token);
                                if unset, /api/vision auth is not enforced
    OLLAMA_URL                  local Ollama base URL for the ollama backend
                                (default: http://127.0.0.1:11434)

Version: 0.2.0
Updated: 2026-06-12
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, AsyncIterator

import httpx
from anthropic import AsyncAnthropic
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse

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

CLAUDE_BIN = os.environ.get("CLAUDE_BIN", "claude")
CONCURRENCY = int(os.environ.get("CLAUDE_PROXY_CONCURRENCY", "3"))
DEFAULT_MODEL = os.environ.get("CLAUDE_PROXY_MODEL", "claude-code")
TIMEOUT_SECONDS = int(os.environ.get("CLAUDE_PROXY_TIMEOUT", "300"))
IMG_DIR = Path(os.environ.get("CLAUDE_PROXY_IMG_DIR", "/tmp/claude_proxy_imgs"))
OAUTH_TOKEN = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "")
VISION_TOKEN = os.environ.get("CLAUDE_PROXY_VISION_TOKEN", "")
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434").rstrip("/")
VISION_DEFAULT_MODEL = "claude-haiku-4-5"

# Claude Haiku 4.5 list pricing (USD per million tokens). Used only to fill the
# cost_usd tracking field. Actual marginal cost is $0 because the OAuth token
# bills against the Claude subscription, not pay-per-token.
HAIKU_IN_PER_MTOK = 1.00
HAIKU_OUT_PER_MTOK = 5.00
HAIKU_CACHE_READ_PER_MTOK = 0.10
HAIKU_CACHE_WRITE_PER_MTOK = 1.25

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("claude-proxy")

app = FastAPI(title="claude-code-proxy")
_semaphore = asyncio.Semaphore(CONCURRENCY)

IMG_DIR.mkdir(parents=True, exist_ok=True)


# --- Helpers ----------------------------------------------------------------

def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def _save_images(images: list[str]) -> list[Path]:
    """Write base64-encoded image blobs to temp files; return their paths."""
    paths: list[Path] = []
    for b64 in images:
        # Strip data-URI prefix if present (data:image/png;base64,...)
        if "," in b64:
            b64 = b64.split(",", 1)[1]
        try:
            data = base64.b64decode(b64)
        except Exception:
            log.warning("skipping malformed base64 image blob")
            continue
        # Detect format from magic bytes
        ext = ".jpg"
        if data[:8] == b"\x89PNG\r\n\x1a\n":
            ext = ".png"
        elif data[:4] in (b"RIFF", b"WEBP"):
            ext = ".webp"
        elif data[:4] == b"GIF8":
            ext = ".gif"
        tmp = tempfile.NamedTemporaryFile(
            suffix=ext, prefix="img_", dir=IMG_DIR, delete=False
        )
        tmp.write(data)
        tmp.close()
        paths.append(Path(tmp.name))
    return paths


async def _run_claude(prompt: str, image_paths: list[Path] | None = None) -> str:
    """Run `claude -p <prompt>` and return stdout as a string.

    When image_paths is provided the prompt already contains [image: path]
    references and we pass --allowedTools to permit Read on those specific files.
    """
    async with _semaphore:
        log.info(
            "claude -p invoked (prompt %d chars, images %d) | preview: %s",
            len(prompt),
            len(image_paths) if image_paths else 0,
            prompt[:120].replace("\n", " "),
        )

        cmd = [CLAUDE_BIN, "-p", prompt]
        if image_paths:
            cmd += ["--allowedTools", "Read"]

        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=os.environ.copy(),
        )
        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=TIMEOUT_SECONDS
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            raise RuntimeError(f"claude -p timed out after {TIMEOUT_SECONDS}s")
        finally:
            # Always clean up temp image files.
            if image_paths:
                for p in image_paths:
                    try:
                        p.unlink(missing_ok=True)
                    except Exception:
                        pass

        if proc.returncode != 0:
            err = stderr.decode("utf-8", errors="replace")[:1000]
            raise RuntimeError(f"claude -p exited {proc.returncode}: {err}")

        result = stdout.decode("utf-8", errors="replace")
        log.info("claude -p response: %d chars | preview: %s", len(result), result[:120].replace("\n", " "))
        return result


def _build_prompt_from_messages(messages: list[dict]) -> tuple[str, list[Path]]:
    """Flatten Ollama-style messages into a prompt string and extract any images.

    Returns (prompt_string, list_of_temp_image_paths).
    Image references are appended to the user turn that contained them so
    claude's Read tool knows to load them as vision inputs.
    """
    system_parts = [m["content"] for m in messages if m.get("role") == "system"]
    convo: list[str] = []
    all_image_paths: list[Path] = []

    for m in messages:
        role = m.get("role", "user")
        if role == "system":
            continue
        prefix = "User" if role == "user" else "Assistant"
        content = m.get("content", "")

        # Handle vision: save images and embed file references in the turn.
        images_in_msg = m.get("images", [])
        if images_in_msg:
            paths = _save_images(images_in_msg)
            all_image_paths.extend(paths)
            refs = "\n".join(f"[image: {p}]" for p in paths)
            content = f"{content}\n{refs}" if content else refs

        convo.append(f"{prefix}: {content}")

    convo.append("Assistant:")
    body = "\n\n".join(convo)
    if system_parts:
        prompt = "[System]\n" + "\n\n".join(system_parts) + "\n\n" + body
    else:
        prompt = body

    return prompt, all_image_paths


# --- Streaming generators ---------------------------------------------------

async def _stream_generate(base: dict, text: str) -> AsyncIterator[bytes]:
    """Emit Ollama-style NDJSON for /api/generate: incremental chunks then done."""
    chunk_size = 64
    started = time.time()
    for i in range(0, len(text), chunk_size):
        frame = {**base, "response": text[i:i + chunk_size], "done": False}
        yield (json.dumps(frame) + "\n").encode("utf-8")
        await asyncio.sleep(0)
    final = {
        **base,
        "response": "",
        "done": True,
        "done_reason": "stop",
        "total_duration": int((time.time() - started) * 1e9),
    }
    yield (json.dumps(final) + "\n").encode("utf-8")


async def _stream_chat(base: dict, text: str) -> AsyncIterator[bytes]:
    """Emit Ollama-style NDJSON for /api/chat: each frame carries a message."""
    chunk_size = 64
    started = time.time()
    for i in range(0, len(text), chunk_size):
        frame = {
            **base,
            "message": {"role": "assistant", "content": text[i:i + chunk_size]},
            "done": False,
        }
        yield (json.dumps(frame) + "\n").encode("utf-8")
        await asyncio.sleep(0)
    final = {
        **base,
        "message": {"role": "assistant", "content": ""},
        "done": True,
        "done_reason": "stop",
        "total_duration": int((time.time() - started) * 1e9),
    }
    yield (json.dumps(final) + "\n").encode("utf-8")


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

@app.get("/")
async def root() -> dict:
    return {"status": "ok", "service": "claude-code-proxy"}


@app.get("/api/version")
async def version() -> dict:
    return {"version": "0.2.0-claude-proxy"}


@app.get("/api/tags")
async def tags() -> dict:
    """Ollama-style model list. Many clients hit this to verify the endpoint."""
    return {
        "models": [{
            "name": DEFAULT_MODEL,
            "model": DEFAULT_MODEL,
            "modified_at": _now_iso(),
            "size": 0,
            "digest": "sha256:claude-code",
            "details": {
                "parent_model": "",
                "format": "claude",
                "family": "claude",
                "families": ["claude"],
                "parameter_size": "unknown",
                "quantization_level": "none",
            },
        }]
    }


@app.post("/api/show")
async def show(req: Request) -> dict:
    body = await req.json()
    name = body.get("name", DEFAULT_MODEL)
    return {
        "modelfile": f"FROM {name}",
        "parameters": "",
        "template": "",
        "details": {
            "format": "claude",
            "family": "claude",
            "parameter_size": "unknown",
            "quantization_level": "none",
        },
    }


@app.post("/api/generate")
async def generate(req: Request) -> Any:
    body = await req.json()
    model = body.get("model", DEFAULT_MODEL)
    prompt = body.get("prompt", "")
    system = body.get("system")
    stream = bool(body.get("stream", True))

    full_prompt = f"[System]\n{system}\n\n{prompt}" if system else prompt

    started = time.time()
    try:
        text = await _run_claude(full_prompt)
    except Exception as e:
        log.exception("claude invocation failed")
        return JSONResponse({"error": str(e)}, status_code=500)

    base = {"model": model, "created_at": _now_iso()}
    if stream:
        return StreamingResponse(
            _stream_generate(base, text),
            media_type="application/x-ndjson",
        )
    return {
        **base,
        "response": text,
        "done": True,
        "done_reason": "stop",
        "total_duration": int((time.time() - started) * 1e9),
    }


@app.post("/api/chat")
async def chat(req: Request) -> Any:
    body = await req.json()
    model = body.get("model", DEFAULT_MODEL)
    messages = body.get("messages", [])
    stream = bool(body.get("stream", True))

    prompt, image_paths = _build_prompt_from_messages(messages)

    started = time.time()
    try:
        text = await _run_claude(prompt, image_paths or None)
    except Exception as e:
        log.exception("claude invocation failed")
        return JSONResponse({"error": str(e)}, status_code=500)

    base = {"model": model, "created_at": _now_iso()}
    if stream:
        return StreamingResponse(
            _stream_chat(base, text),
            media_type="application/x-ndjson",
        )
    return {
        **base,
        "message": {"role": "assistant", "content": text},
        "done": True,
        "done_reason": "stop",
        "total_duration": int((time.time() - started) * 1e9),
    }


# --- Vision (Garden Buddy Buddy Check) ---------------------------------------

def _media_type(data: bytes) -> str:
    """Detect image media type from magic bytes (default JPEG)."""
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return "image/png"
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "image/webp"
    if data[:4] == b"GIF8":
        return "image/gif"
    return "image/jpeg"


def _haiku_cost_usd(usage: Any) -> float:
    """API-equivalent Haiku 4.5 cost. Informational only (subscription-billed)."""
    return round(
        getattr(usage, "input_tokens", 0) * HAIKU_IN_PER_MTOK / 1e6
        + getattr(usage, "output_tokens", 0) * HAIKU_OUT_PER_MTOK / 1e6
        + (getattr(usage, "cache_read_input_tokens", 0) or 0) * HAIKU_CACHE_READ_PER_MTOK / 1e6
        + (getattr(usage, "cache_creation_input_tokens", 0) or 0) * HAIKU_CACHE_WRITE_PER_MTOK / 1e6,
        6,
    )


async def _ollama_model_available(model: str) -> bool:
    try:
        async with httpx.AsyncClient(timeout=5) as cx:
            r = await cx.get(f"{OLLAMA_URL}/api/tags")
            r.raise_for_status()
            names = [m.get("name", "") for m in r.json().get("models", [])]
            return any(n == model or n.split(":")[0] == model.split(":")[0] for n in names)
    except Exception:
        return False


async def _vision_via_anthropic(
    images: list[bytes], system_prompt: str, user_prompt: str, model: str
) -> dict:
    """Direct /v1/messages call using the Claude Code OAuth token."""
    content: list[dict] = [
        {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": _media_type(img),
                "data": base64.standard_b64encode(img).decode("ascii"),
            },
        }
        for img in images
    ]
    content.append({"type": "text", "text": user_prompt})

    client = AsyncAnthropic(
        auth_token=OAUTH_TOKEN,
        default_headers={"anthropic-beta": "oauth-2025-04-20"},
    )
    started = time.time()
    try:
        msg = await client.messages.create(
            model=model,
            max_tokens=4096,
            system=system_prompt or "",
            messages=[{"role": "user", "content": content}],
        )
    finally:
        await client.close()

    text = "".join(b.text for b in msg.content if b.type == "text")
    return {
        "content": text,
        "backend": "anthropic",
        "model": msg.model,
        "duration_ms": int((time.time() - started) * 1000),
        "prompt_eval_count": msg.usage.input_tokens,
        "eval_count": msg.usage.output_tokens,
        "cache_read_tokens": msg.usage.cache_read_input_tokens or 0,
        "cache_creation_tokens": msg.usage.cache_creation_input_tokens or 0,
        "cost_usd": _haiku_cost_usd(msg.usage),
    }


async def _vision_via_ollama(
    images: list[bytes], system_prompt: str, user_prompt: str, model: str
) -> dict:
    """Local Ollama vision call (only used when the model is installed)."""
    payload = {
        "model": model,
        "stream": False,
        "messages": [
            {"role": "system", "content": system_prompt or ""},
            {
                "role": "user",
                "content": user_prompt,
                "images": [base64.standard_b64encode(i).decode("ascii") for i in images],
            },
        ],
    }
    started = time.time()
    async with httpx.AsyncClient(timeout=TIMEOUT_SECONDS) as cx:
        r = await cx.post(f"{OLLAMA_URL}/api/chat", json=payload)
        r.raise_for_status()
        data = r.json()
    return {
        "content": data.get("message", {}).get("content", ""),
        "backend": "ollama",
        "model": data.get("model", model),
        "duration_ms": int((time.time() - started) * 1000),
        "prompt_eval_count": data.get("prompt_eval_count", 0),
        "eval_count": data.get("eval_count", 0),
        "cache_read_tokens": 0,
        "cache_creation_tokens": 0,
        "cost_usd": 0.0,
    }


@app.post("/api/vision")
async def vision(req: Request) -> Any:
    """Photo diagnosis endpoint matching the gardenbuddy-buddy-check contract.

    Request:  multipart/form-data with file fields image, image2, image3 and
              text fields prompt, system_prompt, backend (anthropic|ollama),
              model. Bearer auth required when CLAUDE_PROXY_VISION_TOKEN set.
    Response: 200 JSON {content, backend, model, duration_ms,
              prompt_eval_count, eval_count, cache_read_tokens,
              cache_creation_tokens, cost_usd}.
    """
    if VISION_TOKEN:
        auth = req.headers.get("authorization", "")
        if auth != f"Bearer {VISION_TOKEN}":
            return JSONResponse({"error": "unauthorized"}, status_code=401)

    form = await req.form()
    user_prompt = str(form.get("prompt", "") or "")
    system_prompt = str(form.get("system_prompt", "") or "")
    backend = str(form.get("backend", "anthropic") or "anthropic").lower()
    model = str(form.get("model", "") or "")

    images: list[bytes] = []
    for key, value in form.multi_items():
        if key.startswith("image") and hasattr(value, "read"):
            data = await value.read()
            if data:
                images.append(data)

    if not images:
        return JSONResponse({"error": "at least one image file is required"}, status_code=400)
    if len(images) > 3:
        return JSONResponse({"error": "up to 3 images allowed"}, status_code=400)
    if not user_prompt:
        return JSONResponse({"error": "prompt is required"}, status_code=400)

    # Ollama is only honoured when the requested model is actually installed;
    # otherwise we fall back to Anthropic Haiku so a diagnosis never 404s on a
    # missing local model. (qwen2.5vl was retired from this GPU on 2026-04-19.)
    if backend == "ollama":
        ollama_model = model or "qwen2.5vl:7b"
        if await _ollama_model_available(ollama_model):
            try:
                async with _semaphore:
                    return await _vision_via_ollama(images, system_prompt, user_prompt, ollama_model)
            except Exception:
                log.exception("ollama vision failed, falling back to anthropic")
        else:
            log.warning("ollama model %s not installed; falling back to anthropic", ollama_model)
        model = ""  # discard the ollama model name for the fallback

    anthropic_model = model if model.startswith("claude") else VISION_DEFAULT_MODEL
    try:
        async with _semaphore:
            result = await _vision_via_anthropic(images, system_prompt, user_prompt, anthropic_model)
    except Exception as e:
        log.exception("anthropic vision call failed")
        return JSONResponse({"error": f"vision backend error: {e}"}, status_code=502)

    log.info(
        "vision ok: backend=%s model=%s images=%d in=%d out=%d %dms",
        result["backend"], result["model"], len(images),
        result["prompt_eval_count"], result["eval_count"], result["duration_ms"],
    )
    return result