"""
app.py
──────
Interfaz web Gradio para el sistema RAG de corrección de castellano s.XVI.

Arranque:
    python app.py

Requiere:
    - .env con OPENAI_API_KEY
    - (opcional) corpus en ./corpus/
"""

import os
import json
import random
import gradio as gr
from dotenv import load_dotenv

from knowledge_base import SAMPLE_PAIRS
from corpus_loader import CorpusLoader
from vector_store import VectorStore
from rag_corrector import RAGCorrector
from evaluator import Evaluator

load_dotenv()

# ── Inicialización ────────────────────────────────────────────────────────────

print(" Inicializando Scriptorium RAG...")

loader     = CorpusLoader(os.getenv("CORPUS_PATH", "./corpus"))
disk_pairs = loader.load()
all_pairs  = SAMPLE_PAIRS + disk_pairs

current_embed_model = "openai"
vs = VectorStore(embedding_model="openai")

if vs.count() == 0:
    print(f"⏳ Índice vacío — indexando {len(all_pairs)} pares, espera...")
    vs.index(all_pairs)
    print(f"✓ Índice listo: {vs.count()} documentos")
else:
    print(f"✓ Índice existente cargado: {vs.count()} documentos")

corrector = RAGCorrector(vs)
evaluator = Evaluator()

print(f" Sistema listo. Documentos en vector store: {vs.count()}")

# ── Ejemplos de demostración ──────────────────────────────────────────────────

DEMO_EXAMPLES = [
    "q̃ fizo merçed al dho lugar de las alcaualas del anno de mill e quinientos",
    "el escriuano del cabildo faze fe y da testimouio verdadero de todo lo sobredho",
    "en la muy noble çibdad de burgos a veynte dias del mes de março anno dho",
    "yo juan de la torre vezino desta uilla de toledo otorgo e conosco",
    "sepan quantos esta carta de poder vieren como yo pero lopez vezino dela villa",
    "fizo pareçer ante si a los testigos q̃ dixeron ser mayores de veynte annos",
]

# ── Caligrafías disponibles en el corpus ──────────────────────────────────────

def get_caligrafias():
    cals = set()
    for p in all_pairs:
        c = p.get("caligrafia", "desconocida")
        if c: cals.add(c)
    return ["Todas"] + sorted(cals)


# ── Funciones ─────────────────────────────────────────────────────────────────

def cambiar_embedding(embed_model: str):
    global vs, corrector, current_embed_model
    if embed_model == current_embed_model:
        return f"ℹ Ya estás usando **{embed_model}**"
    try:
        current_embed_model = embed_model
        vs = VectorStore(embedding_model=embed_model)
        if vs.count() == 0:
            vs.index(all_pairs)
            msg = f" Re-indexado con **{embed_model}** · {vs.count()} docs"
        else:
            msg = f" Cargado índice existente **{embed_model}** · {vs.count()} docs"
        corrector = RAGCorrector(vs)
        return msg
    except Exception as e:
        return f" Error cambiando embedding: {e}"


def corregir(htr_text: str, top_k: int, mostrar_prompt: bool, model: str):
    if not htr_text.strip():
        return "", "", "", "", " Introduce un texto HTR para corregir.", ""
    if not os.getenv("OPENAI_API_KEY"):
        return "", "", "", "", " Falta OPENAI_API_KEY en el fichero .env", ""
    try:
        result = corrector.correct(htr_text, top_k=int(top_k), model=model)
    except Exception as e:
        return "", "", "", "", f" Error al llamar a la API: {e}", ""

    corrected  = result["corrected"]
    retrieved  = result["retrieved"]
    htr_errors = result["htr_errors"]
    grafia_w   = result["grafia_warns"]

    docs_md = f"### Top-{len(retrieved)} documentos recuperados\n\n"
    for i, doc in enumerate(retrieved, 1):
        docs_md += (
            f"**{i}. [{doc['type']} · {doc['region']} · {doc['date']}]**  "
            f"*similitud: {doc['score']}*\n\n"
            f"- **HTR:** `{doc['htr']}`\n"
            f"- **GT:**  `{doc['gt']}`\n"
        )
        if doc["corrections"]:
            docs_md += f"- **Correcciones:** {', '.join(doc['corrections'])}\n"
        docs_md += "\n---\n"

    analysis_md = "### Análisis del texto\n\n"
    if htr_errors:
        analysis_md += "**⚠ Posibles errores HTR detectados:**\n"
        for e in htr_errors:
            examples = e.get("examples", e.get("example", ""))
            if isinstance(examples, list):
                examples = "; ".join(examples[:2])
            sev = e.get("severity", "")
            sev_label = f" `[{sev.upper()}]`" if sev else ""
            analysis_md += f"- `{e['htr']}` → `{e['gt']}`{sev_label}: {e['context']}  \n  *Ej: {examples}*\n"
        analysis_md += "\n"
    if grafia_w:
        analysis_md += "**✦ Alertas de grafía (NO modernizar):**\n"
        for g in grafia_w:
            analysis_md += f"- `{g['modern']}` → mantener `{g['ancient']}`: {g['rule']}\n"
        analysis_md += "\n"
    if not htr_errors and not grafia_w:
        analysis_md += "*No se detectaron patrones conocidos de error en el texto.*\n"

    diff_md = "### Diferencias HTR → Corregido\n\n"
    orig_words = htr_text.split()
    corr_words = corrected.split()
    diff_parts = []
    max_len = max(len(orig_words), len(corr_words))
    changed = 0
    for i in range(max_len):
        o = orig_words[i] if i < len(orig_words) else "—"
        c = corr_words[i] if i < len(corr_words) else "—"
        if o != c:
            diff_parts.append(f"~~{o}~~ → **{c}**")
            changed += 1
        else:
            diff_parts.append(c)
    diff_md += " ".join(diff_parts)
    diff_md += f"\n\n*{changed} palabra(s) modificada(s) de {len(orig_words)} totales.*"

    status = f" Corrección completada con **{result['model']}** · {vs.count()} docs en índice"

    prompt_visible = ""
    if mostrar_prompt:
        prompt_visible = (
            "### System Prompt\n\n"
            f"```\n{result.get('_system', '(ver rag_corrector.py)')}\n```\n\n"
            "### User Prompt (dinámico)\n\n"
            f"```\n{result['prompt']}\n```"
        )

    return corrected, docs_md, analysis_md, diff_md, status, prompt_visible


def evaluar_par(htr_text: str, gt_text: str):
    if not htr_text.strip() or not gt_text.strip():
        return "⚠ Introduce tanto el texto HTR como el groundtruth."
    try:
        result  = corrector.correct(htr_text)
        metrics = evaluator.evaluate_pair(htr_text, result["corrected"], gt_text)
        report  = evaluator.format_pair_report(metrics)
        report += f"\n\n**Texto corregido por RAG:**\n> {result['corrected']}"
        return report
    except Exception as e:
        return f" Error: {e}"


def evaluar_batch(n_samples: int, caligrafia_filtro: str, model: str):
    """
    Evalúa el sistema sobre N pares aleatorios del corpus que tengan GT.
    Muestra las tres comparaciones: GT vs HTR, GT vs Corregido, HTR vs Corregido.
    """
    # Filtrar pares con HTR y GT no vacíos
    pares_validos = [
        p for p in all_pairs
        if p.get("htr", "").strip() and p.get("gt", "").strip()
    ]

    # Filtrar por caligrafía si se especifica
    if caligrafia_filtro and caligrafia_filtro != "Todas":
        pares_validos = [
            p for p in pares_validos
            if p.get("caligrafia", "") == caligrafia_filtro
        ]

    if not pares_validos:
        return "⚠ No hay pares con GT disponibles para evaluar con ese filtro."

    # Muestra aleatoria
    n = min(int(n_samples), len(pares_validos))
    muestra = random.sample(pares_validos, n)

    yield f"⏳ Evaluando {n} pares con modelo **{model}**...\n\n"

    results = []
    errores = []
    for i, pair in enumerate(muestra, 1):
        try:
            out     = corrector.correct(pair["htr"], model=model)
            metrics = evaluator.evaluate_pair(
                htr=pair["htr"],
                corrected=out["corrected"],
                gt=pair["gt"],
            )
            metrics["id"]        = pair.get("id", f"par_{i}")
            metrics["htr"]       = pair["htr"]
            metrics["corrected"] = out["corrected"]
            metrics["gt"]        = pair["gt"]
            metrics["caligrafia"]= pair.get("caligrafia", "desconocida")
            results.append(metrics)
        except Exception as e:
            errores.append(f"  - {pair.get('id','?')}: {e}")

        # Progreso intermedio cada 5 pares
        if i % 5 == 0:
            yield f"⏳ Procesados {i}/{n}...\n\n"

    if not results:
        yield "❌ No se obtuvieron resultados.\n" + "\n".join(errores)
        return

    # ── Resumen global ────────────────────────────────────────────────────────
    def avg(key):
        return sum(r[key] for r in results) / len(results)

    n_res       = len(results)
    mejorados   = sum(1 for r in results if r["cer_improvement"] >  0.02)
    empeorados  = sum(1 for r in results if r["cer_improvement"] < -0.02)
    sin_cambio  = n_res - mejorados - empeorados

    md  = f"## 📊 Evaluación por lotes — {n_res} pares\n\n"
    if caligrafia_filtro != "Todas":
        md += f"**Filtro caligrafía:** {caligrafia_filtro}  |  "
    md += f"**Modelo:** {model}\n\n"
    md += "---\n\n"

    # Comparación 1 — GT vs HTR
    md += "### ① Error de partida  (GT vs HTR original)\n\n"
    md += f"| CER medio | WER medio |\n|---|---|\n"
    md += f"| {avg('cer_before'):.2%} | {avg('wer_before'):.2%} |\n\n"

    # Comparación 2 — GT vs Corregido
    md += "### ② Error final  (GT vs Texto corregido)\n\n"
    md += f"| CER medio | WER medio | Mejora CER | Mejora WER |\n|---|---|---|---|\n"
    md += (
        f"| {avg('cer_after'):.2%} "
        f"| {avg('wer_after'):.2%} "
        f"| {avg('cer_improvement'):+.2%} "
        f"| {avg('wer_improvement'):+.2%} |\n\n"
    )
    md += (
        f"| ✓ Mejorados | ✗ Empeorados | ~ Sin cambio |\n|---|---|---|\n"
        f"| {mejorados} ({mejorados/n_res:.0%}) "
        f"| {empeorados} ({empeorados/n_res:.0%}) "
        f"| {sin_cambio} ({sin_cambio/n_res:.0%}) |\n\n"
    )

    # Comparación 3 — HTR vs Corregido (modernismos)
    md += "### ③ Modernismos introducidos  (HTR vs Corregido)\n\n"
    md += f"Score promedio: **{avg('modernism_score'):.2%}** (1.0 = sin modernismos)\n\n"
    modernismos_total = sum(r["modernism"]["count"] for r in results)
    if modernismos_total == 0:
        md += "✓ El LLM no introdujo modernismos en ningún par.\n\n"
    else:
        md += f"✗ {modernismos_total} modernismo(s) introducidos en total.\n\n"

    # ── Desglose por caligrafía ───────────────────────────────────────────────
    from collections import defaultdict
    by_cal = defaultdict(list)
    for r in results:
        by_cal[r["caligrafia"]].append(r)

    if len(by_cal) > 1:
        md += "### Desglose por caligrafía\n\n"
        md += "| Caligrafía | N | CER antes | CER después | Mejora CER | Modernismos |\n"
        md += "|---|---|---|---|---|---|\n"
        for cal, rs in sorted(by_cal.items()):
            a_cer  = sum(r["cer_before"] for r in rs) / len(rs)
            d_cer  = sum(r["cer_after"]  for r in rs) / len(rs)
            imp    = a_cer - d_cer
            mods   = sum(r["modernism"]["count"] for r in rs)
            md += f"| {cal} | {len(rs)} | {a_cer:.2%} | {d_cer:.2%} | {imp:+.2%} | {mods} |\n"
        md += "\n"

    # ── Detalle por par ───────────────────────────────────────────────────────
    md += "### Detalle por par\n\n"
    md += "| ID | Cal | CER antes | CER después | Mejora | Veredicto | Modernismos |\n"
    md += "|---|---|---|---|---|---|---|\n"
    for r in results:
        md += (
            f"| `{r['id'][:25]}` "
            f"| {r['caligrafia'][:12]} "
            f"| {r['cer_before']:.2%} "
            f"| {r['cer_after']:.2%} "
            f"| {r['cer_improvement']:+.2%} "
            f"| {r['verdict']} "
            f"| {r['modernism']['count']} |\n"
        )

    # Errores
    if errores:
        md += f"\n\n⚠ {len(errores)} pares fallaron:\n" + "\n".join(errores)

    yield md


def add_to_corpus(htr_text, gt_text, doc_type, region, date, caligrafia):
    if not htr_text.strip() or not gt_text.strip():
        return "⚠ HTR y GT son obligatorios."
    try:
        pair_id  = f"user_{abs(hash(htr_text)) % 100000:05d}"
        new_pair = {
            "id":          pair_id,
            "htr":         htr_text.strip(),
            "gt":          gt_text.strip(),
            "type":        doc_type or "desconocido",
            "region":      region   or "desconocida",
            "date":        date     or "",
            "caligrafia":  caligrafia or "desconocida",
            "corrections": [],
            "source":      "user_added",
        }
        added = vs.index([new_pair])
        if added:
            return f" Par añadido al corpus con id `{pair_id}`. Total: {vs.count()} docs."
        else:
            return f" Par ya existía en el corpus (id: `{pair_id}`)."
    except Exception as e:
        return f" Error: {e}"


# ── Interfaz Gradio ───────────────────────────────────────────────────────────

with gr.Blocks(
    title="Scriptorium RAG",
    theme=gr.themes.Base(
        primary_hue="amber",
        secondary_hue="stone",
        neutral_hue="stone",
        font=gr.themes.GoogleFont("IM Fell English"),
    ),
    css="""
    .header { text-align: center; padding: 20px 0 10px; }
    .header h1 { font-size: 2.2em; color: #92400e; letter-spacing: 0.15em; }
    .header p  { color: #78716c; font-style: italic; }
    .status-bar { font-size: 0.85em; padding: 6px 12px; border-radius: 6px; }
    """,
) as demo:

    gr.HTML("""
    <div class="header">
      <h1>RAG CODEX for Historical Spanish</h1>
      <p>RAG system of Spanish correction from the 16th century</p>
    </div>
    """)

    with gr.Tabs():

        # ── Pestaña 1: Corrección ─────────────────────────────────────────────
        with gr.TabItem(" HTR Correction"):
            with gr.Row():
                with gr.Column(scale=2):
                    htr_input = gr.Textbox(
                        label="HTR text (recognizer input)",
                        placeholder="Paste the HTR result here…",
                        lines=6,
                    )
                    with gr.Row():
                        top_k_slider = gr.Slider(
                            minimum=1, maximum=10, value=5, step=1,
                            label="Documents retrieved (k)",
                        )
                        model_selector = gr.Dropdown(
                            label="Modelo LLM",
                            choices=["llama-3.3-70b-versatile", "openai/gpt-oss-120b"],
                            value="llama-3.3-70b-versatile",
                        )
                        embedding_selector = gr.Dropdown(
                            label="Modelo de Embedding",
                            choices=["openai", "mpnet", "mt5"],
                            value="openai",
                        )
                        show_prompt = gr.Checkbox(label="Show RAG prompt", value=False)
                    btn_corregir = gr.Button("✦ Correct with RAG", variant="primary")
                    gr.Examples(examples=DEMO_EXAMPLES, inputs=htr_input, label="Demonstration examples")

                with gr.Column(scale=2):
                    corrected_out = gr.Textbox(label="Corrected text (RAG output)", lines=6, interactive=False)
                    status_out    = gr.Markdown(elem_classes=["status-bar"])

            with gr.Row():
                with gr.Column():
                    docs_out     = gr.Markdown(label="Documents recovered from the corpus")
                with gr.Column():
                    analysis_out = gr.Markdown(label="Pattern analysis")

            diff_out   = gr.Markdown(label="Word-by-word differences")
            prompt_out = gr.Markdown(label="Prompt sent to the LLM", visible=True)

            btn_corregir.click(
                fn=corregir,
                inputs=[htr_input, top_k_slider, show_prompt, model_selector],
                outputs=[corrected_out, docs_out, analysis_out, diff_out, status_out, prompt_out],
            )
            embed_status = gr.Markdown()
            embedding_selector.change(fn=cambiar_embedding, inputs=[embedding_selector], outputs=[embed_status])

        # ── Pestaña 2: Evaluación individual ─────────────────────────────────
        with gr.TabItem(" Evaluation with GT"):
            gr.Markdown(
                "Compara la corrección del RAG con el groundtruth real para medir "
                "CER/WER y detectar modernismos introducidos por el LLM."
            )
            with gr.Row():
                eval_htr = gr.Textbox(label="HTR text", lines=4)
                eval_gt  = gr.Textbox(label="Groundtruth (reference)", lines=4)
            btn_eval = gr.Button("Evaluate", variant="primary")
            eval_out = gr.Markdown()
            btn_eval.click(fn=evaluar_par, inputs=[eval_htr, eval_gt], outputs=eval_out)

        # ── Pestaña 3: Evaluación por lotes ──────────────────────────────────
        with gr.TabItem(" Batch Evaluation"):
            gr.Markdown(
                "Evalúa el sistema sobre una muestra aleatoria del corpus. "
                "Muestra las tres comparaciones: **GT vs HTR** (error de partida), "
                "**GT vs Corregido** (error final) y **HTR vs Corregido** (modernismos)."
            )
            with gr.Row():
                batch_n = gr.Slider(
                    minimum=5, maximum=100, value=20, step=5,
                    label="Número de pares a evaluar",
                )
                batch_cal = gr.Dropdown(
                    label="Filtrar por caligrafía",
                    choices=get_caligrafias(),
                    value="Todas",
                )
                batch_model = gr.Dropdown(
                    label="Modelo LLM",
                    choices=["llama-3.3-70b-versatile", "openai/gpt-oss-120b"],
                    value="llama-3.3-70b-versatile",
                )

            with gr.Row():
                gr.Markdown(
                    f"ℹ Corpus disponible: **{len([p for p in all_pairs if p.get('gt','').strip()])} "
                    f"pares con GT** de {len(all_pairs)} totales."
                )

            btn_batch = gr.Button("▶ Ejecutar evaluación por lotes", variant="primary")
            batch_out = gr.Markdown()

            btn_batch.click(
                fn=evaluar_batch,
                inputs=[batch_n, batch_cal, batch_model],
                outputs=batch_out,
            )

        # ── Pestaña 4: Añadir al corpus ───────────────────────────────────────
        with gr.TabItem("➕ Add to corpus"):
            gr.Markdown("Add new pairs to the vector store to improve the RAG continuously.")
            with gr.Row():
                add_htr = gr.Textbox(label="Texto HTR", lines=4)
                add_gt  = gr.Textbox(label="Groundtruth corregido", lines=4)
            with gr.Row():
                add_type      = gr.Textbox(label="Document type", placeholder="notarial / judicial / eclesiastico")
                add_region    = gr.Textbox(label="Region", placeholder="Castilla, Andalucía…")
                add_date      = gr.Textbox(label="Date", placeholder="1542")
                add_caligrafia = gr.Dropdown(
                    label="Caligrafía",
                    choices=["desconocida", "Procesal", "Encadenada", "Italica_cursiva", "Redonda"],
                    value="desconocida",
                )
            btn_add = gr.Button("Add to corpus", variant="primary")
            add_out = gr.Markdown()
            btn_add.click(
                fn=add_to_corpus,
                inputs=[add_htr, add_gt, add_type, add_region, add_date, add_caligrafia],
                outputs=add_out,
            )

        # ── Pestaña 5: Info del sistema ───────────────────────────────────────
        with gr.TabItem("INFO: System"):
            gr.Markdown(f"""
## System status

- **Modelo LLM:** {os.getenv('OPENAI_MODEL', 'gpt-4o')}
- **Vector store:** ChromaDB (persistente en `{os.getenv('CHROMA_PATH','./chroma_db')}`)
- **Documentos indexados:** {vs.count()}
- **Corpus cargado desde disco:** {len(disk_pairs)} pares
- **Pares de ejemplo embebidos:** {len(SAMPLE_PAIRS)}

## Arquitectura

```
Texto HTR
   │
   ├─► Detector de patrones HTR (knowledge_base.py)
   ├─► Detector de grafías modernas (knowledge_base.py)
   │
   ├─► Embedding (text-embedding-3-small)
   │        │
   │        └─► Búsqueda top-k en ChromaDB ──► Few-shot dinámico
   │
   └─► Prompt constructor ──► LLM ──► Texto corregido
```

## Evaluación por lotes

Las tres comparaciones:
1. **GT vs HTR** — error de partida (cuánto se equivocó el HTR)
2. **GT vs Corregido** — error final (cuánto mejoró el RAG)
3. **HTR vs Corregido** — modernismos (qué cambió el LLM que no debía)
""")

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        auth=("admin", "admin"),
        share=False,
        show_error=True,
    )