from __future__ import annotations

import os
import sys
from functools import lru_cache
from pathlib import Path
from typing import Dict, List

import gradio as gr
from PIL import Image

ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.inference import ClipPredictor, CustomModelPredictor, OpenAIVisionPredictor


labels = ["charizard", "charmander", "charmeleon", "ditto", "eevee", "ekans"]


@lru_cache(maxsize=1)
def get_predictors() -> tuple[CustomModelPredictor, ClipPredictor, OpenAIVisionPredictor]:
    custom = CustomModelPredictor(str(ROOT / "models" / "custom_resnet18.pth"))
    predictor_labels = custom.labels if custom.available() else labels
    clip_model = ClipPredictor(predictor_labels)
    openai_model = OpenAIVisionPredictor(predictor_labels)
    return custom, clip_model, openai_model


def _format_preds(result: Dict[str, object]) -> str:
    if not result.get("available", False):
        return f"Unavailable: {result.get('error', 'unknown error')}"

    lines: List[str] = []
    top = result.get("top_prediction", {})
    label = top.get("label", "-")
    confidence = float(top.get("confidence", 0.0))
    lines.append(f"Top prediction: {label} ({confidence:.2%})")

    for pred in result.get("predictions", []):
        lines.append(f"- {pred['label']}: {pred['confidence']:.2%}")

    raw_response = result.get("raw_response")
    if isinstance(raw_response, dict) and raw_response.get("reason"):
        lines.append(f"Reason: {raw_response['reason']}")

    return "\n".join(lines)


def classify_image(image: Image.Image):
    if image is None:
        return "No image provided.", "No image provided.", "No image provided."

    custom, clip_model, openai_model = get_predictors()
    custom_pred = custom.predict(image)
    clip_pred = clip_model.predict(image)
    openai_pred = openai_model.predict(image)

    return _format_preds(custom_pred), _format_preds(clip_pred), _format_preds(openai_pred)


def get_examples() -> List[List[str]]:
    examples_dir = ROOT / "app" / "examples"
    if not examples_dir.exists():
        return []
    image_paths = sorted(
        [p for p in examples_dir.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}]
    )
    return [[str(p)] for p in image_paths]


description = """
Upload an image and compare predictions from three models:
1) Custom transfer learning model (ResNet18)
2) Open-source CLIP model
3) Closed-source OpenAI vision model

If OPENAI_API_KEY is not set, OpenAI predictions are shown as unavailable.
"""

with gr.Blocks(title="Computer Vision Model Comparison") as demo:
    gr.Markdown("# Computer Vision Classification & Model Comparison")
    gr.Markdown(description)

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload image")

    classify_button = gr.Button("Classify")

    with gr.Row():
        custom_output = gr.Textbox(label="Custom Transfer Learning", lines=8)
        clip_output = gr.Textbox(label="Open-Source CLIP", lines=8)
        openai_output = gr.Textbox(label="Closed-Source OpenAI Vision", lines=8)

    classify_button.click(classify_image, inputs=[image_input], outputs=[custom_output, clip_output, openai_output])

    gr.Examples(examples=get_examples(), inputs=image_input, label="Example images")

if __name__ == "__main__":
    in_space = bool(os.getenv("SPACE_ID"))
    if in_space:
        demo.launch(ssr_mode=False)
    else:
        print("Launching local app on http://127.0.0.1:7860")
        demo.launch(server_name="127.0.0.1", server_port=7860, share=False, ssr_mode=False)