import os
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import time
import spaces

# Prevent VRAM memory fragmentation
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

base_model_id = "haoranxu/X-ALMA-13B-Group7"
default_adapter_id = "sandipghimire/X-ALMA-13B-Nepali-Legal-Translator" 
adapter_path = default_adapter_id if not os.path.exists("./filtered_adapter") else "./filtered_adapter"

# Global model and tokenizer variables
model = None
tokenizer = None

# Check if CUDA is available in the container
cuda_available = torch.cuda.is_available()

@spaces.GPU
def translate(text):
    global model, tokenizer
    
    if not text.strip():
        return "Please enter some text to translate."
        
    if not cuda_available:
        return (
            "⚠️ Error: This Space is currently running on the CPU basic tier. "
            "To enable the translation model, please go to the Settings tab of this Space, "
            "scroll down to 'Space Hardware', and select 'ZeroGPU' or a GPU tier (requires credit card verification or a PRO account)."
        )
        
    # Lazy load model and tokenizer on ZeroGPU
    if model is None or tokenizer is None:
        print("Initializing tokenizer and loading model on ZeroGPU...")
        
        tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        print("Configuring 4-bit Quantization (QLoRA)...")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.bfloat16
        )
        
        print(f"Loading base model {base_model_id} on GPU...")
        base_model = AutoModelForCausalLM.from_pretrained(
            base_model_id,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        
        print(f"Loading LoRA adapter from {adapter_path}...")
        model = PeftModel.from_pretrained(base_model, adapter_path)
        model.eval()
        print("Model initialization complete!")

    prompt_content = f"Translate this from English to Nepali:\nEnglish: {text}\nNepali:"
    messages = [{"role": "user", "content": prompt_content}]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    
    start_time = time.time()
    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False
        )
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)]
    translation = tokenizer.decode(generated_ids[0], skip_special_tokens=True).strip()
    elapsed = time.time() - start_time
    
    print(f"Translated in {elapsed:.2f} seconds.")
    return translation

# Build Gradio UI
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")) as demo:
    gr.Markdown(
        """
        # 🇳🇵 X-ALMA-13B English-to-Nepali Translation Model (PRO)
        This translation model is fine-tuned on administrative and general English-to-Nepali parallel datasets.
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="English Source Text",
                placeholder="Enter English sentence or paragraph here...",
                lines=8
            )
            translate_btn = gr.Button("Translate to Nepali ➔", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(
                label="Nepali Translation",
                placeholder="Translation will appear here...",
                lines=8,
                interactive=False
            )
            
    translate_btn.click(fn=translate, inputs=input_text, outputs=output_text)

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