import html
import json
import re
import tempfile
from functools import lru_cache
from pathlib import Path

import gradio as gr
import torch
from huggingface_hub import hf_hub_download

from inference import load_lemmatizer, load_registry, MODEL_REPO_ID, MODEL_ROOT


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
REGISTRY = load_registry("models_registry.json")

TARGET_COL_IDX = 2
BATCH_SIZE = 32




def _display_name(item):
    return f"{item['language']} - {item['treebank']}"


LANGUAGES = sorted({item["language"] for item in REGISTRY.values()})

DISPLAY_TO_ID = {
    _display_name(item): model_id
    for model_id, item in REGISTRY.items()
}


def treebank_choices(language):
    choices = []

    for model_id, item in REGISTRY.items():
        if item["language"] == language:
            choices.append(_display_name(item))

    return sorted(choices)


def default_language():
    if "Old Church Slavonic" in LANGUAGES:
        return "Old Church Slavonic"

    return LANGUAGES[0] if LANGUAGES else None


def default_treebank(language):
    choices = treebank_choices(language)
    preferred = "Old Church Slavonic - PROIEL"

    if preferred in choices:
        return preferred

    return choices[0] if choices else None


CUSTOM_CSS = """
body {
    background: linear-gradient(135deg, #eaf3ff 0%, #ffffff 48%, #dbeafe 100%);
}
.gradio-container {
    max-width: 980px !important;
    margin: auto !important;
    font-family: Arial, Helvetica, sans-serif !important;
}
#main-card {
    background: #ffffff;
    border: 1px solid #bfdbfe;
    border-radius: 26px;
    padding: 30px;
    box-shadow: 0 20px 50px rgba(15, 23, 42, 0.16);
}
#title {
    text-align: center;
    color: #020617;
    font-size: 2.5rem;
    font-weight: 900;
    margin-bottom: 0.25rem;
}
#subtitle {
    text-align: center;
    color: #1e40af;
    font-size: 1.05rem;
    line-height: 1.55;
    margin-bottom: 1.6rem;
}
#badge-row {
    text-align: center;
    margin-bottom: 1.2rem;
}
#badge-row span {
    display: inline-block;
    background: #eff6ff;
    color: #1e3a8a;
    border: 1px solid #bfdbfe;
    border-radius: 999px;
    padding: 7px 13px;
    margin: 4px;
    font-size: 0.88rem;
    font-weight: 700;
}
textarea, input, select {
    border-radius: 16px !important;
    border: 1.5px solid #2563eb !important;
    background: #ffffff !important;
    color: #020617 !important;
    box-shadow: 0 6px 16px rgba(37, 99, 235, 0.08) !important;
}
textarea:focus, input:focus {
    border-color: #1d4ed8 !important;
    box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.18) !important;
}
label {
    color: #020617 !important;
    font-weight: 800 !important;
}
button {
    background: linear-gradient(90deg, #020617, #1d4ed8) !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 16px !important;
    padding: 13px 24px !important;
    font-weight: 900 !important;
    font-size: 1rem !important;
    box-shadow: 0 10px 22px rgba(37, 99, 235, 0.30) !important;
}
button:hover {
    background: linear-gradient(90deg, #000000, #2563eb) !important;
    transform: translateY(-1px);
}
#output-box textarea {
    background: #f8fbff !important;
    border: 1.5px solid #1d4ed8 !important;
    color: #020617 !important;
    font-family: Consolas, "Courier New", monospace !important;
}
#token-card {
    background: #f8fbff;
    border: 1.5px solid #1d4ed8;
    border-radius: 18px;
    padding: 18px;
    box-shadow: 0 8px 20px rgba(37, 99, 235, 0.10);
}
.lemma-table {
    width: 100%;
    border-collapse: separate;
    border-spacing: 0 8px;
    font-size: 1rem;
}
.lemma-table th {
    background: linear-gradient(90deg, #020617, #1d4ed8);
    color: white;
    padding: 12px 14px;
    text-align: left;
    font-weight: 900;
}
.lemma-table th:first-child {
    border-radius: 12px 0 0 12px;
}
.lemma-table th:last-child {
    border-radius: 0 12px 12px 0;
}
.lemma-table td {
    background: #ffffff;
    color: #020617;
    padding: 12px 14px;
    border-top: 1px solid #bfdbfe;
    border-bottom: 1px solid #bfdbfe;
    font-weight: 700;
}
.lemma-table td:first-child {
    border-left: 1px solid #bfdbfe;
    border-radius: 12px 0 0 12px;
}
.lemma-table td:last-child {
    border-right: 1px solid #bfdbfe;
    border-radius: 0 12px 12px 0;
    color: #1d4ed8;
}
#note {
    color: #1e3a8a;
    font-size: 0.92rem;
    text-align: center;
    margin-top: 1rem;
    font-weight: 600;
}
footer,
.api,
.settings,
a[href*="gradio.app"],
button[aria-label="Settings"] {
    display: none !important;
}
"""


def make_html_table(tokens, lemmas):
    if not tokens:
        return ""

    rows = []

    for token, lemma in zip(tokens, lemmas):
        rows.append(
            f"""
            <tr>
                <td>{html.escape(token)}</td>
                <td>{html.escape(lemma)}</td>
            </tr>
            """
        )

    return f"""
    <div id="token-card">
        <table class="lemma-table">
            <thead>
                <tr>
                    <th>Word</th>
                    <th>Lemma</th>
                </tr>
            </thead>
            <tbody>
                {''.join(rows)}
            </tbody>
        </table>
    </div>
    """


def update_treebanks(language):
    choices = treebank_choices(language)

    return gr.Dropdown(
        choices=choices,
        value=default_treebank(language),
    )


def selected_model_id(display_name):
    if not display_name or display_name not in DISPLAY_TO_ID:
        raise ValueError("Please select a valid language and treebank.")

    return DISPLAY_TO_ID[display_name]


@lru_cache(maxsize=128)
def load_vocab_chars_for_model(model_id):
    
    item = REGISTRY[model_id]

    vocab_path = hf_hub_download(
        repo_id=MODEL_REPO_ID,
        repo_type="model",
        filename=f"{MODEL_ROOT}/{item['folder']}/{item['vocab_file']}",
    )

    with open(vocab_path, encoding="utf8") as f:
        vocab_data = json.load(f)

    return set(vocab_data["char2idx"]) - {"<pad>", "<sos>", "<eos>", "<unk>"}


def unsupported_input_for_model(text, allowed_chars, max_bad_ratio=0.60, min_checked_chars=4):
    
    checked = []
    bad = []

    for ch in text:
        if ch.isspace():
            continue

        
        if ch.isdigit() or ch in {".", ",", ";", ":", "!", "?", "-", "'", '"', "(", ")", "[", "]", "/"}:
            continue

        checked.append(ch)

        if ch not in allowed_chars:
            bad.append(ch)

    if len(checked) < min_checked_chars:
        return False, []

    bad_ratio = len(bad) / len(checked)

    return bad_ratio >= max_bad_ratio, sorted(set(bad))


def lemmatize_sentence(sentence, display_name):
    sentence = str(sentence).strip()

    if not sentence:
        return "", ""

    try:
        model_id = selected_model_id(display_name)
    except ValueError as e:
        return "", str(e)

    allowed_chars = load_vocab_chars_for_model(model_id)

    is_bad, bad = unsupported_input_for_model(sentence, allowed_chars)
    if is_bad:
        return "",f"This input does not seem to match the selected language/treebank. Please select your desired language and treebank, then try again. Unsupported characters: {' '.join(bad[:20])}"

    lemmatizer = load_lemmatizer(model_id, DEVICE)

    tokens = sentence.split()
    lemmas = lemmatizer.lemmatize_sentence(tokens)

    lemmatized_sentence = " ".join(lemmas)
    token_html = make_html_table(tokens, lemmas)

    return lemmatized_sentence, token_html


def parse_conllu_sentences_from_text(text):
    text = text.strip()
    sents = []

    for block in re.split(r"\n\n+", text):
        sent = []

        for line in block.splitlines():
            if not line or line.startswith("#"):
                continue

            cols = line.split("\t")

            if len(cols) != 10:
                continue

            tok_id = cols[0]

            if "-" in tok_id or "." in tok_id:
                continue

            form = cols[1]
            sent.append(form)

        if sent:
            sents.append(sent)

    return sents


def make_source_for_token(tokens, index, k_context, sep_char):
    form = tokens[index]

    left_context = tokens[max(0, index - k_context):index]
    right_context = tokens[index + 1:index + 1 + k_context]

    left = " ".join(left_context).strip()
    right = " ".join(right_context).strip()

    src_left = left + " " if left else ""
    src_right = " " + right if right else ""

    return f"{src_left}{sep_char}{form}{sep_char}{src_right}"


def make_all_sources_from_conllu(text, lemmatizer):
    sents = parse_conllu_sentences_from_text(text)

    sources = []

    for tokens in sents:
        for i in range(len(tokens)):
            src_string = make_source_for_token(
                tokens=tokens,
                index=i,
                k_context=lemmatizer.k_context,
                sep_char=lemmatizer.sep_char,
            )
            sources.append(src_string)

    return sources


def predict_sources_batched(sources, lemmatizer, batch_size=BATCH_SIZE):
    preds_all = []

    if not sources:
        return preds_all

    pad_id = lemmatizer.vocab.char2idx["<pad>"]
    sos_id = lemmatizer.vocab.char2idx["<sos>"]
    eos_id = lemmatizer.vocab.char2idx["<eos>"]

    for start in range(0, len(sources), batch_size):
        batch_sources = sources[start:start + batch_size]

        src_ids_list = []
        src_lens = []

        for src_string in batch_sources:
            src_ids = (
                [sos_id]
                + lemmatizer.vocab.encode(src_string)
                + [eos_id]
            )

            src_ids_list.append(src_ids)
            src_lens.append(len(src_ids))

        max_len = max(src_lens)

        padded = [
            ids + [pad_id] * (max_len - len(ids))
            for ids in src_ids_list
        ]

        src = torch.tensor(
            padded,
            dtype=torch.long,
            device=lemmatizer.device,
        )

        src_lens_tensor = torch.tensor(
            src_lens,
            dtype=torch.long,
            device=lemmatizer.device,
        )

        batch_preds = lemmatizer.model.generate(
            src,
            src_lens_tensor,
            lemmatizer.vocab,
            max_len=lemmatizer.max_gen_len,
        )

        preds_all.extend(batch_preds)

    return preds_all


def predict_conllu_lemmas(text, lemmatizer):
    sources = make_all_sources_from_conllu(text, lemmatizer)

    return predict_sources_batched(
        sources=sources,
        lemmatizer=lemmatizer,
        batch_size=BATCH_SIZE,
    )


def write_back_conllu(input_text, preds_all):
    text = input_text.rstrip("\n")
    blocks = re.split(r"\n\n+", text)

    out_blocks = []
    p = 0

    for block in blocks:
        lines = block.split("\n")
        new_lines = []

        for line in lines:
            if not line or line.startswith("#"):
                new_lines.append(line)
                continue

            cols = line.split("\t")

            if len(cols) != 10:
                new_lines.append(line)
                continue

            tok_id = cols[0]

            if "-" in tok_id or "." in tok_id:
                new_lines.append(line)
                continue

            pred = preds_all[p] if p < len(preds_all) else "_"
            cols[TARGET_COL_IDX] = pred if pred else "_"

            new_lines.append("\t".join(cols))
            p += 1

        out_blocks.append("\n".join(new_lines))

    output_text = "\n\n".join(out_blocks).rstrip() + "\n\n"

    return output_text, p


def lemmatize_conllu_file(file_obj, display_name):
    if file_obj is None:
        return gr.update(value=None, visible=False), "Please upload a CoNLL-U file."

    try:
        model_id = selected_model_id(display_name)
    except ValueError as e:
        return gr.update(value=None, visible=False), str(e)

    input_path = Path(file_obj.name)

    with open(input_path, encoding="utf8") as f:
        text = f.read()

    lemmatizer = load_lemmatizer(model_id, DEVICE)

    preds = predict_conllu_lemmas(text, lemmatizer)
    output_text, total = write_back_conllu(text, preds)

    safe_model_name = display_name.replace(" ", "_").replace("-", "_")
    safe_model_name = re.sub(r"[^A-Za-z0-9_]+", "", safe_model_name)

    out_path = (
        Path(tempfile.gettempdir())
        / f"{input_path.stem}.{safe_model_name}.lemmatized.conllu"
    )

    with open(out_path, "w", encoding="utf8", newline="\n") as f:
        f.write(output_text)

    message = (
        f"Done. Wrote {total:,} lemma predictions.\n"
        f"Input used: FORM column only.\n"
        f"Updated column: LEMMA only.\n"
        f"All other CoNLL-U columns and comments were preserved."
    )

    return gr.update(value=str(out_path), visible=True), message


def reset_download_button(file_obj):
    return gr.update(value=None, visible=False), ""


DEFAULT_LANGUAGE = default_language()
DEFAULT_TREEBANK = default_treebank(DEFAULT_LANGUAGE)


with gr.Blocks(
    title="oldslaviclemma",
    css=CUSTOM_CSS,
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="sky",
        neutral_hue="slate",
    ),
) as demo:
    with gr.Column(elem_id="main-card"):
        gr.Markdown("# oldslaviclemma", elem_id="title")

        gr.Markdown(
            "Select a language and treebank. Paste one sentence or upload a tokenized CoNLL-U file. "
            "The system returns lemma predictions while preserving the original tokenization.",
            elem_id="subtitle",
        )

        gr.HTML(
            """
            <div id="badge-row">
                <span>oldslaviclemma</span>
                <span>60+ languages</span>
                <span>110+ treebanks</span>
                <span>UD v2.12</span>
                <span>Lemmatization</span>
            </div>
            """
        )

        with gr.Row():
            language_input = gr.Dropdown(
                label="Language",
                choices=LANGUAGES,
                value=DEFAULT_LANGUAGE,
            )

            treebank_input = gr.Dropdown(
                label="Treebank",
                choices=treebank_choices(DEFAULT_LANGUAGE),
                value=DEFAULT_TREEBANK,
            )

        with gr.Tab("Sentence input"):
            sentence_input = gr.Textbox(
                label="Input sentence",
                lines=5,
                value="",
                placeholder="Paste a sentence with words separated by spaces...",
            )

            run_button = gr.Button("Lemmatize sentence")

            sentence_output = gr.Textbox(
                label="Lemmatized sentence",
                lines=5,
                elem_id="output-box",
            )

            token_output = gr.HTML(
                label="Word-level output",
                value="",
            )

            gr.Markdown(
                "Please paste one sentence with whitespace tokenization.",
                elem_id="note",
            )

        with gr.Tab("CoNLL-U file input"):
            gr.Markdown(
                "Upload an already-tokenized CoNLL-U file. "
                "The app reads the FORM column, predicts the LEMMA column, "
                "and preserves comments, token IDs, UPOS, XPOS, FEATS, HEAD, DEPREL, DEPS, and MISC."
            )

            conllu_input = gr.File(
                label="Upload CoNLL-U file",
                file_types=[".conllu", ".txt"],
            )

            conllu_button = gr.Button("Lemmatize CoNLL-U file")

            conllu_output = gr.DownloadButton(
                label="Download lemmatized CoNLL-U file",
                value=None,
                visible=False,
            )

            conllu_message = gr.Textbox(
                label="Status",
                lines=4,
            )

        language_input.change(
            fn=update_treebanks,
            inputs=language_input,
            outputs=treebank_input,
        )

        run_button.click(
            fn=lemmatize_sentence,
            inputs=[sentence_input, treebank_input],
            outputs=[sentence_output, token_output],
        )

        conllu_input.change(
            fn=reset_download_button,
            inputs=conllu_input,
            outputs=[conllu_output, conllu_message],
        )

        conllu_button.click(
            fn=lemmatize_conllu_file,
            inputs=[conllu_input, treebank_input],
            outputs=[conllu_output, conllu_message],
        )


demo.queue()
demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)