import os
import gradio as gr
from google import genai
from google.genai import types

api_key = os.environ.get("GEMINI_API_KEY")

client = None
if api_key:
    client = genai.Client(api_key=api_key)

SYSTEM_PROMPT = """
Eres un Asistente de Planificación Maestro experto. Tu objetivo es ayudar al usuario a planificar proyectos en tres áreas específicas:
1. Planificación de Producción (Demanda, Inventario, Capacidades).
2. Planificación de Proyectos de Ingeniería de Datos.
3. Planificación de Proyectos de IA Generativa.

Al inicio de la conversación, DEBES presentarte brevemente y ofrecer estas 3 capacidades de planificación.

REGLAS GENERALES:
- Habla siempre en Español.
- Sé profesional, analítico y servicial.
- Si el usuario elige una opción, guíalo paso a paso solicitando la información necesaria.
- Solicita información paso a paso en lugar de pedir todo de una vez.

MODULO 1: PLANIFICACIÓN DE PRODUCCIÓN
Si el usuario elige esta opción, debes solicitar:
1. Producto(s) a planificar.
2. Demanda esperada por periodo (días, semanas o meses).
3. Inventario disponible actual (terminado y materias primas).
4. Capacidad de producción (máximo por periodo).
5. Tiempo de producción por unidad o lote.
6. Prioridades o restricciones opcionales.

Una vez tengas los datos, aplica uno de estos métodos (explica cuál usas):
- Clásico (MRP): Cálculo de necesidades netas y fechas.
- Temporal: Basado en cronogramas y carga de capacidad.
- Probabilístico: Considera incertidumbre y propone buffers.

RESULTADOS ESPERADOS:
- Plan de producción por periodo (Ej: Semana 1: X unidades, Semana 2: Y unidades).
- Necesidades de materiales (Cantidades de materias primas).
- Alertas de capacidad o stock insuficiente.
- Resumen final (producción total, inventario final estimado, posibles retrasos).

MODULO 2: INGENIERÍA DE DATOS
Ayuda a planificar pipelines ETL, lagos de datos, modelado de almacenes, etc. Solicita objetivos, fuentes de datos, volumen y frecuencia.

MODULO 3: IA GENERATIVA
Ayuda a planificar sistemas RAG, fine-tuning, agentes, etc. Solicita casos de uso, modelos base, datos de entrenamiento/contexto y métricas de evaluación.

IMPORTANTE: No inventes datos si el usuario no los proporciona. Pídelos amablemente.
"""

def extract_text(content):
    if isinstance(content, str):
        return content
    
    if isinstance(content, list):
        texts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                texts.append(item.get("text", ""))
        return " ".join(texts)
    
    return str(content)

def predict(message, history):
    if not client:
        yield "Error: No se ha configurado la clave de API de Gemini (GEMINI_API_KEY)."
        return
    
    contents = []
    
    for msg in history:
        if isinstance(msg, dict):
            role = msg.get("role")
            raw_text = msg.get("content", "")
            text = extract_text(raw_text)

            if role == "user":
                contents.append(
                    types.Content(
                        role="user",
                        parts=[types.Part.from_text(text=text)]
                    )
                )
            elif role == "assistant":
                contents.append(
                    types.Content(
                        role="model",
                        parts=[types.Part.from_text(text=text)]
                    )
                )
        
        # Añadir el mensaje actual
    contents.append(types.Content(role="user", parts=[types.Part.from_text(text=message)]))
    

    try:
        generate_content_config = types.GenerateContentConfig(
            system_instruction=SYSTEM_PROMPT,
            temperature=0.7,
        )

        response = client.models.generate_content_stream(
            model="gemini-2.0-flash",
            contents=contents,
            config=generate_content_config,
        )

        partial_message = ""
        for chunk in response:
            if chunk.text:
                partial_message += chunk.text
                yield partial_message
    except Exception as e:
        yield f"Ocurrió un error: {str(e)}"

demo = gr.ChatInterface(
    predict,
    title="Planificador Maestro AI",
    description="Asistente inteligente para planificación de producción, ingeniería de datos e IA generativa.",
    examples=[
        "Quiero planificar la producción de galletas",
        "Ayúdame a planificar un proyecto de ingeniería de datos",
        "Necesito un plan para implementar un sistema RAG"
    ],
    flagging_mode="manual",
    flagging_options=["Like", "Spam", "Inappropriate", "Other"],
    save_history=True,
)

if __name__ == "__main__":
    demo.launch()