import gradio as gr
from transformers import pipeline
from deep_translator import GoogleTranslator

# Cargar el modelo de Hugging Face para text-to-SQL
model = pipeline("text2text-generation", model="tscholak/text-to-sql")

# Esquema de base de datos
SCHEMA = """
Database schema:
Table bodegas(Id, Nombre, Encargado, Telefono, Email, Direccion, Horario, Regional, Latitud, Longitud)
Table maestra(CodigoSap, Descripcion, Grupo, Agrupador, Marca, Parte, Operacion, Componente)
"""

# Función para generar la consulta SQL
def generar_sql(pregunta_espanol):
    try:
        # Traducir pregunta a inglés
        pregunta_ingles = GoogleTranslator(source="es", target="en").translate(pregunta_espanol)
        
        # Crear el prompt para el modelo
        prompt = f"{SCHEMA}\ntranslate English to SQL: {pregunta_ingles}"

        # Generar la consulta SQL usando el modelo de Hugging Face
        result = model(prompt)
        sql = result[0]['generated_text']
        
        return sql
    except Exception as e:
        return f"Error: {str(e)}"

# Interfaz Gradio para desplegar
iface = gr.Interface(
    fn=generar_sql,
    inputs=gr.Textbox(lines=3, label="Pregunta en español"),
    outputs=gr.Textbox(label="Consulta SQL generada"),
    title="Generador de SQL",
    description="Escribe una pregunta en español sobre la base de datos y obtén la consulta SQL.",
)

iface.launch()