import os
import gradio as gr
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# ============================================================
# 🔐 Autenticación con Hugging Face
# ============================================================
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 y Carga ÚNICA del Modelo
# ============================================================
MODEL_NAME = "bigcode/santacoder"

print("🔄 Cargando modelo y tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# 🚀 Inicialización del pipeline UNA SOLA VEZ
# Esto hace que la generación sea mucho más rápida en la app web.
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# ============================================================
# 🤖 Generación de texto (Función ahora más rápida)
# ============================================================
def generate_text(prompt):
    try:
        # ⚠️ Usa el 'generator' que ya está cargado globalmente
        output = generator(prompt, max_new_tokens=200, temperature=0.7, top_p=0.9)
        return output[0]["generated_text"]
    except Exception as e:
        return f"❌ Error al generar texto: {e}"

# ============================================================
# 💻 Interfaz con Gradio
# ============================================================
with gr.Blocks(title="💙 AmorCoderAI") as demo:
    gr.Markdown("# 💙 AmorCoderAI")
    gr.Markdown("Genera texto con un modelo BigCode libre (`santacoder`)")

    prompt = gr.Textbox(label="🧠 Escribe un prompt (ej: 'def fibonacci(n):')", lines=4) # Más espacio para prompts
    generate_button = gr.Button("💬 Generar texto")
    output_box = gr.Textbox(label="✨ Resultado", lines=10) # Más espacio para el código generado

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

if __name__ == "__main__":
    # Usa `share=True` si necesitas compartirlo temporalmente
    demo.launch(server_name="0.0.0.0", server_port=7860)