import gradio as gr
import os
import torch
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
from peft import PeftModel

# === Step 1: Login to Hugging Face ===
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("❌ HF_TOKEN not set! Add it in Settings > Secrets.")

login(token=HF_TOKEN)

# === Step 2: Model IDs ===
BASE_MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
ADAPTER_MODEL_ID = "sakibzaman/llama3.2-bengaliPunc-adapter-3b"

# Create a folder for offloading (required for low VRAM)
OFFLOAD_FOLDER = "./offload"
os.makedirs(OFFLOAD_FOLDER, exist_ok=True)

# === Step 3: 4-bit Quantization Config ===
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

# === Step 4: Load Model Safely ===
try:
    print("🔍 Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
    tokenizer.pad_token = tokenizer.eos_token  # Important for generation

    print("🧠 Loading base model with 4-bit quantization...")
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",           # Automatically splits layers across GPU/CPU
        torch_dtype=torch.float16,
        token=HF_TOKEN,
        offload_folder=OFFLOAD_FOLDER,  # Required for CPU offloading
    )

    print("📎 Loading LoRA adapter...")
    model = PeftModel.from_pretrained(
        base_model,
        ADAPTER_MODEL_ID,
        token=HF_TOKEN,
        offload_folder=OFFLOAD_FOLDER,
    )

    print("🚀 Creating generation pipeline...")
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,

        max_new_tokens=64,
        temperature=0.1,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
    )

    print("✅ Model loaded successfully!")

except Exception as e:
    raise RuntimeError(f"❌ Failed to load model: {e}")


# === Step 5: Inference Function ===
def punctuate_text(text):
    if not text.strip():
        return "⚠️ Please enter some Bengali text."

    try:

        messages = [
            {
                "role": "system",
                "content": "You are an expert in Bengali. Restore punctuation. Only return corrected text, nothing else."
            },
            {
                "role": "user",
                "content": text.strip()
            }
        ]

        outputs = pipe(messages)
        response = outputs[0]["generated_text"]

        # Extract only the assistant's part
        for msg in reversed(response):
            if msg["role"] == "assistant":
                return msg["content"].strip()

        return "�� Could not extract response."












    except Exception as e:
        return f"❌ Error: {str(e)}"


# === Step 6: Gradio UI & API ===
with gr.Blocks(theme=gr.themes.Soft(), title="🇧🇩 Bengali Punctuation") as demo:
    gr.HTML("""
    <h1 style="text-align: center;">🇧🇩 Bengali Punctuation Assistant</h1>
    <p style="text-align: center;">
        Powered by Llama-3.2-3B + LoRA • API enabled at <code>/api/predict</code>

    </p>
    """)

    inp = gr.Textbox(label="📝 Input Bengali Text", placeholder="আমি ঢাকায় থাকি আমার বাড়ি সিলেটে", lines=4)
    out = gr.Textbox(label="✅ Output", lines=4, interactive=False)






    btn = gr.Button("🔤 Add Punctuation", variant="primary")

    gr.Examples(
        examples=[
            ["আপনি কেমন আছেন"],
            ["আমি ঢাকায় থাকি আমার বাড়ি সিলেটে"],
            ["তুমি কখন আসবে আমি অপেক্ষা করছি"]
        ],
        inputs=inp,
        outputs=out,
        fn=punctuate_text,
        label="Try Examples"
    )

    btn.click(fn=punctuate_text, inputs=inp, outputs=out, api_name="punctuate")






    gr.HTML("""
    <div style="text-align: center; margin-top: 20px; color: #555;">
        💡 API: Use <code>/api/punctuate</code> for programmatic access


    </div>
    """)

# === Step 7: Launch ===
if __name__ == "__main__":
    demo.queue(max_size=200).launch(debug=True, show_api=True)