#!/usr/bin/env python3
from __future__ import annotations

import os
from functools import lru_cache
from typing import Any

import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

TITLE = "Crack Detection Studio"
DEFAULT_MODEL_REPO_ID = "Mezosky/cracker-yolo26l-baseline"
DEFAULT_MODEL_FILENAME = "best.pt"
DEFAULT_DEVICE = os.getenv("DEVICE", "cpu")
DEFAULT_IMGSZ = int(os.getenv("IMGSZ", "640"))
DEFAULT_CONF = float(os.getenv("CONF", "0.25"))
DEFAULT_IOU = float(os.getenv("IOU", "0.45"))


@lru_cache(maxsize=1)
def get_model() -> tuple[YOLO, str]:
    model_repo_id = os.getenv("MODEL_REPO_ID", DEFAULT_MODEL_REPO_ID)
    model_filename = os.getenv("MODEL_FILENAME", DEFAULT_MODEL_FILENAME)
    model_path = hf_hub_download(repo_id=model_repo_id, filename=model_filename, repo_type="model")
    return YOLO(model_path), model_path


def build_rows(result: Any) -> list[list[Any]]:
    rows: list[list[Any]] = []
    boxes = getattr(result, "boxes", None)
    if boxes is None or len(boxes) == 0:
        return rows

    xyxy = boxes.xyxy.cpu().tolist()
    confs = boxes.conf.cpu().tolist()
    class_ids = boxes.cls.cpu().tolist()
    names = getattr(result, "names", {})

    for idx, (box, conf, class_id) in enumerate(zip(xyxy, confs, class_ids), start=1):
        class_id_int = int(class_id)
        class_name = names.get(class_id_int, f"class_{class_id_int}") if isinstance(names, dict) else str(class_id_int)
        x1, y1, x2, y2 = [round(float(value), 1) for value in box]
        rows.append([idx, class_name, round(float(conf), 4), x1, y1, x2, y2])
    return rows


def infer(image, conf_threshold, iou_threshold, image_size):
    if image is None:
        return None, [], "Upload an image to start detection."

    model, _ = get_model()
    results = model.predict(
        source=image,
        conf=float(conf_threshold),
        iou=float(iou_threshold),
        imgsz=int(image_size),
        device=DEFAULT_DEVICE,
        verbose=False,
    )
    result = results[0]

    rendered = result.plot()
    rendered_rgb = rendered[..., ::-1]
    rows = build_rows(result)

    if not rows:
        summary = "No crack detections found at this confidence threshold."
    else:
        avg_conf = sum(row[2] for row in rows) / len(rows)
        summary = f"Detected **{len(rows)} crack box(es)**. Mean confidence: **{avg_conf:.3f}**."

    return rendered_rgb, rows, summary


def build_demo() -> gr.Blocks:
    css = """
    .gradio-container {
      background: radial-gradient(1200px 600px at 5% 0%, #0f172a 0%, #111827 35%, #030712 100%);
    }
    .hero {
      background: linear-gradient(135deg, #0ea5e9 0%, #22d3ee 45%, #34d399 100%);
      border-radius: 18px;
      padding: 18px 22px;
      color: #00111a;
      box-shadow: 0 12px 30px rgba(34, 211, 238, 0.25);
      margin-bottom: 12px;
    }
    .glass {
      background: rgba(255, 255, 255, 0.06);
      border: 1px solid rgba(255, 255, 255, 0.15);
      border-radius: 14px;
      padding: 10px;
    }
    """

    with gr.Blocks(title=TITLE, css=css) as demo:
        gr.HTML(
            """
            <div class='hero'>
              <h1 style='margin:0;font-size:28px'>Crack Detection Studio</h1>
              <p style='margin:6px 0 0 0;font-size:14px'>Upload one image and get instant crack localization from YOLO26.</p>
            </div>
            """
        )

        gr.Markdown(
            "Model repo: `" + os.getenv("MODEL_REPO_ID", DEFAULT_MODEL_REPO_ID) + "`  \\n"
            "Model file: `" + os.getenv("MODEL_FILENAME", DEFAULT_MODEL_FILENAME) + "`  \\n"
            "Device: `" + DEFAULT_DEVICE + "`"
        )

        with gr.Row():
            with gr.Column(scale=1, elem_classes=["glass"]):
                input_image = gr.Image(type="pil", label="Upload Image")
                conf_slider = gr.Slider(0.05, 0.95, value=DEFAULT_CONF, step=0.01, label="Confidence Threshold")
                iou_slider = gr.Slider(0.10, 0.90, value=DEFAULT_IOU, step=0.01, label="IoU Threshold")
                imgsz_slider = gr.Slider(320, 1280, value=DEFAULT_IMGSZ, step=32, label="Inference Image Size")
                run_button = gr.Button("Detect Cracks", variant="primary")

            with gr.Column(scale=1, elem_classes=["glass"]):
                output_image = gr.Image(type="numpy", label="Predicted Crack Positions")
                output_table = gr.Dataframe(
                    headers=["id", "class", "confidence", "x1", "y1", "x2", "y2"],
                    datatype=["number", "str", "number", "number", "number", "number", "number"],
                    row_count=8,
                    column_count=(7, "fixed"),
                    label="Detections",
                )
                summary_md = gr.Markdown("Upload an image to start detection.")

        run_button.click(
            fn=infer,
            inputs=[input_image, conf_slider, iou_slider, imgsz_slider],
            outputs=[output_image, output_table, summary_md],
        )
        input_image.change(
            fn=infer,
            inputs=[input_image, conf_slider, iou_slider, imgsz_slider],
            outputs=[output_image, output_table, summary_md],
        )
    return demo


demo = build_demo()

if __name__ == "__main__":
    demo.launch()