import os
import torch
import gradio as ui
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_ID = "coder-vansh/cypher_llm_model"
HF_TOKEN = os.getenv("HF_TOKEN")

print("⚙️ System: Initializing 4-bit Quantization engine to save Cloud RAM...")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True
)

print("📥 Loading Model weights natively into memory space...")
# Load tokenizer and model weights inside the local basic CPU instance
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    low_cpu_mem_usage=True,
    token=HF_TOKEN
)
print("✅ CYPHER Brain is now natively alive inside the Space environment!")

def predict(message, history):
    system_prompt = "You are CYPHER, a casual, witty Hinglish AI companion by Vansh & Aditya. Tagline: \"Not just an AI — YOUR AI.\""
    
    # Building exact chat block matching schema
    messages = [{"role": "system", "content": system_prompt}]
    if history:
        for user_msg, bot_msg in history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})
            
    messages.append({"role": "user", "content": message})

    try:
        inputs = tokenizer.apply_chat_template(
            messages, 
            tokenize=True, 
            add_generation_prompt=True, 
            return_tensors="pt"
        )
        
        with torch.no_grad():
            outputs = model.generate(
                input_ids=inputs,
                max_new_tokens=250,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
            
        response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
        return response

    except Exception as e:
        return f"⚠️ Execution Glitch: {str(e)}"

# Elegant interface container
with ui.Blocks() as demo:
    ui.Markdown("# 🤖 CYPHER AI Live Production Space")
    ui.Markdown("### *Not just an AI — YOUR AI.* | Developed by Vansh & Aditya")
    
    ui.ChatInterface(
        fn=predict,
        textbox=ui.Textbox(placeholder="Bhai se kuch bhi poocho...", container=False, scale=7),
        description="CYPHER is running 100% Natively inside your cloud container.",
    )

demo.launch()