import re
from functools import lru_cache
from pathlib import Path

import gradio as gr
from pie.tagger import Tagger, simple_tokenizer

MODELS_DIR = Path(__file__).parent / "models"
MAX_WORDS = 1000

TASK_LABELS = {
    "CAS": "Cas",
    "DEGRE": "Degré",
    "GENRE": "Genre",
    "MODE": "Mode",
    "NOMB": "Nombre",
    "PERS": "Personne",
    "POS": "Partie du discours",
    "TEMPS": "Temps",
    "lemma": "Lemme",
}


def discover_tasks():
    tasks = {}
    pattern = re.compile(r"^mf-([A-Za-z_]+)-")
    for path in sorted(MODELS_DIR.glob("*.tar")):
        m = pattern.match(path.name)
        if not m:
            continue
        task = m.group(1)
        tasks[task] = str(path)
    return tasks


TASKS = discover_tasks()

TASK_CHOICES = [(TASK_LABELS.get(k, k), k) for k in TASKS]


@lru_cache(maxsize=16)
def get_tagger(model_path: str) -> Tagger:
    tagger = Tagger(device="cpu", batch_size=8, lower=False)
    tagger.add_model(model_path)
    return tagger


def tag_text(selected_tasks, text):
    if not selected_tasks:
        return "Choisir au moins une tâche."
    if not text or not text.strip():
        return "Texte vide."

    n_words = len(text.split())
    if n_words > MAX_WORDS:
        return f"Texte trop long : {n_words} mots (max {MAX_WORDS})."

    sents = list(simple_tokenizer(text))
    if not sents:
        return "Aucun token détecté."
    lengths = [len(s) for s in sents]

    all_tasks = []
    all_runs = []
    for task_name in selected_tasks:
        path = TASKS[task_name]
        tagger = get_tagger(path)
        tagged, tasks = tagger.tag(sents=sents, lengths=lengths)
        all_tasks.extend(tasks)
        all_runs.append(tagged)

    header = "token\t" + "\t".join(all_tasks)
    lines = [header]
    for sent_idx, sent in enumerate(all_runs[0]):
        for tok_idx, (token, _) in enumerate(sent):
            row = [token]
            for tagged in all_runs:
                _, tags = tagged[sent_idx][tok_idx]
                row.extend(tags)
            lines.append("\t".join(row))
    return "\n".join(lines)


default_tasks = list(TASKS.keys())

CITATION = (
    "L'annotation du corpus et son entraînement ont été effectuées dans le cadre du "
    "Centre de ressources computationnelles pour les langues à variation graphique, créé "
    "au sein du Centre Jean-Mabillon (EA 3624) de l'Ecole des chartes. Le financement du "
    "Centre de ressources a été assuré par Biblissima+ (Equipex+ au titre du Programme "
    "d'investissements d'avenir intégré à France 2030, portant la référence ANR-21-ESRE-0005)."
)

with gr.Blocks(title="Annotateur moyen français") as demo:
    gr.Markdown("## Annotateur moyen français")
    tasks_cb = gr.CheckboxGroup(
        choices=TASK_CHOICES,
        value=default_tasks,
        label="Tâches",
    )
    text = gr.Textbox(lines=15, label=f"Texte (≤ {MAX_WORDS} mots)")
    run = gr.Button("Annoter")
    out = gr.Textbox(lines=20, label="TSV")

    run.click(tag_text, inputs=[tasks_cb, text], outputs=out)

    gr.Markdown(f"---\n_{CITATION}_")


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