import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# Configurar el dispositivo (CPU)
device = torch.device("cpu")

# Cargar el modelo y tokenizer
print("Cargando modelo code-autocomplete-gpt2-base...")
model_name = "shibing624/code-autocomplete-gpt2-base"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# Mover modelo a CPU y ponerlo en modo evaluación
model.to(device)
model.eval()

# Configurar pad_token si no existe
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def autocomplete_text(input_text, max_tokens=20):
    """
    Autocompleta el texto/código de entrada usando code-autocomplete-gpt2-base
    
    Args:
        input_text (str): Texto/código inicial a completar
        max_tokens (int): Número máximo de tokens a generar
    
    Returns:
        str: Solo la parte nueva generada (sin el input original)
    """
    if not input_text.strip():
        return "Por favor, ingresa algún texto para completar."
    
    try:
        # Tokenizar el texto de entrada
        inputs = tokenizer.encode(input_text, return_tensors="pt", padding=True)
        inputs = inputs.to(device)
        
        # Generar texto
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=max_tokens,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
                attention_mask=torch.ones_like(inputs)
            )
        
        # Decodificar el resultado completo
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extraer solo la parte nueva (sin el input original)
        new_text = generated_text[len(input_text):].strip()
        
        if not new_text:
            return "No se pudo generar texto adicional."
        
        return new_text
        
    except Exception as e:
        return f"Error al generar texto: {str(e)}"

def create_autocomplete_interface():
    """
    Crea la interfaz de autocompletar dentro de gr.Blocks()
    """
    
    with gr.Blocks(title="Autocompletar Código") as demo:
        
        gr.Markdown("# 🤖 Autocompletar Código")
        gr.Markdown("Escribe el inicio de tu código y la IA lo completará por ti.")
        
        with gr.Tab("Autocompletar"):
            with gr.Row():
                with gr.Column():
                    input_textbox = gr.Textbox(
                        label="Código a completar",
                        placeholder="def fibonacci(n):",
                        lines=5,
                        max_lines=10
                    )
                    
                    generate_btn = gr.Button("Completar Código", variant="primary")
                    
                with gr.Column():
                    output_textbox = gr.Textbox(
                        label="Código generado",
                        placeholder="Aquí aparecerá la continuación del código...",
                        lines=5,
                        max_lines=10,
                        interactive=False
                    )
            
            # Conectar el botón con la función
            generate_btn.click(
                fn=autocomplete_text,
                inputs=[input_textbox],
                outputs=[output_textbox]
            )
            
            # También permitir Enter para generar
            input_textbox.submit(
                fn=autocomplete_text,
                inputs=[input_textbox],
                outputs=[output_textbox]
            )
            
        # Pestaña adicional con ejemplos
        with gr.Tab("Ejemplos"):
            gr.Markdown("""
            ### Ejemplos de uso:
            
            **Entrada:** "def fibonacci(n):"
            **Salida:** "\\n    if n <= 1:\\n        return n\\n    return fibonacci(n-1) + fibonacci(n-2)"
            
            **Entrada:** "for i in range("
            **Salida:** "10):\\n    print(i)"
            
            **Entrada:** "import pandas as pd\\ndf = pd.read_csv("
            **Salida:** "'data.csv')\\nprint(df.head())"
            
            **Entrada:** "class Calculator:"
            **Salida:** "\\n    def __init__(self):\\n        pass"
            """)
    
    return demo

# Crear y lanzar la aplicación
if __name__ == "__main__":
    print("Iniciando aplicación de autocompletar código...")
    
    # Crear la interfaz
    app = create_autocomplete_interface()
    
    # Lanzar la aplicación
    app.launch(
        share=False,  # Cambiar a True si quieres compartir públicamente
        server_name="0.0.0.0",  # Permite acceso desde otras máquinas en la red local
        server_port=7860,  # Puerto por defecto de Gradio
        show_error=True,
        debug=False
    )