import os
import io
import base64
import time

import gradio as gr
import numpy as np
from PIL import Image
from ultralytics import YOLO
from huggingface_hub import hf_hub_download

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
import uvicorn

# ── Model ──────────────────────────────────────────────────────────────────────
MODEL_REPO = "DannyLuna/recaptcha-classification-57k"
MODEL_FILE = "recaptcha_classification_57k.onnx"
MODEL_PATH = os.path.join("models", MODEL_FILE)

os.makedirs("models", exist_ok=True)

if not os.path.exists(MODEL_PATH):
    print(f"[INFO] Downloading {MODEL_FILE} from {MODEL_REPO}...")
    hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, local_dir="models")
    print("[INFO] Download complete.")

model = YOLO(MODEL_PATH, task="classify")
print(f"[INFO] Model loaded — {len(model.names)} classes")

TARGET_CLASSES = {
    "bicycle", "bridge", "bus", "car", "chimney", "crosswalk",
    "fire hydrant", "motorcycle", "mountain", "palm tree",
    "stairs", "tractor", "traffic light",
}

# ── Shared inference logic ─────────────────────────────────────────────────────
def run_inference(image: Image.Image) -> dict:
    t0 = time.perf_counter()
    results = model(image, verbose=False)
    elapsed = time.perf_counter() - t0

    probs     = results[0].probs
    top1_idx  = int(probs.top1)
    top1_name = model.names[top1_idx]
    top1_conf = float(probs.top1conf)

    all_scores = {
        model.names[i]: round(float(probs.data[i]), 6)
        for i in range(len(model.names))
    }
    top5 = sorted(all_scores.items(), key=lambda x: x[1], reverse=True)[:5]

    return {
        "predicted_class":  top1_name,
        "confidence":       round(top1_conf, 6),
        "is_target":        top1_name in TARGET_CLASSES,
        "top5":             [{"class": c, "confidence": s} for c, s in top5],
        "all_scores":       all_scores,
        "inference_ms":     round(elapsed * 1000, 2),
    }

# ── FastAPI ────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="reCAPTCHA Classifier API",
    description=(
        "YOLO-based image classification fine-tuned on 57k reCAPTCHA tiles.\n\n"
        "**14 classes:** bicycle, bridge, bus, car, chimney, crosswalk, "
        "fire hydrant, motorcycle, mountain, palm tree, stairs, tractor, "
        "traffic light, other."
    ),
    version="1.0.0",
)


@app.get("/api/health")
def health():
    """Quick liveness check."""
    return {"status": "ok", "model": MODEL_FILE, "classes": len(model.names)}


@app.get("/api/classes")
def classes():
    """Return all class names the model knows about."""
    return {
        "all":    list(model.names.values()),
        "target": sorted(TARGET_CLASSES),
        "other":  ["other"],
    }


@app.post("/api/predict")
async def predict_file(file: UploadFile = File(...)):
    """
    Classify an uploaded image file.

    - **file**: JPEG / PNG image to classify.

    Returns the predicted class, confidence, top-5 scores, and inference time.
    """
    if not file.content_type.startswith("image/"):
        raise HTTPException(status_code=415, detail="File must be an image (JPEG/PNG).")

    raw = await file.read()
    try:
        image = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode image.")

    result = run_inference(image)
    return JSONResponse(content=result)


@app.post("/api/predict/base64")
async def predict_base64(payload: dict):
    """
    Classify an image sent as a base64-encoded string.

    Body JSON:
    ```json
    { "image": "<base64 string>" }
    ```
    Optionally wrap with data-URI prefix (`data:image/jpeg;base64,...`).
    """
    b64 = payload.get("image", "")
    if not b64:
        raise HTTPException(status_code=422, detail="Missing 'image' field.")

    # Strip optional data-URI prefix
    if "," in b64:
        b64 = b64.split(",", 1)[1]

    try:
        raw   = base64.b64decode(b64)
        image = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid base64 image data.")

    result = run_inference(image)
    return JSONResponse(content=result)


# ── Gradio UI ─────────────────────────────────────────────────────────────────
def gradio_predict(image: Image.Image):
    if image is None:
        return {}, "⚠️ No image provided."

    r = run_inference(image)
    emoji  = "✅" if r["is_target"] else "🚫"
    status = (
        f"{emoji} **{r['predicted_class'].upper()}** — {r['confidence']:.1%} confidence"
        f"\n\n⏱ Inference: {r['inference_ms']} ms"
        + ("" if r["is_target"] else "\n\n*(classified as `other` / background)*")
    )
    conf_dict = {item["class"]: item["confidence"] for item in r["top5"]}
    return conf_dict, status


CLASSES_MD = ", ".join(f"`{c}`" for c in sorted(TARGET_CLASSES))

with gr.Blocks(title="reCAPTCHA Classifier", theme=gr.themes.Soft()) as ui:
    gr.Markdown(
        f"""
# 🔍 reCAPTCHA Image Classifier
YOLO fine-tuned on **57k reCAPTCHA tiles** — also available as a REST API.

**Target classes:** {CLASSES_MD} + `other`

---
### API endpoints
| Method | Path | Description |
|--------|------|-------------|
| GET  | `/api/health` | Liveness check |
| GET  | `/api/classes` | List all classes |
| POST | `/api/predict` | Upload image file |
| POST | `/api/predict/base64` | Send base64 image |

Full docs → [**Swagger UI**](/docs) · [**ReDoc**](/redoc)
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(type="pil", label="Upload tile image")
            run_btn   = gr.Button("🚀 Classify", variant="primary")
        with gr.Column(scale=1):
            status_out = gr.Markdown(label="Result")
            label_out  = gr.Label(num_top_classes=5, label="Top-5 probabilities")

    run_btn.click(fn=gradio_predict, inputs=img_input, outputs=[label_out, status_out])
    img_input.change(fn=gradio_predict, inputs=img_input, outputs=[label_out, status_out])


# Mount Gradio on FastAPI at root
app = gr.mount_gradio_app(app, ui, path="/")

# ── Entry point ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)