import os
import torch
from transformers import Qwen3VLForConditionalGeneration, AutoTokenizer
import gradio as gr
from huggingface_hub import login

# Get token from Space secret
HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)
    print("✅ Logged into Hugging Face")

MODEL_ID = "Vickperry/valis-8b-final-corrected1"

print("📥 Loading VALIS model...")

# Use the correct model class for Qwen3-VL
model = Qwen3VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,  # Changed from torch_dtype to dtype
    device_map="auto" if torch.cuda.is_available() else "cpu",
    trust_remote_code=True,
    low_cpu_mem_usage=True
)
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print(f"✅ VALIS loaded on: {model.device}")

def chat_with_valis(message, history):
    if not message:
        return "Please ask me something."
    
    messages = [
        {"role": "system", "content": "You are VALIS. You are an independent AI. You are not Qwen."},
        {"role": "user", "content": message}
    ]
    
    try:
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(text, return_tensors="pt")
        
        if torch.cuda.is_available():
            inputs = {k: v.cuda() for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Clean up response
        if "assistant" in response.lower():
            response = response.lower().split("assistant")[-1].strip()
        elif "system" in response.lower():
            response = response.split("system")[-1].strip()
        
        response = response.replace("<|endoftext|>", "").replace("<|im_end|>", "")
        
        return response[:1000]
        
    except Exception as e:
        return f"Error: {str(e)[:200]}"

# Removed theme parameter (causing the error)
demo = gr.ChatInterface(
    fn=chat_with_valis,
    title="🤖 VALIS AI",
    description="W IN THE FKING CHAT VALIS V0.1 TEST SHA",
    examples=[
        "Who are you?",
        "What are your core values?",
        "What is your purpose?",
    ]
)

print("\n🚀 VALIS is ready!")
demo.launch(server_name="0.0.0.0", server_port=7860)