from __future__ import annotations

import os
import tempfile
from pathlib import Path

import gradio as gr

from ig_classifier.model_hub import load_model, model_version_string
from ig_classifier.pipeline import run_file

MODEL_ID = os.environ.get("IG_CLASSIFIER_MODEL", "bravo-pena/ig-classifier-1.0")

_MODEL = None
_META: dict = {}

DEPTH_CHOICES = [
    ("1 — Rule type (AGG / BOU / CHO / INF / PAY / POS / SCO)", 1),
    ("2 — Subtype", 2),
    ("3 — Category", 3),
    ("4 — Full detail (recommended)", 4),
]


def get_model():
    global _MODEL, _META
    if _MODEL is None:
        _MODEL, _META = load_model(MODEL_ID)
    return _MODEL, _META


def classify(file_obj, technical_obj, depth):
    if file_obj is None:
        raise gr.Error("Upload the Paso 1 Excel first (the IS Identifier output).")
    input_path = Path(file_obj.name)
    if input_path.suffix.lower() != ".xlsx":
        raise gr.Error("Expected a .xlsx file — the Excel produced by IS Identifier (Paso 1).")
    technical_path = None
    if technical_obj is not None:
        technical_path = Path(technical_obj.name)
        if technical_path.suffix.lower() != ".json":
            raise gr.Error("The technical sidecar must be the *_technical.json file from Paso 1.")

    model, meta = get_model()
    with tempfile.TemporaryDirectory() as tmpdir:
        output_path = Path(tmpdir) / f"{input_path.stem}_paso2.xlsx"
        try:
            written = run_file(
                input_path=input_path,
                output_path=output_path,
                model=model,
                model_version=model_version_string(MODEL_ID, meta),
                depth=int(depth),
                technical_path=technical_path,
            )
        except ValueError as exc:
            raise gr.Error(
                f"{exc} — is this the Paso 1 Excel (sheet 'segments' with AIM candidates)?"
            ) from exc
        final_path = Path(tempfile.gettempdir()) / written.name
        final_path.write_bytes(written.read_bytes())
        return str(final_path)


with gr.Blocks(title="IG Classifier — Paso 2") as demo:
    gr.Markdown("# IG Classifier — Paso 2 (interim model)")
    gr.Markdown(
        "Upload the Excel produced by "
        "[IS Identifier (Paso 1)](https://huggingface.co/spaces/bravo-pena/is-identifier) "
        "and download it back with the official Rules-taxonomy applied to every "
        "institutional-statement candidate: TYPE / TAXON / LINK.typ with "
        "per-class probabilities (e.g. `PAY 81% | CHO 12% | BOU 4%`).\n\n"
        "Sheets: `planilla` (team-style wide, one row per segment), `aims` "
        "(one row per AIM with probabilities), `schema`, `summary`. Amber rows "
        "carry `needs_review` flags.\n\n"
        "⚠️ **Interim model**: trained while the annotation base is still being "
        "completed (several taxa have very few examples). Depth 1–2 aggregates "
        "are reliable; deeper taxa are suggestions to review, not final codes."
    )
    file_input = gr.File(label="Paso 1 Excel (.xlsx)", file_types=[".xlsx"])
    technical_input = gr.File(
        label="Optional: Paso 1 technical sidecar (*_technical.json) — improves "
              "verb normalization",
        file_types=[".json"],
    )
    depth_input = gr.Dropdown(
        label="Taxonomy depth",
        choices=[label for label, _ in DEPTH_CHOICES],
        value=DEPTH_CHOICES[3][0],
    )
    run_button = gr.Button("Run Paso 2", variant="primary")
    output_file = gr.File(label="Excel result (Paso 2)")

    def _run(file_obj, technical_obj, depth_label):
        depth = dict(DEPTH_CHOICES)[depth_label]
        return classify(file_obj, technical_obj, depth)

    run_button.click(_run, inputs=[file_input, technical_input, depth_input],
                     outputs=output_file)


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