import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from peft import PeftModel
from threading import Thread

# 1. Map both coordinates
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
ADAPTER_MODEL = "Cydercoder/qwen2.5-coder-3b"

print("Loading official base tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Configuring aggressive 4-bit CPU/GPU quantization parameters...")
# This config compresses the weights from 32-bit down to 4-bit integers to fit in RAM
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float32,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    llm_int8_enable_fp32_cpu_offload=True # Crucial fallback for free CPU spaces
)

print("Loading compressed base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=quantization_config,
    device_map="auto"
)

print("Merging your custom fine-tuned engineering weights...")
model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL)

def chat_function(message, history):
    messages = [
        {"role": "system", "content": "You are an expert full-stack developer assistant fine-tuned for frontend, backend, animations, and debugging."}
    ]
    
    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})
    
    # Fix: return_dict=False forces the tokenizer to deliver pure tensor arrays to the generation layer
    inputs = tokenizer.apply_chat_template(
        messages, 
        add_generation_prompt=True, 
        return_tensors="pt",
        return_dict=False
    )
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.6,
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text


demo = gr.ChatInterface(
    fn=chat_function,
    title="🤖 Cydercoder Qwen 3B AI Chatbot",
    description="Your custom fine-tuned assistant running compressed for speed in the cloud.",
    examples=["Write a login form using React and Tailwind.", "Fix this code error: Cannot read properties of undefined"]
)

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