import torch
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
import os

# UI
st.title("AI Tutor (Fine-tuned LLM)")
st.write("This AI tutor is fine-tuned on Python-related questions.")

# Load base model and tokenizer
model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 🔥 Fix: Add padding token if missing
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})

# Check if fine-tuned model exists
model_path = "./models"
if os.path.exists(model_path):
    st.write("✅ Loading fine-tuned model...")
    model = AutoModelForCausalLM.from_pretrained(model_path)
else:
    st.write("⚡ Fine-tuning the model (this will take time)...")

    # Load model on CPU
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32,  # Use float32 for CPU compatibility
        device_map={"": "cpu"}  # Force CPU usage
    )

    # Resize model embeddings
    model.resize_token_embeddings(len(tokenizer))

    # Apply LoRA
    lora_config = LoraConfig(
        r=8,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora_config)

    # Load dataset (Choose any one)
    dataset = load_dataset("lvwerra/codeparrot-clean", split="train")  # ✅ Free dataset

    # 🔥 Fix: Set `labels` properly
    def tokenize_function(examples):
        inputs = tokenizer(examples["content"], padding="max_length", truncation=True, max_length=512)
        inputs["labels"] = inputs["input_ids"].copy()  # ✅ Ensure labels exist
        return inputs

    tokenized_dataset = dataset.map(tokenize_function, batched=True)

    # Data collator
    data_collator = DataCollatorForSeq2Seq(tokenizer, return_tensors="pt")

    # Training arguments
    training_args = TrainingArguments(
        per_device_train_batch_size=1,
        num_train_epochs=1,  # Reduce epochs for quick training
        learning_rate=3e-4,
        output_dir=model_path,
        save_strategy="epoch",
        logging_dir="./logs",
        logging_steps=10,
        save_total_limit=2,
        evaluation_strategy="no",  # ✅ No eval dataset needed
        load_best_model_at_end=False  # ✅ Prevents conflicts
    )

    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=data_collator,
    )

    # Train
    trainer.train()

    # Save model
    model.save_pretrained(model_path)
    tokenizer.save_pretrained("./tokenizer")

    st.write("🎉 Fine-tuning complete! Model saved.")

# Chat Interface
user_input = st.text_input("Ask a coding question:")
if user_input:
    inputs = tokenizer(user_input, return_tensors="pt").to("cpu")
    outputs = model.generate(**inputs, max_length=150)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    st.write("🤖 AI Tutor:", response)