import base64
import json
import mimetypes
import os
from pathlib import Path
from typing import Dict, List, Optional

import gradio as gr
from dotenv import load_dotenv
from transformers import pipeline

try:
    from openai import OpenAI
except ModuleNotFoundError:
    OpenAI = None

load_dotenv()

# Configure environment variables for easy Hugging Face Space setup.
CUSTOM_MODEL_REPO_ID = os.getenv("CUSTOM_MODEL_REPO_ID", "your-hf-username/sports-vit-transfer")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TOP_K = int(os.getenv("TOP_K", "5"))

openai_client = OpenAI(api_key=OPENAI_API_KEY) if OpenAI is not None and OPENAI_API_KEY else None

# Extended label set for zero-shot and OpenAI comparison.
EXTENDED_SPORT_LABELS = [
    "football",
    "tennis",
    "golf",
    "baseball",
    "basketball",
    "volleyball",
    "handball",
    "rugby",
    "cricket",
    "hockey",
    "table tennis",
    "badminton",
    "swimming",
    "skiing",
    "surfing",
    "boxing",
    "athletics",
    "american football",
]

_custom_classifier = None
_clip_classifier = None


def _parse_json_from_text(text: str) -> Optional[Dict]:
    cleaned = text.strip()

    if cleaned.startswith("```"):
        lines = cleaned.splitlines()
        if len(lines) >= 3:
            cleaned = "\n".join(lines[1:-1]).strip()

    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        pass

    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start != -1 and end != -1 and end > start:
        candidate = cleaned[start : end + 1]
        try:
            return json.loads(candidate)
        except json.JSONDecodeError:
            return None

    return None


def _load_custom_classifier():
    global _custom_classifier
    if _custom_classifier is None:
        _custom_classifier = pipeline(
            "image-classification",
            model=CUSTOM_MODEL_REPO_ID,
            top_k=TOP_K,
        )
    return _custom_classifier


def _load_clip_classifier():
    global _clip_classifier
    if _clip_classifier is None:
        _clip_classifier = pipeline(
            task="zero-shot-image-classification",
            model="openai/clip-vit-large-patch14",
        )
    return _clip_classifier


def _to_score_dict(results: List[Dict]) -> Dict[str, float]:
    return {row["label"]: round(float(row["score"]), 6) for row in results}


def _encode_image(image_path: str) -> str:
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def classify_with_openai(image_path: str) -> Dict:
    if OpenAI is None:
        return {
            "error": "OpenAI package is not installed in this environment.",
            "hint": "Add 'openai' to requirements.txt and redeploy the Space.",
        }

    if openai_client is None:
        return {
            "error": "OPENAI_API_KEY fehlt. Setze es in den Space Secrets.",
            "hint": "In Hugging Face Space unter Settings -> Secrets als OPENAI_API_KEY eintragen.",
        }

    mime_type, _ = mimetypes.guess_type(image_path)
    mime_type = mime_type or "image/jpeg"
    base64_image = _encode_image(image_path)

    prompt = (
        "Classify the sport shown in this image. "
        f"Choose exactly one label from this list: {', '.join(EXTENDED_SPORT_LABELS)}. "
        "Return valid JSON only with keys: label, confidence, reasoning. "
        "confidence must be a number between 0 and 1."
    )

    response = openai_client.responses.create(
        model=OPENAI_MODEL,
        input=[
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": prompt},
                    {
                        "type": "input_image",
                        "image_url": f"data:{mime_type};base64,{base64_image}",
                    },
                ],
            }
        ],
    )

    text = response.output_text.strip()
    parsed = _parse_json_from_text(text)
    if parsed is None:
        parsed = {
            "warning": "OpenAI antwortete nicht als valides JSON.",
            "raw_response": text,
        }
    return parsed


def classify_sport(image_path: Optional[str]):
    if image_path is None:
        return {"error": "Bitte ein Bild hochladen."}

    output = {}

    if CUSTOM_MODEL_REPO_ID == "your-hf-username/sports-vit-transfer":
        output["Custom Transfer Model"] = {
            "error": "CUSTOM_MODEL_REPO_ID ist noch der Platzhalter.",
            "hint": "Setze CUSTOM_MODEL_REPO_ID in Space Variables auf dein HF-Modellrepo.",
        }
    else:
        try:
            custom_results = _load_custom_classifier()(image_path)
            output["Custom Transfer Model"] = _to_score_dict(custom_results)
        except Exception as exc:  # noqa: BLE001
            output["Custom Transfer Model"] = {"error": f"Custom model failed: {exc}"}

    try:
        clip_results = _load_clip_classifier()(image_path, candidate_labels=EXTENDED_SPORT_LABELS)
        output["CLIP Zero-Shot"] = _to_score_dict(clip_results)
    except Exception as exc:  # noqa: BLE001
        output["CLIP Zero-Shot"] = {"error": f"CLIP failed: {exc}"}

    output["OpenAI Vision"] = classify_with_openai(image_path)
    return output


def _collect_examples() -> List[List[str]]:
    examples_dir = Path("example_images")
    if not examples_dir.exists():
        return []

    allowed_ext = {".jpg", ".jpeg", ".png", ".webp"}
    paths = sorted([p for p in examples_dir.iterdir() if p.suffix.lower() in allowed_ext])
    return [[str(p)] for p in paths]


iface = gr.Interface(
    fn=classify_sport,
    inputs=gr.Image(type="filepath", label="Sportbild hochladen"),
    outputs=gr.JSON(label="Modellvergleich"),
    title="Sport Classification Comparison",
    examples=_collect_examples(),
)


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