#!/usr/bin/env python3
"""HTTP API that serves the pipeline (preprocess -> ocr -> structure) to
the Android app. Runs the same in-process functions the numbered host
scripts drive via `docker run` per stage, but skips the per-stage
container start-up since a live server has no need for that
inspectability - the host scripts remain the way to inspect each stage
independently during pipeline development.

Jobs are processed asynchronously so the app can show real per-stage
progress: POST /jobs returns immediately with a job id; GET /jobs/{id}
reports which of preprocessing/ocr/extracting is currently running, and
the final result once done. Job state lives in memory only - this is a
personal, single-instance server with no need for persistence here (the
Android app is the system of record for confirmed receipts).

Auth: set FUELISH_API_KEY to require a matching X-API-Key header on
/jobs (GET /health stays open for reverse-proxy health checks). Leave it
unset for the original LAN-only/no-auth setup. This server does not
terminate TLS itself - see require_api_key()'s docstring.
"""
import os
import shutil
import tempfile
import threading
import time
import uuid

import requests
from fastapi import Depends, FastAPI, HTTPException, Security, UploadFile
from fastapi.responses import JSONResponse
from fastapi.security.api_key import APIKeyHeader

from adaptive import preprocess_and_ocr_adaptive
from extract import structure

OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://127.0.0.1:11434")
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "qwen2.5:3b")
API_KEY = os.environ.get("FUELISH_API_KEY", "")

app = FastAPI(title="fuelish")

_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)


def require_api_key(provided: str | None = Security(_api_key_header)) -> None:
    """No-op when FUELISH_API_KEY isn't set - the unchanged default for a
    LAN-only deployment with no key configured. When it is set, every
    request to a protected route must carry a matching X-API-Key header.
    This only checks the key; it doesn't terminate TLS - a reverse proxy
    in front of this server (e.g. Traefik) is expected to do that for any
    deployment reachable outside a trusted LAN."""
    if not API_KEY:
        return
    if provided != API_KEY:
        raise HTTPException(status_code=401, detail="invalid or missing API key")

_jobs: dict[str, dict] = {}
_jobs_lock = threading.Lock()


def _set_job(job_id: str, **fields) -> None:
    with _jobs_lock:
        _jobs[job_id].update(fields)


def _ensure_model_pulled() -> None:
    """Best-effort: wait for Ollama to come up (compose starts `api` and
    `ollama` together, with no ordering guarantee beyond container start)
    and pull OLLAMA_MODEL if it isn't already present, so a bare `docker
    compose up` is enough without a manual `ollama pull` step. Runs in a
    background thread - jobs submitted before this finishes will simply
    fail with a clear error from Ollama."""
    tags = None
    for _ in range(60):
        try:
            tags = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=5).json()
            break
        except requests.RequestException:
            time.sleep(2)
    if tags is None:
        return

    have = {m["name"] for m in tags.get("models", [])}
    if OLLAMA_MODEL in have or any(m.startswith(OLLAMA_MODEL + ":") for m in have):
        return
    try:
        requests.post(
            f"{OLLAMA_BASE_URL}/api/pull",
            json={"name": OLLAMA_MODEL},
            timeout=1800,
        )
    except requests.RequestException:
        pass


@app.on_event("startup")
def _on_startup() -> None:
    threading.Thread(target=_ensure_model_pulled, daemon=True).start()


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


def _process_job(job_id: str, image_path: str, scratch_dir: str) -> None:
    try:
        _set_job(job_id, status="preprocessing")
        adaptive_result = preprocess_and_ocr_adaptive(image_path, scratch_dir)

        _set_job(job_id, status="extracting")
        result = structure(adaptive_result["ocr_text"])

        _set_job(job_id, status="done", result=result)
    except Exception as exc:  # noqa: BLE001 - reported to the client, not swallowed
        _set_job(job_id, status="error", error=str(exc))
    finally:
        shutil.rmtree(scratch_dir, ignore_errors=True)


@app.post("/jobs", status_code=201, dependencies=[Depends(require_api_key)])
async def create_job(image: UploadFile) -> JSONResponse:
    job_id = uuid.uuid4().hex
    scratch_dir = tempfile.mkdtemp(prefix=f"fuelish-{job_id}-")
    image_path = os.path.join(scratch_dir, "upload" + os.path.splitext(image.filename or "")[1])
    with open(image_path, "wb") as f:
        f.write(await image.read())

    with _jobs_lock:
        _jobs[job_id] = {"status": "queued", "result": None, "error": None, "_created": time.time()}

    thread = threading.Thread(target=_process_job, args=(job_id, image_path, scratch_dir), daemon=True)
    thread.start()

    return JSONResponse({"job_id": job_id}, status_code=201)


@app.get("/jobs/{job_id}", dependencies=[Depends(require_api_key)])
def get_job(job_id: str) -> dict:
    with _jobs_lock:
        job = _jobs.get(job_id)
        if job is None:
            raise HTTPException(status_code=404, detail="job not found")
        return {"status": job["status"], "result": job["result"], "error": job["error"]}