import os
import sys
import logging
import traceback
import warnings
import re

warnings.filterwarnings("ignore")

import numpy as np
import gradio as gr
from PIL import Image, ImageDraw, ImageFont

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("cocoscan-classifier")
logger.info("Logger initialised.")
logger.info(f"Gradio version: {gr.__version__}")

MODEL_PATH = "best.pt"
model = None

try:
    from ultralytics import YOLO

    if os.path.exists(MODEL_PATH):
        model = YOLO(MODEL_PATH)
        logger.info(f"Model loaded: {MODEL_PATH}")
    else:
        logger.warning(f"best.pt not found at '{MODEL_PATH}'. Running in fallback mode.")
except Exception as e:
    logger.error(f"Failed to load model: {e}\n{traceback.format_exc()}")
    model = None

DEFAULT_CLASSES = [
    "infested by CRB",
    "infestation from other pest",
    "not infested",
    "unspecified",
]

CONFIDENCE_THRESHOLDS = {
    "not infested": 0.50,
    "infested by CRB": 0.45,
    "infestation from other pest": 0.45,
    "unspecified": 0.30,
}

CLASS_COLORS = {
    "not infested": "#2ecc71",
    "infested by CRB": "#e74c3c",
    "infestation from other pest": "#3498db",
    "unspecified": "#95a5a6",
}

# Keep this small and safe:
# - normalize common spacing/case
# - map any legacy/variant labels into the 4 final labels
NAME_NORMALIZATION = {
    "crb infestation": "infested by CRB",
    "crb-infestation": "infested by CRB",
    "crb_infestation": "infested by CRB",
    "infested-by-crb": "infested by CRB",
    "infested_by_crb": "infested by CRB",
    "unhealthy": "infestation from other pest",
    "other-pest-damage": "infestation from other pest",
    "other_pest_damage": "infestation from other pest",
    "healthy": "not infested",
    "not-infested": "not infested",
    "not_infested": "not infested",
    "unknown": "unspecified",
    "uncertain": "unspecified",
    "unclear": "unspecified",
}


def normalize_class_name(name: str) -> str:
    if not isinstance(name, str):
        return str(name)
    key = name.strip().lower()
    key = re.sub(r"\s+", " ", key)
    return NAME_NORMALIZATION.get(key, name.strip())


def health_cascade(probs: dict) -> tuple:
    ranked = sorted(probs.items(), key=lambda x: float(x[1]), reverse=True)
    if not ranked:
        return "not infested", 0.0

    for cls_name, conf in ranked:
        threshold = CONFIDENCE_THRESHOLDS.get(cls_name, 0.30)
        if float(conf) >= float(threshold):
            return cls_name, float(conf)
    return ranked[0]


def multi_run_predict(image: Image.Image, runs: int = 3) -> dict:
    if model is None:
        return {}

    accumulated = {}
    imgsz_list = [224, 256, 192]

    for i in range(runs):
        imgsz = imgsz_list[i % len(imgsz_list)]
        try:
            result = model(image, imgsz=imgsz, verbose=False)[0]
            names = result.names
            probs = result.probs.data.cpu().numpy()

            for idx, prob in enumerate(probs):
                cls_name = names.get(idx, f"class_{idx}")
                cls_name = normalize_class_name(cls_name)
                accumulated[cls_name] = accumulated.get(cls_name, 0.0) + float(prob)
        except Exception as e:
            logger.warning(f"Run {i+1} failed: {e}")
            continue

    if not accumulated:
        return {}

    averaged = {k: v / runs for k, v in accumulated.items()}

    for c in DEFAULT_CLASSES:
        averaged.setdefault(c, 0.0)

    return averaged


def predict_classification(image: Image.Image) -> dict:
    if image is None:
        return {
            "success": False,
            "class": "not infested",
            "confidence": 0.0,
            "all_probs": {},
            "message": "No image provided.",
        }

    image = image.convert("RGB")

    if model is None:
        return {
            "success": True,
            "class": "not infested",
            "confidence": 0.0,
            "all_probs": {c: 0.0 for c in DEFAULT_CLASSES},
            "message": "Model not available. Please put best.pt beside app.py and restart.",
        }

    try:
        avg_probs = multi_run_predict(image, runs=3)
        if not avg_probs:
            raise ValueError("No probabilities returned from model.")

        predicted_class, confidence = health_cascade(avg_probs)
        predicted_class = normalize_class_name(predicted_class)

        logger.info(f"Prediction: {predicted_class} ({confidence:.4f})")

        return {
            "success": True,
            "class": predicted_class,
            "confidence": round(float(confidence), 4),
            "all_probs": {k: round(float(v), 4) for k, v in avg_probs.items()},
            "message": "Classification successful.",
        }
    except Exception as e:
        logger.error(f"Prediction error: {e}\n{traceback.format_exc()}")
        return {
            "success": True,
            "class": "not infested",
            "confidence": 0.0,
            "all_probs": {c: 0.0 for c in DEFAULT_CLASSES},
            "message": f"Prediction failed: {str(e)}",
        }


def _escape_html(s: str) -> str:
    return str(s).replace("&", "&​amp;").replace("<", "&​lt;").replace(">", "&​gt;")


def predict_on_image(input_image):
    if input_image is None:
        blank = Image.new("RGB", (420, 220), color="#1a1a2e")
        draw = ImageDraw.Draw(blank)
        draw.text((80, 100), "Please upload an image.", fill="white")
        return blank, "<div style='color:#fff;'>No image uploaded.</div>"

    if isinstance(input_image, np.ndarray):
        pil_image = Image.fromarray(input_image.astype(np.uint8))
    elif isinstance(input_image, Image.Image):
        pil_image = input_image
    else:
        pil_image = Image.fromarray(np.array(input_image).astype(np.uint8))

    result = predict_classification(pil_image)
    cls_name = normalize_class_name(result["class"])
    confidence = float(result["confidence"])
    all_probs = result.get("all_probs", {}) or {}
    message = result.get("message", "")

    img_display = pil_image.convert("RGB").copy()
    w, h = img_display.size
    draw = ImageDraw.Draw(img_display)

    bar_h = max(50, h // 8)
    bar_color = CLASS_COLORS.get(cls_name, "#888888")
    draw.rectangle([0, h - bar_h, w, h], fill=bar_color)

    label = f"{cls_name.upper()}  {confidence * 100:.1f}%"
    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
            max(14, bar_h // 2),
        )
    except Exception:
        font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), label, font=font)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    text_x = (w - text_w) // 2
    text_y = h - bar_h + (bar_h - text_h) // 2
    draw.text((text_x, text_y), label, fill="white", font=font)

    emoji = {
        "not infested": "✅",
        "infested by CRB": "🪲",
        "infestation from other pest": "🐛",
        "unspecified": "❓",
    }.get(cls_name, "🔍")

    lines = [
        f"{emoji}  Predicted Class : {cls_name.upper()}",
        f"Confidence         : {confidence * 100:.2f}%",
        "",
        "── All Class Probabilities ──",
    ]

    seen = set()
    for c in DEFAULT_CLASSES:
        p = float(all_probs.get(c, 0.0))
        bar = "█" * int(max(0.0, min(1.0, p)) * 20)
        lines.append(f"  {c:<28} {p * 100:5.1f}%  {bar}")
        seen.add(c)

    extras = [(k, v) for k, v in all_probs.items() if k not in seen]
    for c, p in sorted(
        extras, key=lambda x: float(x[1]) if x[1] is not None else 0.0, reverse=True
    ):
        try:
            p = float(p)
        except Exception:
            p = 0.0
        bar = "█" * int(max(0.0, min(1.0, p)) * 20)
        lines.append(f"  {str(c):<28} {p * 100:5.1f}%  {bar}")

    lines += ["", f"Info: {message}"]

    text_color = CLASS_COLORS.get(cls_name, "#ffffff")
    safe_lines = "<br>".join(_escape_html(line) for line in lines)

    html = f"""
<div style="
  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace;
  white-space: normal;
  line-height: 1.35;
  color: {text_color};
">
  {safe_lines}
</div>
"""
    return img_display, html


with gr.Blocks(title="Capstone CocoScan — 4-Class Classifier") as demo:
    gr.HTML(
        """
        <div style="text-align:center; padding:16px 0;">
            <h1>Capstone CocoScan — 4-Class Classifier</h1>
            <p>Upload an image to classify it into:</p>
            <p>
              <b>INFESTED BY CRB</b>,
              <b>INFESTATION FROM OTHER PEST</b>,
              <b>NOT INFESTED</b>,
              <b>UNSPECIFIED</b>
            </p>
            <p style="color:#888; font-size:13px;">
                Model: YOLOv8-cls · 4 classes · Multi-run averaging
            </p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Upload Image",
                type="numpy",
                height=350,
            )
            classify_btn = gr.Button(
                value="Classify",
                variant="primary",
            )

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Result",
                type="pil",
                height=350,
            )
            output_text = gr.HTML(label="Details")

    classify_btn.click(
        fn=predict_on_image,
        inputs=input_image,
        outputs=[output_image, output_text],
    )

    input_image.change(
        fn=predict_on_image,
        inputs=input_image,
        outputs=[output_image, output_text],
    )

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