import gradio as gr 
import os
from typing import TypedDict
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI

class MigrationState(TypedDict):
    codigo_original: str
    codigo_traducido: str
    analisis_seguridad: str

# Nota: El código trabaja actualmente con el modelo DeepSeek V3.1. 
# Sin embargo, se recomienda trabajar con Qwen2.5-Coder 32B para un mejor análisis de código. 
llm = ChatOpenAI(
    api_key = os.environ.get("API_KEY"),
    base_url="https://api.fireworks.ai/inference/v1",
    model="accounts/fireworks/models/deepseek-v3p1" 
)

def nodo_traductor(state: MigrationState):
    prompt = f"Actúa como un experto en AMD ROCm. Traduce este código NVIDIA CUDA a AMD HIP. Solo devuelve el código final, sin explicaciones:\n\n{state['codigo_original']}"
    respuesta = llm.invoke(prompt)
    
    codigo_limpio = respuesta.content
    if codigo_limpio.startswith("```"):
         lineas = codigo_limpio.split("\n")
         codigo_limpio = "\n".join(lineas[1:-1])
         
    return {"codigo_traducido": codigo_limpio}

def nodo_auditor(state: MigrationState):
    prompt = f"Revisa este código AMD HIP. Confirma si la migración fue correcta y advierte sobre posibles fugas de memoria. Sé breve (máximo 3 líneas):\n\n{state['codigo_traducido']}"
    respuesta = llm.invoke(prompt)
    return {"analisis_seguridad": respuesta.content}

workflow = StateGraph(MigrationState)
workflow.add_node("traductor", nodo_traductor)
workflow.add_node("auditor", nodo_auditor)
workflow.add_edge(START, "traductor")
workflow.add_edge("traductor", "auditor")
workflow.add_edge("auditor", END)

agente_rocm = workflow.compile()

DUMMY_CUDA_CODE = """#include <iostream>

__global__ void vectorAdd(const float *A, const float *B, float *C, int numElements) {
    int i = blockDim.x * blockIdx.x + threadIdx.x;
    if (i < numElements) {
        C[i] = A[i] + B[i];
    }
}

int main() {
    return 0;
}
"""

def leer_archivo(archivo):
    if archivo is not None:
        try:
            with open(archivo.name, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            return f"Error al leer el archivo: {str(e)}"
    return ""

def chat_con_agente(mensaje_usuario, historial):
    if not historial:
        historial = []
    
    prompt_chat = f"""Eres el Ingeniero Líder de ROCm-Sentinel. 
    El usuario te da esta instrucción para la migración: '{mensaje_usuario}'. 
    Responde en una frase corta y técnica cómo integrarás esta instrucción en el proceso de traducción a AMD HIP."""
    
    try:
        respuesta = llm.invoke(prompt_chat)
        historial.append({"role": "user", "content": mensaje_usuario})
        historial.append({"role": "assistant", "content": respuesta.content})
    except Exception:
        historial.append({"role": "user", "content": mensaje_usuario})
        historial.append({"role": "assistant", "content": "Entendido. Aplicaré esas optimizaciones técnicas al código."})
    
    return "", historial

def simular_migracion(cuda_code, historial):
    if not historial: historial = []
    if not cuda_code:
         yield "Error: Ingresa código.", historial, "Desconocido", ""
         return
        
    historial.append({"role": "assistant", "content": "🚀 **Proceso iniciado.** El Agente Traductor está analizando los kernels de CUDA..."})
    yield "Estado: Analizando Kernels...", historial, "Analizando...", ""
    
    try:
        resultado = agente_rocm.invoke({"codigo_original": cuda_code})
        rocm_code = resultado["codigo_traducido"]
        reporte_auditor = resultado["analisis_seguridad"]
        
        historial.append({"role": "assistant", "content": f"✅ **Migración Exitosa.**\n\n🛡️ **Auditoría de Seguridad:**\n{reporte_auditor}"})
        yield "Estado: Completado con éxito", historial, "NVIDIA CUDA (nvcc) -> AMD HIP (hipcc)", rocm_code
        
    except Exception as e:
        historial.append({"role": "assistant", "content": f"❌ Error en el flujo: {str(e)}"})
        yield "Estado: Error crítico", historial, "Error", ""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        <div style='text-align: center;'>
            <h1 style='color: #FF3232;'>ROCm-Sentinel</h1>
            <p><strong>Advanced CUDA to AMD ROCm Migration Agent</strong></p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            archivo_subida = gr.File(
                label="📁 Subir archivo fuente", 
                file_types=[".cu", ".cpp", ".txt", ".h"]
            )
            gr.Markdown("### 📜 Código Origen (CUDA)")
            input_code = gr.Code(
                label="Editor de entrada", 
                language="cpp", 
                lines=15, 
                value=DUMMY_CUDA_CODE
            )
            
        with gr.Column(scale=1):
            lenguaje_detectado = gr.Textbox(
                label="🔍 Firma / Lenguaje Detectado", 
                interactive=False
            )
            estado_migracion = gr.Textbox(
                label="⚡ Monitor de Proceso (Estado)",
                interactive=False,
                placeholder="Listo para procesar..."
            )
            gr.Markdown("### 🛠️ Código Migrado (AMD HIP)")
            output_code = gr.Code(
                label="Resultado de la migración", 
                language="cpp", 
                lines=15, 
                interactive=False
            )
            
    with gr.Row():
        migrate_btn = gr.Button(
            "🚀 INICIAR MIGRACIÓN SEGURA A ROCm", 
            variant="primary", 
            size="lg"
        )
            
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🤖 Centro de Control del Agente")
            chatbot = gr.Chatbot(
                label="Interacción con el Agente de IA", 
                height=300
            )
            chat_input = gr.Textbox(
                label="Instrucciones adicionales para el Agente", 
                placeholder="Ej: 'Usa memoria unificada' o 'Optimiza para hilos específicos'..."
            )
            
    archivo_subida.change(fn=leer_archivo, inputs=archivo_subida, outputs=input_code)
    
    chat_input.submit(
        fn=chat_con_agente, 
        inputs=[chat_input, chatbot], 
        outputs=[chat_input, chatbot]
    )
    
    migrate_btn.click(
        fn=simular_migracion, 
        inputs=[input_code, chatbot], 
        outputs=[estado_migracion, chatbot, lenguaje_detectado, output_code]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Monochrome())