import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ---------------------------------------------------------------------------
# CONFIGURATION
# ---------------------------------------------------------------------------

MODEL_PATH = "./model" 

# Automatically detect hardware
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on: {device}")

# ---------------------------------------------------------------------------
# MODEL LOADING
# ---------------------------------------------------------------------------
try:
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_PATH,
        use_fast=False, # Essential for SentencePiece tokenizers (like Gemma/Llama)
        trust_remote_code=True
    )
    
    # Load model with appropriate precision
    # GPU = float16 (faster, less VRAM)
    # CPU = float32 (required for compatibility on basic CPU Spaces)
    torch_dtype = torch.float16 if device == "cuda" else torch.float32
    
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch_dtype,
        trust_remote_code=True
    ).to(device)
    
    model.eval()
    print("Model loaded successfully.")

except Exception as e:
    print(f"FATAL ERROR loading model: {e}")
    raise e

# ---------------------------------------------------------------------------
# INFERENCE FUNCTION
# ---------------------------------------------------------------------------
def classify_spam(text):
    # 1. Input Validation
    if not text or not text.strip():
        return "⚠️ Please enter a message."

    # 2. Prepare Prompt
    messages = [
        {"role": "user", "content": text},
    ]
    
    try:
        # Apply template (returns a string)
        prompt_str = tokenizer.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )
        
        # Convert string to PyTorch tensors and move to device
        model_inputs = tokenizer(prompt_str, return_tensors="pt").to(device)
        
        # 3. Generate
        with torch.no_grad():
            outputs = model.generate(
                **model_inputs,
                max_new_tokens=20,   # Keep strictly short for classification
                do_sample=False,     # STRICTLY REQUIRED: Deterministic output (no randomness)
                pad_token_id=tokenizer.eos_token_id
            )

        # 4. Decode and Clean
        # The model returns [Prompt + Answer]. We slice off the prompt.
        input_length = model_inputs["input_ids"].shape[1]
        generated_tokens = outputs[0][input_length:]
        
        result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
        # FIX: Manually remove the specific tag and whitespace
        cleaned_result = result.replace("<end_of_turn>", "").strip()
        
        return cleaned_result

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

# ---------------------------------------------------------------------------
# GRADIO INTERFACE
# ---------------------------------------------------------------------------
with gr.Blocks(title="Selling Intent Classifier") as demo:
    gr.Markdown("# 📧 Selling Intent Classifier")
    gr.Markdown("Enter a youtube comment section message selling or asking to subscribe channel")
    
    with gr.Row():
        inp = gr.Textbox(label="Message Content", placeholder="Paste text here...", lines=4)
        out = gr.Textbox(label="Classification Result")
    
    btn = gr.Button("Classify", variant="primary")
    
    # Connecting the button to the function
    btn.click(fn=classify_spam, inputs=inp, outputs=out)

    # --- ADDED EXAMPLES SECTION HERE ---
    gr.Examples(
        examples=[
            ["ye teacher bakar padhata hai mujha sub karlo pwalternative on sc"],
            ["""Hello Dosto / Bhaiyo / Bahno agar aapko Notes banaane me problem aa rahi hai ya aapko Notes Chaihiye school ke liye to ye ytube channel dekhiye .......  ya aapko Notes , Formula sheet , and exam related chize milengi ... jisse aapki Help ho sake  ............  

            Channel name : -  Edu Notes 

            (See logo of ytube channel )"""],
            ["ye sub toh bahut hard hai maths is tough"],
            ["Click here for a free prize and subscribe to my channel!"]
        ],
        inputs=inp,
        label="Try these examples:"
    )

# ---------------------------------------------------------------------------
# LAUNCH
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.queue().launch()