import base64
import json
import os
from typing import Annotated

import httpx
from fastapi import FastAPI, File, Form, HTTPException, UploadFile


app = FastAPI(title="Gemini image proxy", docs_url=None, redoc_url=None)

GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-3.6-flash")
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models"
MAX_IMAGE_BYTES = 10 * 1024 * 1024
ALLOWED_TYPES = {"image/jpeg", "image/png", "image/webp"}


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/gemini")
async def recognize(
    image: Annotated[UploadFile, File(...)],
    prompt: Annotated[str, Form(...)],
    generation_config: Annotated[str, Form()] = "{}",
) -> dict:
    if image.content_type not in ALLOWED_TYPES:
        raise HTTPException(status_code=415, detail="Unsupported image type")

    image_data = await image.read(MAX_IMAGE_BYTES + 1)
    if len(image_data) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=413, detail="Image is too large")

    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise HTTPException(status_code=503, detail="Gemini is not configured")

    try:
        config = json.loads(generation_config)
    except json.JSONDecodeError as exc:
        raise HTTPException(status_code=422, detail="Invalid generation_config") from exc
    if not isinstance(config, dict):
        raise HTTPException(status_code=422, detail="generation_config must be an object")

    payload = {
        "contents": [{
            "parts": [
                {"text": prompt},
                {
                    "inline_data": {
                        "mime_type": image.content_type,
                        "data": base64.b64encode(image_data).decode("ascii"),
                    }
                },
            ]
        }],
        "generationConfig": config,
    }

    url = f"{GEMINI_API_URL}/{GEMINI_MODEL}:generateContent"
    try:
        async with httpx.AsyncClient(timeout=180) as client:
            response = await client.post(
                url,
                params={"key": api_key},
                json=payload,
            )
    except httpx.HTTPError as exc:
        raise HTTPException(status_code=503, detail="Gemini unavailable") from exc

    if response.status_code != 200:
        try:
            provider_error = response.json().get("error", {}).get("message")
        except ValueError:
            provider_error = None
        detail = provider_error or "Gemini request failed"
        raise HTTPException(status_code=502, detail=detail)

    data = response.json()
    try:
        candidate = data["candidates"][0]
        text = candidate["content"]["parts"][0]["text"]
    except (KeyError, IndexError, TypeError) as exc:
        raise HTTPException(status_code=502, detail="Invalid Gemini response") from exc

    return {"text": text, "usage": data.get("usageMetadata", {})}