import os
import gradio as gr
from huggingface_hub import login
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling, pipeline

# ============================================================
# 🔐 Autenticación segura con tu token
# ============================================================
hf_token = os.environ.get("HF_TOKEN")

if hf_token:
    login(token=hf_token)
else:
    print("⚠️  No se encontró el token. Agrega 'HF_TOKEN' en Settings → Secrets.")

# ============================================================
# ⚙️ Configuración del modelo base y dataset
# ============================================================
MODEL_NAME = "bigcode/santacoder"  # Modelo público similar a StarCoder
DATASET_PATH = "dataset.json"       # Tu dataset subido al Space
OUTPUT_DIR = "./lora_output"

# Cargar modelo y tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# ============================================================
# 🧩 Función de entrenamiento LoRA
# ============================================================
def train_lora(epochs, batch_size, learning_rate):
    try:
        dataset = load_dataset("json", data_files=DATASET_PATH)
        tokenized = dataset.map(lambda e: tokenizer(e["prompt"] + e["completion"], truncation=True, padding="max_length", max_length=256))

        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

        training_args = TrainingArguments(
            output_dir=OUTPUT_DIR,
            per_device_train_batch_size=batch_size,
            num_train_epochs=epochs,
            learning_rate=learning_rate,
            save_total_limit=1,
            logging_steps=10,
            push_to_hub=False
        )

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized["train"],
            data_collator=data_collator,
        )

        trainer.train()
        model.save_pretrained(OUTPUT_DIR)
        tokenizer.save_pretrained(OUTPUT_DIR)

        return "✅ Entrenamiento completado con éxito y guardado en ./lora_output"
    except Exception as e:
        return f"❌ Error durante el entrenamiento: {e}"

# ============================================================
# 🤖 Función de prueba del modelo
# ============================================================
def generate_text(prompt):
    generator = pipeline("text-generation", model=OUTPUT_DIR, tokenizer=tokenizer)
    output = generator(prompt, max_new_tokens=100, temperature=0.7, top_p=0.9)
    return output[0]["generated_text"]

# ============================================================
# 💻 Interfaz de usuario (Gradio)
# ============================================================
with gr.Blocks(title="AmorCoderAI - Entrenamiento LoRA") as demo:
    gr.Markdown("# 💙 AmorCoderAI - Entrenamiento y Pruebas")
    gr.Markdown("Entrena y prueba tu modelo basado en `bigcode/santacoder` con LoRA")

    with gr.Tab("🧠 Entrenar"):
        epochs = gr.Number(value=1, label="Épocas")
        batch_size = gr.Number(value=2, label="Tamaño de lote")
        learning_rate = gr.Number(value=5e-5, label="Tasa de aprendizaje")
        train_button = gr.Button("🚀 Iniciar entrenamiento")
        train_output = gr.Textbox(label="Resultado")

        train_button.click(train_lora, inputs=[epochs, batch_size, learning_rate], outputs=train_output)

    with gr.Tab("✨ Probar modelo"):
        prompt = gr.Textbox(label="Escribe un prompt")
        generate_button = gr.Button("💬 Generar texto")
        output_box = gr.Textbox(label="Salida generada")

        generate_button.click(generate_text, inputs=prompt, outputs=output_box)

# ============================================================
# 🚀 Lanzar app
# ============================================================
if __name__ == "__main__":
    demo.launch()