import gradio as gr
import spaces
import torch
import os
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, set_peft_model_state_dict

# Configurations
MODEL_ID = "huihui-ai/Huihui-Qwen3.5-9B-abliterated"
DATASET_NAME = "LvcidPsyche/webreaper-deep-crawl" # We will want to filter this for 'reasoning'
CHECKPOINT_DIR = "./qwen_checkpoints"
BATCH_SIZE = 2
GRADIENT_ACCUMULATION = 8
LEARNING_RATE = 2e-5

os.makedirs(CHECKPOINT_DIR, exist_ok=True)

# 1. Setup QLoRA Configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

lora_config = LoraConfig(
    r=16, 
    lora_alpha=32, 
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Initialize globally so we don't reload the massive base weights every 120s
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token

print("Loading base Huihui-Qwen3.5-9B model in 4-bit...")
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto"
)
model = get_peft_model(base_model, lora_config)
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, fused=True)

@spaces.GPU(duration=120)
def train_chunk_qwen(steps=20, checkpoint_path=None):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    start_step = 0
    if checkpoint_path and os.path.exists(checkpoint_path):
        checkpoint = torch.load(checkpoint_path, map_location='cpu')
        # Only load the LoRA adapter weights, NOT the 7B base parameters
        set_peft_model_state_dict(model, checkpoint['lora_weights'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_step = checkpoint.get('step', 0)
        
    model.train()
    
    # We stream the reasoning-filtered data
    ds = load_dataset(DATASET_NAME, split="train", streaming=True)
    ds_iter = iter(ds)
    
    logs = []
    logs.append(f"--- Resuming Huihui-Qwen3.5-9B LoRA training at Step {start_step} ---")
    
    for step in range(start_step, start_step + steps):
        try:
            row = next(ds_iter)
            # Ideally this is a deeply structured reasoning trace.
            text = row.get("content_text", "")
            if len(text) < 100: continue
        except StopIteration:
            break
            
        # Format for Qwen Chat Template
        # In practice, you'd map your curated data to User/Assistant turns
        prompt = f"<|im_start|>user\nAnalyze the underlying architecture of this system.\n<|im_end|>\n<|im_start|>assistant\n<think>\nAnalyzing constraints...\n</think>\n{text[:1000]}<|im_end|>"
        
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
        
        # Forward Pass (bfloat16)
        with torch.autocast(device_type=device, dtype=torch.bfloat16):
            outputs = model(**inputs, labels=inputs["input_ids"])
            loss = outputs.loss / GRADIENT_ACCUMULATION
            
        loss.backward()
        
        if (step + 1) % GRADIENT_ACCUMULATION == 0:
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            
        if step % 5 == 0:
            logs.append(f"Step {step} | Loss: {loss.item() * GRADIENT_ACCUMULATION:.4f}")
            
    # Save Checkpoint BEFORE ZeroGPU shuts down (Only save LoRA weights ~50MB)
    out_ckpt = os.path.join(CHECKPOINT_DIR, f"qwen_lora_step_{start_step + steps}.pt")
    
    from peft import get_peft_model_state_dict
    torch.save({
        'lora_weights': get_peft_model_state_dict(model),
        'optimizer': optimizer.state_dict(),
        'step': start_step + steps
    }, out_ckpt)
    
    if checkpoint_path and os.path.exists(checkpoint_path) and checkpoint_path != out_ckpt:
        os.remove(checkpoint_path)
        
    return "\n".join(logs), out_ckpt

current_ckpt = None
def run_ui(steps):
    global current_ckpt
    logs, new_ckpt = train_chunk_qwen(steps=int(steps), checkpoint_path=current_ckpt)
    current_ckpt = new_ckpt
    return logs, f"Active LoRA Checkpoint: {current_ckpt}"

with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Huihui-Qwen3.5-9B ZeroGPU Reasoning Trainer")
    gr.Markdown("Fine-tunes the uncensored/abliterated Qwen 9B using QLoRA (4-bit). Designed for chunked state execution.")
    steps_slider = gr.Slider(5, 50, 20, step=5, label="Steps per burst")
    train_btn = gr.Button("Train Qwen 9B Chunk")
    log_out = gr.Textbox(lines=10)
    ckpt_out = gr.Textbox()
    train_btn.click(run_ui, [steps_slider], [log_out, ckpt_out])

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