import fasttext
from huggingface_hub import hf_hub_download
import regex
import gradio as gr
import os
import asyncio
import atexit
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# ==================== Constants ====================
MAX_INPUT_LENGTH = 10000  # OpenLID character limit
COMMONLINGUA_MAX_BYTES = 512  # CommonLingua byte limit

# ==================== OpenLID Setup ====================
print("Loading OpenLID-v3 model...")
openlid_path = hf_hub_download(
    repo_id="HPLT/OpenLID-v3",
    filename="openlid-v3.bin"
)
openlid_model = fasttext.load_model(openlid_path)
print("OpenLID-v3 loaded successfully!")

# Preprocessing patterns for OpenLID
NONWORD_REPLACE_STR = r"[^\p{Word}\p{Zs}]|\d"
NONWORD_REPLACE_PATTERN = regex.compile(NONWORD_REPLACE_STR)
SPACE_PATTERN = regex.compile(r"\s\s+")

def openlid_preprocess(text):
    """Preprocess text for OpenLID-v3."""
    text = text.strip().replace('\n', ' ').lower()
    text = regex.sub(SPACE_PATTERN, " ", text)
    text = regex.sub(NONWORD_REPLACE_PATTERN, "", text)
    return text

# ==================== CommonLingua Setup ====================
# Inline model architecture (from model.py) so no extra file is needed
class ByteNgramEmbed(nn.Module):
    def __init__(self, num_buckets=4096, embed_dim=64, n=3):
        super().__init__()
        self.n = n
        self.num_buckets = num_buckets
        self.embed = nn.Embedding(num_buckets, embed_dim)

    def forward(self, byte_ids):
        B, T = byte_ids.shape
        clamped = byte_ids.clamp(max=255)
        padded = F.pad(clamped, (0, self.n - 1), value=0)
        h = torch.zeros(B, T, dtype=torch.long, device=byte_ids.device)
        for i in range(self.n):
            h = h * 257 + padded[:, i:i + T]
        return self.embed(h % self.num_buckets)


class ByteConvBlock(nn.Module):
    def __init__(self, d_model, kernel_size=15, expand=2):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.pad = kernel_size - 1
        self.conv = nn.Conv1d(d_model, d_model, kernel_size, groups=d_model)
        self.norm2 = nn.LayerNorm(d_model)
        ffn = d_model * expand
        self.ffn_gate = nn.Linear(d_model, ffn, bias=False)
        self.ffn_up = nn.Linear(d_model, ffn, bias=False)
        self.ffn_down = nn.Linear(ffn, d_model, bias=False)

    def forward(self, x):
        residual = x
        x = self.norm1(x).transpose(1, 2)
        x = F.pad(x, (self.pad, 0))
        x = F.silu(self.conv(x)).transpose(1, 2)
        x = residual + x

        residual = x
        x = self.norm2(x)
        x = self.ffn_down(F.silu(self.ffn_gate(x)) * self.ffn_up(x))
        return residual + x


def _rope(q, k):
    head_dim = q.shape[-1]
    seq_len = q.shape[-2]
    freqs = 1.0 / (10000.0 ** (torch.arange(0, head_dim, 2, device=q.device).float() / head_dim))
    t = torch.arange(seq_len, device=q.device)
    a = torch.outer(t, freqs)
    cos = a.cos().to(q.dtype)
    sin = a.sin().to(q.dtype)

    def rot(x):
        x1, x2 = x[..., : head_dim // 2], x[..., head_dim // 2:]
        return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)

    return rot(q), rot(k)


class ByteAttnBlock(nn.Module):
    def __init__(self, d_model, n_heads=4, expand=2):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.norm1 = nn.LayerNorm(d_model)
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.norm2 = nn.LayerNorm(d_model)
        ffn = d_model * expand
        self.ffn_gate = nn.Linear(d_model, ffn, bias=False)
        self.ffn_up = nn.Linear(d_model, ffn, bias=False)
        self.ffn_down = nn.Linear(ffn, d_model, bias=False)

    def forward(self, x):
        B, T, D = x.shape
        residual = x
        h = self.norm1(x)
        qkv = self.qkv(h).reshape(B, T, 3, self.n_heads, self.head_dim)
        q, k, v = (t.transpose(1, 2) for t in qkv.unbind(dim=2))
        q, k = _rope(q, k)
        attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = attn.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, D)
        x = residual + self.out_proj(out)

        residual = x
        h = self.norm2(x)
        h = self.ffn_down(F.silu(self.ffn_gate(h)) * self.ffn_up(h))
        return residual + h


class ByteHybrid(nn.Module):
    def __init__(
        self,
        num_classes,
        d_model=256,
        n_conv=3,
        n_attn=1,
        n_heads=4,
        ffn_expand=2,
        max_len=512,
        conv_kernel=15,
        ngram_buckets=0,
        ngram_dim=64,
    ):
        super().__init__()
        self.max_len = max_len
        self.embed = nn.Embedding(257, d_model, padding_idx=256)

        self.ngram_embed = None
        if ngram_buckets > 0:
            self.ngram_embed = ByteNgramEmbed(ngram_buckets, ngram_dim, n=3)
            self.ngram_proj = nn.Linear(ngram_dim, d_model, bias=False)

        self.conv_layers = nn.ModuleList(
            [ByteConvBlock(d_model, conv_kernel, ffn_expand) for _ in range(n_conv)]
        )
        self.attn_layers = nn.ModuleList(
            [ByteAttnBlock(d_model, n_heads, ffn_expand) for _ in range(n_attn)]
        )
        self.final_norm = nn.LayerNorm(d_model)
        self.head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(d_model, num_classes),
        )

    def forward(self, byte_ids):
        pad_mask = byte_ids != 256
        x = self.embed(byte_ids)
        if self.ngram_embed is not None:
            x = x + self.ngram_proj(self.ngram_embed(byte_ids))
        for layer in self.conv_layers:
            x = layer(x)
        for layer in self.attn_layers:
            x = layer(x)
        x = self.final_norm(x)
        mask = pad_mask.unsqueeze(-1).to(x.dtype)
        x = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        return self.head(x)


CONFIGS = {
    "base_ngram": dict(
        d_model=256, n_conv=3, n_attn=1, n_heads=4, conv_kernel=15,
        ngram_buckets=4096, ngram_dim=64,
    ),
}

def commonlingua_encode(texts, max_len):
    out = np.full((len(texts), max_len), 256, dtype=np.int64)
    for i, t in enumerate(texts):
        if not isinstance(t, str):
            t = "" if t is None else str(t)
        raw = t.encode("utf-8", errors="replace")[:max_len]
        if raw:
            out[i, :len(raw)] = np.frombuffer(raw, dtype=np.uint8)
    return torch.from_numpy(out)


@torch.no_grad()
def commonlingua_predict(model, texts, idx2lang, max_len, device, top_k=3):
    """Returns a list of [(lang, prob), ...] (one list per text, top-k entries each)."""
    out = []
    batch = commonlingua_encode(texts, max_len).to(device)
    probs = torch.softmax(model(batch).float(), dim=-1)
    top_p, top_idx = probs.topk(top_k, dim=-1)
    for p_row, idx_row in zip(top_p.cpu().tolist(), top_idx.cpu().tolist()):
        out.append([(idx2lang[j], float(p)) for p, j in zip(p_row, idx_row)])
    return out


print("Loading CommonLingua model...")
commonlingua_path = hf_hub_download(
    repo_id="PleIAs/CommonLingua",
    filename="model.pt"
)
ckpt = torch.load(commonlingua_path, map_location="cpu", weights_only=False)
commonlingua_model = ByteHybrid(
    num_classes=ckpt["num_classes"],
    max_len=ckpt["max_len"],
    **CONFIGS[ckpt["config"]]
)
commonlingua_model.load_state_dict(ckpt["model_state_dict"])
commonlingua_model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
commonlingua_model = commonlingua_model.to(device)
commonlingua_idx2lang = {v: k for k, v in ckpt["lang2idx"].items()}
commonlingua_max_len = ckpt["max_len"]
print(f"CommonLingua loaded successfully! ({len(commonlingua_idx2lang)} languages, device={device})")

# ==================== Prediction Functions ====================
def predict_openlid(text, top_k=3, threshold=0.5):
    """Predict language using OpenLID-v3."""
    if not text or not text.strip():
        return "Please enter some text to analyze."
    
    processed_text = openlid_preprocess(text)
    if not processed_text.strip():
        return "Text contains no valid characters for language identification."
    
    predictions = openlid_model.predict(
        text=processed_text,
        k=min(top_k, 10),
        threshold=threshold,
        on_unicode_error="strict",
    )
    
    labels, scores = predictions
    results = []
    for label, score in zip(labels, scores):
        lang_code = label.replace("__label__", "")
        confidence = float(score) * 100
        results.append(f"**{lang_code}**: {confidence:.2f}%")
    
    return "\n\n".join(results) if results else "No predictions above threshold."


def predict_commonlingua(text, top_k=3):
    """Predict language using CommonLingua."""
    if not text or not text.strip():
        return "Please enter some text to analyze."
    
    results = commonlingua_predict(
        commonlingua_model, [text], commonlingua_idx2lang,
        commonlingua_max_len, device, top_k=min(top_k, 10)
    )
    
    formatted = []
    for lang, prob in results[0]:
        formatted.append(f"**{lang}**: {prob*100:.2f}%")
    return "\n\n".join(formatted)


def predict_both(text, top_k=3, threshold=0.5):
    """
    Run both models and return combined results.
    Returns tuple: (openlid_result, commonlingua_result, status_message)
    """
    # Check OpenLID length limit
    if len(text) > MAX_INPUT_LENGTH:
        return (
            f"**Error**: Input too long ({len(text):,} characters). Maximum allowed is {MAX_INPUT_LENGTH:,} characters.",
            f"**Error**: Input too long ({len(text):,} characters). Maximum allowed is {MAX_INPUT_LENGTH:,} characters.",
            "❌ Input exceeds maximum length."
        )
    
    # Check CommonLingua byte limit
    byte_length = len(text.encode('utf-8'))
    if byte_length > COMMONLINGUA_MAX_BYTES:
        status = f"⚠️ Warning: Input is {byte_length} bytes. CommonLingua works best with ≤{COMMONLINGUA_MAX_BYTES} bytes (first {COMMONLINGUA_MAX_BYTES} bytes will be used)."
    else:
        status = f"✅ Input length: {len(text):,} chars | {byte_length} bytes"
    
    openlid_result = predict_openlid(text, top_k, threshold)
    commonlingua_result = predict_commonlingua(text, top_k)
    
    return openlid_result, commonlingua_result, status


# ==================== Cleanup ====================
def cleanup():
    try:
        loop = asyncio.get_event_loop()
        if loop.is_running():
            loop.stop()
        if not loop.is_closed():
            loop.close()
    except Exception:
        pass

atexit.register(cleanup)

# ==================== Gradio Interface ====================
with gr.Blocks(title="OpenLID-v3 vs CommonLingua") as demo:
    gr.HTML("""
    <h1>🔍 Language Identification: OpenLID-v3 vs CommonLingua</h1>
    <p>Compare two state-of-the-art language identification models side-by-side.</p>
    <p>
      <em>OpenLID-v3</em>: <a href="https://huggingface.co/HPLT/OpenLID-v3" target="_blank">HPLT/OpenLID-v3</a> (fastText, 194+ languages)<br>
      <em>CommonLingua</em>: <a href="https://huggingface.co/PleIAs/CommonLingua" target="_blank">PleIAs/CommonLingua</a> (byte-level CNN+Attention, 334 languages, 2.35M params)
    </p>
    """)
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Input Text",
                placeholder="Enter text to identify its language...",
                lines=5,
                max_lines=10,
                max_length=MAX_INPUT_LENGTH
            )
            with gr.Row():
                top_k = gr.Slider(
                    minimum=1, maximum=10, value=3, step=1,
                    label="Top-K Predictions"
                )
                threshold = gr.Slider(
                    minimum=0.0, maximum=1.0, value=0.5, step=0.05,
                    label="OpenLID Confidence Threshold"
                )
            submit_btn = gr.Button("🔍 Identify Language", variant="primary")
            status = gr.Textbox(label="Status", interactive=False)
        
    with gr.Row():
        with gr.Column():
            openlid_output = gr.Markdown(label="OpenLID-v3 Predictions")
        with gr.Column():
            commonlingua_output = gr.Markdown(label="CommonLingua Predictions")
    
    # Examples
    gr.Examples(
        examples=[
            ["Asebter-a yura s wudem awurman d amagrad s tutlayt taqbaylit."],
            ["L'interès es d'utilizar un sistèma liure, personalizable e en occitan."],
            ["Maskinsjefen er oppteken av å løfta fram dei maritime utdanningane."],
            ["The quick brown fox jumps over the lazy dog."],
            ["Le renard brun rapide saute par-dessus le chien paresseux."],
            ["El rápido zorro marrón salta sobre el perro perezoso."],
            ["Быстрая коричневая лисица прыгает через ленивую собаку."],
            ["快速的棕色狐狸跳过了懒惰的狗。"],
            ["Wikipédia est une encyclopédie universelle, multilingue."],
            ["CommonLingua est un modèle d'identification de langue très léger."],
        ],
        inputs=input_text,
        label="Try these examples"
    )
    
    gr.Markdown(f"""
    ### Tips for best results:
    - **OpenLID-v3**: Text is automatically preprocessed (lowercased, normalized). Longer texts generally give more accurate predictions. Max {MAX_INPUT_LENGTH:,} characters.
    - **CommonLingua**: Operates directly on raw UTF-8 bytes (no tokenizer). Designed for paragraph-level corpus curation. Works best with ≤{COMMONLINGUA_MAX_BYTES} bytes. Not assessed on very short segments.
    - Use the **Top-K** slider to see more alternative predictions.
    - Use the **Threshold** slider to filter out uncertain OpenLID predictions (does not affect CommonLingua).
    """)
    
    # Event handlers
    submit_btn.click(
        fn=predict_both,
        inputs=[input_text, top_k, threshold],
        outputs=[openlid_output, commonlingua_output, status]
    )
    
    input_text.submit(
        fn=predict_both,
        inputs=[input_text, top_k, threshold],
        outputs=[openlid_output, commonlingua_output, status]
    )

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    
    try:
        demo.launch(
            server_name="0.0.0.0",
            server_port=port,
            ssr_mode=False,
            share=False,
            show_error=True
        )
    except KeyboardInterrupt:
        print("\nShutting down gracefully...")
    finally:
        cleanup()