"""
Gemma 4 12B QAT — multimodal inference API
Wraps Ollama's /api/chat endpoint and normalises it into a clean REST API
that the Next.js frontend (and any other client) can consume.

Endpoints
─────────
GET  /health          → liveness + model status
POST /chat            → text-only chat
POST /chat/image      → image-only or text+image (multipart)
POST /chat/openai     → OpenAI-compatible /v1/chat/completions passthrough
"""

import base64
import os
import httpx
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional

# ── Config ────────────────────────────────────────────────────────────────────
OLLAMA_BASE = "http://localhost:11434"
MODEL = os.getenv("MODEL_TAG", "gemma4:12b-it-qat")
TIMEOUT = 300  # seconds — large model, be generous

app = FastAPI(
    title="Gemma 4 12B QAT API",
    description="Multimodal inference via Ollama — text, image, or text+image",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],   # tighten to your Next.js origin in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Pydantic schemas ──────────────────────────────────────────────────────────

class Message(BaseModel):
    role: str          # "user" | "assistant" | "system"
    content: str

class ChatRequest(BaseModel):
    messages: list[Message]
    system: Optional[str] = None
    stream: bool = False
    temperature: float = 1.0
    top_p: float = 0.95
    top_k: int = 64

class ChatResponse(BaseModel):
    model: str
    message: Message
    done: bool
    total_duration: Optional[int] = None
    prompt_eval_count: Optional[int] = None
    eval_count: Optional[int] = None


# ── Helpers ───────────────────────────────────────────────────────────────────

def _build_ollama_payload(messages: list[dict], images: list[str] | None = None,
                          system: str | None = None,
                          temperature: float = 1.0,
                          top_p: float = 0.95,
                          top_k: int = 64,
                          stream: bool = False) -> dict:
    """Build the payload for POST /api/chat on the local Ollama server."""
    payload: dict = {
        "model": MODEL,
        "messages": messages,
        "stream": stream,
        "options": {
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
        },
    }
    if system:
        payload["system"] = system
    # Attach images to the last user message if provided
    if images:
        for msg in reversed(payload["messages"]):
            if msg["role"] == "user":
                msg["images"] = images
                break
    return payload

async def _post_ollama(payload: dict) -> dict:
    """Send payload to Ollama and return the parsed JSON response."""
    async with httpx.AsyncClient(timeout=TIMEOUT) as client:
        try:
            resp = await client.post(f"{OLLAMA_BASE}/api/chat", json=payload)
            resp.raise_for_status()
            return resp.json()
        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=502, detail=f"Ollama error: {e.response.text}")
        except httpx.ConnectError:
            raise HTTPException(status_code=503, detail="Ollama is not reachable — still loading?")


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/health")
async def health():
    """Liveness + model availability check."""
    async with httpx.AsyncClient(timeout=10) as client:
        try:
            r = await client.get(f"{OLLAMA_BASE}/api/tags")
            tags = r.json().get("models", [])
            pulled = any(MODEL in m.get("name", "") for m in tags)
            return {
                "status": "ok",
                "ollama": "running",
                "model": MODEL,
                "model_pulled": pulled,
                "available_models": [m["name"] for m in tags],
            }
        except Exception as e:
            return {"status": "degraded", "error": str(e)}


@app.post("/chat", response_model=ChatResponse)
async def chat_text(req: ChatRequest):
    """
    Text-only chat.

    Body:
      {
        "messages": [{"role": "user", "content": "Hello!"}],
        "system": "You are a helpful assistant.",   // optional
        "temperature": 1.0,
        "top_p": 0.95,
        "top_k": 64
      }
    """
    messages = [m.model_dump() for m in req.messages]
    payload = _build_ollama_payload(
        messages,
        system=req.system,
        temperature=req.temperature,
        top_p=req.top_p,
        top_k=req.top_k,
        stream=False,
    )
    data = await _post_ollama(payload)
    return ChatResponse(
        model=data.get("model", MODEL),
        message=Message(**data["message"]),
        done=data.get("done", True),
        total_duration=data.get("total_duration"),
        prompt_eval_count=data.get("prompt_eval_count"),
        eval_count=data.get("eval_count"),
    )


@app.post("/chat/image")
async def chat_image(
    prompt: str = Form(...),
    system: Optional[str] = Form(None),
    history: Optional[str] = Form(None),   # JSON string of prior messages
    temperature: float = Form(1.0),
    top_p: float = Form(0.95),
    top_k: int = Form(64),
    file: UploadFile = File(...),
):
    """
    Multimodal chat — image (+ optional text).

    Send as multipart/form-data:
      - prompt     : str   — user's text question about the image
      - file       : file  — JPEG / PNG / WEBP
      - system     : str   — optional system prompt
      - history    : str   — JSON array of {role, content} prior turns
      - temperature / top_p / top_k — sampling params
    """
    # Validate MIME type
    allowed = {"image/jpeg", "image/png", "image/webp", "image/gif"}
    if file.content_type not in allowed:
        raise HTTPException(status_code=415, detail=f"Unsupported image type: {file.content_type}")

    raw = await file.read()
    b64_image = base64.b64encode(raw).decode("utf-8")

    # Build message list
    import json
    messages: list[dict] = []
    if history:
        try:
            messages = json.loads(history)
        except json.JSONDecodeError:
            raise HTTPException(status_code=400, detail="history must be valid JSON")

    messages.append({"role": "user", "content": prompt})

    payload = _build_ollama_payload(
        messages,
        images=[b64_image],
        system=system,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        stream=False,
    )
    data = await _post_ollama(payload)
    return {
        "model": data.get("model", MODEL),
        "message": data.get("message", {}),
        "done": data.get("done", True),
        "total_duration": data.get("total_duration"),
        "prompt_eval_count": data.get("prompt_eval_count"),
        "eval_count": data.get("eval_count"),
    }


@app.post("/v1/chat/completions")
async def openai_compat(body: dict):
    """
    OpenAI-compatible completions endpoint.
    Allows the Next.js app (or any OpenAI SDK) to point at this Space
    and use it as a drop-in replacement.

    Just set:
      baseURL = "https://<your-space>.hf.space/v1"
      apiKey  = "ollama"   // any non-empty string
    """
    oai_messages = body.get("messages", [])
    stream = body.get("stream", False)
    options = {
        "temperature": body.get("temperature", 1.0),
        "top_p": body.get("top_p", 0.95),
        "top_k": body.get("top_k", 64),
    }

    # Forward directly to Ollama's OpenAI-compat endpoint
    async with httpx.AsyncClient(timeout=TIMEOUT) as client:
        try:
            r = await client.post(
                f"{OLLAMA_BASE}/v1/chat/completions",
                json={**body, "model": MODEL},
            )
            r.raise_for_status()
            return r.json()
        except httpx.HTTPStatusError as e:
            raise HTTPException(status_code=502, detail=str(e))