from __future__ import annotations

import asyncio
import io
import os
import secrets
import threading
from contextlib import asynccontextmanager

from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from PIL import Image, ImageOps, UnidentifiedImageError

from detectors import Detector, load_detectors

try:
    import spaces
except ImportError:  # Docker and ordinary CPU/GPU hosts do not need the Spaces runtime.
    spaces = None

if spaces is not None:
    @spaces.GPU(duration=1)
    def zerogpu_runtime_probe() -> bool:
        """Declares ZeroGPU compatibility; authenticated detector inference remains on CPU."""
        return True

MAX_IMAGE_BYTES = int(os.getenv("MAX_IMAGE_BYTES", str(15 * 1024 * 1024)))
Image.MAX_IMAGE_PIXELS = int(os.getenv("MAX_IMAGE_PIXELS", "60000000"))

device = "uninitialized"
detectors: list[Detector] = []
load_errors: dict[str, str] = {}
inference_lock = threading.Lock()


def authorize(authorization: str | None) -> None:
    expected = os.getenv("DETECTOR_AUTH_TOKEN", "")
    supplied = (authorization or "").removeprefix("Bearer ").strip()
    if not expected or not secrets.compare_digest(supplied, expected):
        raise HTTPException(status_code=401, detail="Unauthorized")


@asynccontextmanager
async def lifespan(_: FastAPI):
    global device, detectors, load_errors
    device, detectors, load_errors = await asyncio.to_thread(load_detectors)
    yield


app = FastAPI(title="Echo self-hosted AI image detector", version="1.0.0", lifespan=lifespan)


@app.get("/health")
async def health(authorization: str | None = Header(default=None)):
    authorize(authorization)
    ready = {detector.key: {"ready": True, "version": detector.version} for detector in detectors}
    for key, error in load_errors.items():
        ready[key] = {"ready": False, "error": error}
    return {
        "ok": len(detectors) == 2,
        "service_version": app.version,
        "device": device,
        "models": ready,
        "privacy": "public-image-bytes-only",
    }


def _run_inference(image: Image.Image) -> list[dict[str, object]]:
    with inference_lock:
        return [detector.predict(image).as_dict() for detector in detectors]


run_inference = (
    spaces.GPU(duration=120)(_run_inference)
    if spaces is not None and os.getenv("USE_ZEROGPU", "false").lower() == "true"
    else _run_inference
)


@app.post("/v1/detect")
async def detect(
    file: UploadFile = File(...),
    authorization: str | None = Header(default=None),
):
    authorize(authorization)
    if len(detectors) != 2:
        raise HTTPException(status_code=503, detail={"message": "Both detectors must be healthy", "models": load_errors})
    payload = await file.read(MAX_IMAGE_BYTES + 1)
    if not payload or len(payload) > MAX_IMAGE_BYTES:
        raise HTTPException(status_code=413, detail="Image exceeds configured size limit")
    try:
        image = Image.open(io.BytesIO(payload))
        image.verify()
        image = Image.open(io.BytesIO(payload)).convert("RGB")
        image = ImageOps.exif_transpose(image)
    except (UnidentifiedImageError, OSError, ValueError) as exc:
        raise HTTPException(status_code=422, detail="Unsupported or malformed image") from exc
    results = await asyncio.to_thread(run_inference, image)
    return {
        "service_version": app.version,
        "device": device,
        "models": results,
        "policy": {"minimum_models": 2, "service_combines_scores": False},
    }


if spaces is not None:
    import gradio as gr

    with gr.Blocks(title="Echo detector runtime") as platform_status:
        gr.Markdown("# Echo AI image detector\nAuthenticated provenance service. Inference endpoints are not exposed in this panel.")
        platform_output = gr.Textbox(label="Runtime", interactive=False)
        platform_probe = gr.Button("Check ZeroGPU runtime")
        platform_probe.click(zerogpu_runtime_probe, outputs=platform_output)


if __name__ == "__main__":
    if spaces is not None:
        device, detectors, load_errors = load_detectors()
        platform_status.launch(
            server_name="0.0.0.0",
            server_port=int(os.getenv("PORT", "7860")),
            prevent_thread_lock=True,
        )
        # Gradio installs a root catch-all route. Registering API routes normally
        # appends them after it, so the SPA would answer every API request with
        # HTML. Move Echo's routes to the front while keeping native Gradio
        # startup (required for free ZeroGPU Spaces).
        server_app = platform_status.server_app
        existing_route_count = len(server_app.router.routes)
        # A ZeroGPU Space's Node frontend forwards this prefix to Gradio's
        # Python server. Keep the Docker deployment's routes at the root while
        # exposing equivalent Space routes through the forwarded API prefix.
        server_app.add_api_route("/gradio_api/health", health, methods=["GET"])
        server_app.add_api_route("/gradio_api/v1/detect", detect, methods=["POST"])
        echo_routes = server_app.router.routes[existing_route_count:]
        server_app.router.routes = echo_routes + server_app.router.routes[:existing_route_count]
        platform_status.block_thread()
    else:
        import uvicorn

        uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")), workers=1)