import time

import pandas as pd
import streamlit as st
from datasets import load_dataset
from transformers import AutoTokenizer

try:
    from nomeda import NomedaMSATokenizer
except ImportError:
    NomedaMSATokenizer = None

# -- Page Config & CSS --
st.set_page_config(
    page_title="MSA Tokenizer Arena 🏆",
    page_icon="🏆",
    layout="wide",
)

st.markdown(
    """
<style>
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
        max-width: 1200px;
    }
    h1 {
        color: #1f2937;
        font-weight: 800;
        margin-bottom: 0.5rem;
        text-align: center;
    }
    .subtitle {
        color: #6b7280;
        font-size: 1.1rem;
        margin-bottom: 2rem;
        text-align: center;
    }
    .stButton>button {
        background-color: #4f46e5;
        color: white;
        font-weight: 600;
        border-radius: 0.5rem;
        border: none;
        padding: 0.5rem 2rem;
        transition: all 0.2s ease-in-out;
        width: 100%;
    }
    .stButton>button:hover {
        background-color: #4338ca;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    }
    div[data-testid="stExpander"] {
        border-radius: 0.5rem;
        border: 1px solid #e5e7eb;
        box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    }
    .tokenizer-tag {
        display: inline-block;
        background-color: #f3f4f6;
        border: 1px solid #d1d5db;
        border-radius: 0.375rem;
        padding: 0.25rem 0.75rem;
        margin-bottom: 0.5rem;
        font-size: 0.875rem;
        color: #374151;
        font-family: monospace;
    }
</style>
""",
    unsafe_allow_html=True,
)

# -- Title --
st.markdown("<h1>🏆 MSA Tokenizer Arena</h1>", unsafe_allow_html=True)
st.markdown(
    "<div class='subtitle'>Evaluate and compare tokenizers on Modern Standard Arabic (MSA) datasets.</div>",
    unsafe_allow_html=True,
)

# -- State Management --
if "leaderboard" not in st.session_state:
    st.session_state.leaderboard = pd.DataFrame()
if "tokenizer_list" not in st.session_state:
    st.session_state.tokenizer_list = [
        "nomeda-lab/Nomeda-MSA-64Kv1",
        "aubmindlab/bert-base-arabertv02",
        "core42/jais-13b-chat",
    ]

# -- Configuration Section --
col1, col2 = st.columns(2)

with col1:
    with st.expander("📂 Dataset Configuration", expanded=True):
        data_source = st.radio(
            "Dataset Source",
            ["Hugging Face Hub", "Upload File (CSV/JSONL/Parquet)"],
            horizontal=True,
        )

        dataset_name = None
        uploaded_file = None

        if data_source == "Hugging Face Hub":
            dataset_name = st.text_input(
                "🤗 Hugging Face Dataset Path", value="abucoderr/arabic_news_dataset"
            )
            split = st.text_input("Dataset Split", value="train")
        else:
            uploaded_file = st.file_uploader(
                "Upload File", type=["csv", "jsonl", "parquet", "json"]
            )

        c1, c2 = st.columns(2)
        text_column = c1.text_input(
            "Text Column Name",
            value="text",
            help="The name of the column containing the text data.",
        )
        num_samples = c2.number_input(
            "Samples to Evaluate", min_value=10, max_value=50000, value=1000, step=100
        )

with col2:
    with st.expander("🤖 Tokenizers Configuration", expanded=True):
        st.markdown("Add Hugging Face tokenizer IDs to evaluate.")

        # Add new tokenizer using a form to ensure reliable interactions
        with st.form("add_tokenizer_form", clear_on_submit=True):
            add_col1, add_col2 = st.columns([3, 1])
            new_tok = add_col1.text_input(
                "Tokenizer ID",
                placeholder="e.g., openai/clip-vit-base-patch32",
                label_visibility="collapsed",
            )
            submitted = add_col2.form_submit_button("➕ Add")
            if submitted and new_tok:
                tok_clean = new_tok.strip()
                if tok_clean and tok_clean not in st.session_state.tokenizer_list:
                    st.session_state.tokenizer_list.append(tok_clean)
                    st.rerun()

        st.markdown("---")
        st.markdown("**Selected Tokenizers:**")

        # Display list with remove buttons
        if not st.session_state.tokenizer_list:
            st.info("No tokenizers added yet.")

        for i, tok in enumerate(st.session_state.tokenizer_list):
            tc1, tc2 = st.columns([5, 1])
            tc1.markdown(
                f"<div class='tokenizer-tag'>{tok}</div>", unsafe_allow_html=True
            )
            if tc2.button("❌", key=f"del_{i}", help="Remove tokenizer"):
                st.session_state.tokenizer_list.pop(i)
                st.rerun()

st.markdown("<br>", unsafe_allow_html=True)

# Run Button centered
run_col1, run_col2, run_col3 = st.columns([1, 2, 1])
with run_col2:
    run_eval = st.button("🚀 Run Evaluation", type="primary")

st.markdown("---")


# -- Main Logic --
@st.cache_resource
def load_hf_tokenizer(model_id):
    try:
        # Custom logic for Nomeda tokenizer loading
        if "nomeda-msa" in model_id.lower() or "nomeda-lab" in model_id.lower():
            if NomedaMSATokenizer is not None:
                try:
                    # In case the library supports it directly
                    return NomedaMSATokenizer.from_pretrained(model_id)
                except Exception:
                    # Fallback to downloading the tsv if needed
                    from huggingface_hub import hf_hub_download

                    tsv_path = hf_hub_download(
                        repo_id=model_id, filename=model_id.split("/")[-1] + ".tsv"
                    )
                    return NomedaMSATokenizer.from_tsv(tsv_path)
            else:
                return "The 'nomeda' package is required but not installed."

        return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    except Exception as e:
        return str(e)


@st.cache_data
def get_data(file_or_name, source_type, col_name, split, samples):
    try:
        if source_type == "Hugging Face Hub":
            dataset = load_dataset(file_or_name, split=f"{split}[:{samples}]")
            if col_name not in dataset.column_names:
                return (
                    None,
                    f"Column '{col_name}' not found. Available columns: {dataset.column_names}",
                )
            texts = [str(t) for t in dataset[col_name] if t is not None]
        else:
            if file_or_name.name.endswith(".csv"):
                df = pd.read_csv(file_or_name, nrows=samples)
            elif file_or_name.name.endswith(".parquet"):
                df = pd.read_parquet(file_or_name)
                df = df.head(samples)
            elif file_or_name.name.endswith(".jsonl") or file_or_name.name.endswith(
                ".json"
            ):
                df = pd.read_json(file_or_name, lines=True, nrows=samples)
            else:
                return None, "Unsupported file format."

            if col_name not in df.columns:
                return (
                    None,
                    f"Column '{col_name}' not found. Available columns: {list(df.columns)}",
                )
            texts = df[col_name].dropna().astype(str).tolist()

        return texts, None
    except Exception as e:
        return None, str(e)


if run_eval:
    if not st.session_state.tokenizer_list:
        st.error("Please add at least one tokenizer.")
    else:
        with st.status("Evaluating Tokenizers...", expanded=True) as status:
            st.write("Loading dataset...")
            if data_source == "Hugging Face Hub":
                if not dataset_name:
                    status.update(label="Failed", state="error")
                    st.error("Please provide a dataset name.")
                    st.stop()
                texts, data_err = get_data(
                    dataset_name, data_source, text_column, split, num_samples
                )
            else:
                if uploaded_file is None:
                    status.update(label="Failed", state="error")
                    st.error("Please upload a file.")
                    st.stop()
                texts, data_err = get_data(
                    uploaded_file, data_source, text_column, None, num_samples
                )

            if data_err:
                status.update(label="Failed", state="error")
                st.error(f"Error loading dataset: {data_err}")
                st.stop()

            total_chars = sum(len(t) for t in texts)
            st.write(
                f"✅ Loaded {len(texts)} samples. Total characters: **{total_chars:,}**"
            )

            results = []

            # Progress bar
            progress_bar = st.progress(0)
            total_toks = len(st.session_state.tokenizer_list)

            for idx, model_id in enumerate(st.session_state.tokenizer_list):
                st.write(f"🔄 Evaluating `{model_id}`...")
                tok = load_hf_tokenizer(model_id)

                if isinstance(tok, str):
                    st.warning(f"Failed to load `{model_id}`: {tok}")
                    progress_bar.progress((idx + 1) / total_toks)
                    continue

                vocab_size = (
                    len(tok)
                    if hasattr(tok, "__len__")
                    else getattr(tok, "vocab_size", "Unknown")
                )

                start_time = time.time()
                if type(tok).__name__ == "NomedaMSATokenizer":
                    # Custom handling for Nomeda tokenizer
                    total_tokens = 0
                    for t in texts:
                        encoded = tok.encode(t)
                        total_tokens += len(encoded.ids)
                else:
                    encoded = tok(texts, add_special_tokens=False)
                    total_tokens = sum(len(ids) for ids in encoded["input_ids"])
                end_time = time.time()
                time_taken = end_time - start_time
                tokens_per_sec = total_tokens / time_taken if time_taken > 0 else 0
                compression_ratio = (
                    total_chars / total_tokens if total_tokens > 0 else 0
                )

                results.append(
                    {
                        "Tokenizer": model_id,
                        "Compression Ratio (Chars/Token) ⬆️": round(
                            compression_ratio, 2
                        ),
                        "Total Tokens ⬇️": total_tokens,
                        "Speed (Tokens/sec) ⬆️": round(tokens_per_sec, 0),
                        "Vocab Size": vocab_size,
                        "Time Taken (s)": round(time_taken, 2),
                    }
                )
                progress_bar.progress((idx + 1) / total_toks)

            if results:
                df = pd.DataFrame(results)
                df = df.sort_values(
                    by="Compression Ratio (Chars/Token) ⬆️", ascending=False
                ).reset_index(drop=True)
                df.index = df.index + 1
                df.index.name = "Rank"
                st.session_state.leaderboard = df

            status.update(
                label="Evaluation Complete! 🎉", state="complete", expanded=False
            )

# -- Leaderboard Display --
if not st.session_state.leaderboard.empty:
    st.subheader("📊 Arena Leaderboard")
    st.markdown(
        "Higher **Compression Ratio** means the tokenizer is more efficient (requires fewer tokens to represent the same text)."
    )

    def highlight_max(s, props=""):
        return [props if v == s.max() else "" for v in s]

    def highlight_min(s, props=""):
        return [props if v == s.min() else "" for v in s]

    styled_df = (
        st.session_state.leaderboard.style.apply(
            highlight_max,
            props="background-color: #dcfce7; color: #166534; font-weight: bold;",
            subset=["Compression Ratio (Chars/Token) ⬆️", "Speed (Tokens/sec) ⬆️"],
        )
        .apply(
            highlight_min,
            props="background-color: #dcfce7; color: #166534; font-weight: bold;",
            subset=["Total Tokens ⬇️"],
        )
        .format(
            {
                "Total Tokens ⬇️": "{:,}",
                "Speed (Tokens/sec) ⬆️": "{:,.0f}",
                "Compression Ratio (Chars/Token) ⬆️": "{:.2f}",
            }
        )
    )

    st.dataframe(styled_df, use_container_width=True, height=400)

    col1, col2, col3 = st.columns([1, 1, 1])
    with col2:
        csv = st.session_state.leaderboard.to_csv().encode("utf-8")
        st.download_button(
            label="📥 Download Results (CSV)",
            data=csv,
            file_name="tokenizer_arena_results.csv",
            mime="text/csv",
            use_container_width=True,
        )