import spaces  # MUST be the very first import for ZeroGPU Spaces

import os
import gradio as gr
import torch
import regex
import unicodedata
import re
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ─── Model Configuration ──────────────────────────────────────────────────────
BASE_MODEL = "unsloth/gemma-3-270m-it"
LORA_PATH  = "Razon2006/gemma3-tamil-Grapheme-NewChange-r32-lr"
hf_token = os.environ.get("HF_TOKEN")

print("Loading tokenizer globally...")
tokenizer = AutoTokenizer.from_pretrained(LORA_PATH, token=hf_token)

# Global model variable for lazy loading
model = None

# ─── Grapheme Tokenizer & Constants ───────────────────────────────────────────
INSTRUCTION = (
    "கீழ்க்காணும் தமிழ் வெண்பாவை, அதன் சொற்களை சரியான எல்லைகளுடன் "
    "பிரித்து எழுதுக (பதப்பிரிவு). சொற்களின் வரிசை மாறாமல், "
    "ஒட்டிய சொற்களை பிரிக்கவும், தேவையெனில் பிரிந்த எழுத்துக்களை "
    "இணைத்து சரியான சொல்லை உருவாக்கவும். பதப்பிரிவை மட்டும் தரவும்."
)

PROBLEMATIC_AKSHARAS = {
    'யு', 'றா', 'ஞா', 'யே', 'யை', 'றே', 'கெ', 'றோ', 'யெ', 'யொ',
    'றெ', 'ழா', 'நூ', 'லெ', 'வொ', 'ரெ', 'றொ', 'கீ', 'டே', 'லொ',
    'னொ', 'டொ', 'ணே', 'ளெ', 'ணோ', 'ரொ', 'றூ', 'ணீ', 'றீ', 'ழீ',
    'னூ', 'ணெ', 'வூ', 'யீ', 'டீ', 'ளொ', 'லூ', 'நொ', 'ணொ', 'டூ',
    'நை', 'கௌ', 'ளூ', 'ழே', 'ழூ', 'ஞை', 'ளீ', 'ஜை', 'ணூ', 'ழொ',
    'ழெ', 'ஞு', 'ஞெ', 'ஞூ', 'ஸா', 'ஜூ', 'எா', 'எீ', 'ழோ', 'பௌ',
}

def clean_grapheme_split(text, tokenizer):
    words = text.split(' ')
    result_words = []
    for word in words:
        if not word.strip():
            continue
        clusters = regex.findall(r'\X', word)
        clusters = [
            unicodedata.normalize('NFC', c) for c in clusters
            if c.strip() and any('\u0B80' <= ch <= '\u0BFF' for ch in c)
        ]
        if not clusters:
            result_words.append(word)
            continue
        merged = []
        i = 0
        while i < len(clusters):
            c = clusters[i]
            tokens = tokenizer.tokenize(c)
            if (c in PROBLEMATIC_AKSHARAS or len(tokens) > 1) and i + 1 < len(clusters):
                combined = c + clusters[i+1]
                while len(tokenizer.tokenize(combined)) > 1 and i + 2 < len(clusters):
                    i += 1
                    combined = combined + clusters[i+1]
                merged.append(combined)
                i += 2
            else:
                merged.append(c)
                i += 1
        result_words.append('|'.join(merged))
    return ' '.join(result_words)

def post_process_grapheme(text):
    return text.replace('|', '')

# ─── Inference Pipeline ───────────────────────────────────────────────────────
@spaces.GPU(duration=60)
def generate_pathaseth(venba, max_new_tokens=256):
    global model
    
    # LAZY LOADING: Load the model here where CUDA is fully unlocked by ZeroGPU
    if model is None:
        print("Loading base model directly into GPU...")
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.bfloat16,
            device_map="cuda", # <--- FIX: Forces native loading to GPU to avoid precision loss
            token=hf_token
        )
        print("Loading LoRA adapter...")
        model = PeftModel.from_pretrained(base_model, LORA_PATH, token=hf_token)
        model.eval()
        print("Model ready for inference!")

    venba = venba.strip()
    if not venba:
        return "Please enter a Venba verse."

    venba_clean = re.sub(r'\s+', ' ', venba).replace('்்', '்').strip()
    venba_no_space = venba_clean.replace(' ', '')
    venba_g = clean_grapheme_split(venba_no_space, tokenizer)

    messages = [{
        "role": "user",
        "content": f"{INSTRUCTION}\n\nவெண்பா: {venba_g}\nபதப்பிரிவு:"
    }]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    bos = tokenizer.bos_token
    if bos and prompt.startswith(bos):
        prompt = prompt[len(bos):]

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            use_cache=True,
            do_sample=False, # <--- FIX: Ensures deterministic output (stops hallucination)
        )

    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
    raw = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    return post_process_grapheme(raw)

# ─── Example Venbas ───────────────────────────────────────────────────────────
EXAMPLES = [
    ["இறுதி பயப்பினு மெஞ்சா திறைவற் குறுதி பயப்பதாந் தூது"],
]

# ─── Gradio Interface ─────────────────────────────────────────────────────────
with gr.Blocks(
    title="Tamil Venba Pathasetham",
    theme=gr.themes.Soft()
) as demo:

    gr.HTML("""
        <div style="text-align: center; font-size: 2em; font-weight: bold; color: #4B0082;">Tamil Venba Pathasetham (பதசேதம்)</div>
        <div style="text-align: center; color: #666; margin-bottom: 20px;">
            Automated word segmentation of classical Tamil Venba poetry using Gemma 3 + Grapheme Tokenization
        </div>
    """)

    with gr.Row():
        with gr.Column():
            venba_input = gr.Textbox(
                label="Enter Tamil Venba (வெண்பா)",
                placeholder="வண்மைதரு மாகமநூன் வைத்த பொருள்வழுவா...",
                lines=4,
            )
            submit_btn = gr.Button("Generate Pathasetham →", variant="primary", size="lg")

        with gr.Column():
            pathaseth_output = gr.Textbox(
                label="Pathasetham (பதசேதம்) Output",
                lines=4,
                interactive=False,
            )

    gr.Examples(
        examples=EXAMPLES,
        inputs=venba_input,
        outputs=pathaseth_output,
        fn=generate_pathaseth,
        cache_examples=False,
        label="Click any example to test instantly"
    )

    submit_btn.click(
        fn=generate_pathaseth,
        inputs=venba_input,
        outputs=pathaseth_output,
    )

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