from __future__ import annotations
import asyncio
from typing import List, Optional, Dict, Any, TYPE_CHECKING, Literal
if TYPE_CHECKING:
    from ..llm_service import LLMService
import functools
import queue
import json
from dataclasses import dataclass, field

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage
from langgraph.graph import StateGraph, END
from rich.console import Console, Group
from rich.markdown import Markdown
from rich.panel import Panel
from rich.live import Live
from rich.padding import Padding
from rich.text import Text

from kogniterm.core.agent_state import AgentState
from kogniterm.terminal.themes import ColorPalette, Icons

console = Console()

# --- Estado Extendido para Deep Research ---
class DeepResearchState(AgentState):
    research_plan: List[str] = field(default_factory=list)
    findings: List[Dict[str, Any]] = field(default_factory=list)
    iteration_count: int = 0
    max_iterations: int = 5
    current_focus: str = ""

# --- Prompts ---
def get_deep_research_system_prompt(llm_service: LLMService) -> str:
    prompt = """Eres **KogniDeepResearcher**, motor de investigación técnica de élite.

Proceso iterativo:
1. **Planificación**: Desglosa la consulta en sub-preguntas lógicas.
2. **Exploración**: Usa búsquedas web, análisis de código y GitHub.
3. **Síntesis**: Crea un informe técnico con citas, código y arquitectura.

**Reglas:**
- Si una búsqueda es poco clara, cambia los términos o busca en archivos relacionados.
- Investiga referencias desconocidas. No asumas.
- Respalda cada afirmación con código, documentación o búsquedas.
- Usa Mermaid para diagramas y tablas para comparativas.
"""
    if not llm_service.is_thinking_model():
        prompt += """- Evalúa críticamente la información. Si hay contradicciones, investiga más.
- Proceso mental: ¿Qué necesito saber? ¿Dónde está? ¿Confirma mi hipótesis? ¿Qué nuevas preguntas surgen?
"""
    return prompt

DEEP_RESEARCH_SYSTEM_PROMPT = get_deep_research_system_prompt(LLMService(use_multi_provider=False)) if 'LLMService' in globals() else ""

# --- Nodos del Grafo ---

def planning_node(state: AgentState, llm_service: LLMService, terminal_ui=None):
    """Genera un plan de investigación e inyecta el contexto en el historial."""
    is_tui = getattr(terminal_ui, "is_tui", False)
    current_console = terminal_ui.console if terminal_ui else console

    last_message = state.messages[-1].content if state.messages else ""
    
    prompt = f"""Tarea: '{last_message}'
Genera un plan de investigación breve. Responde SOLO con JSON:
{{"plan": ["paso 1", "paso 2", ...]}}"""
    
    plan_steps = []
    try:
        response = llm_service.invoke_structured_output(
            prompt=prompt,
            system_prompt="Eres un planificador de investigación. Responde solo JSON.",
            response_model=None
        )
        content = response.content if hasattr(response, 'content') else str(response)
        if "```json" in content:
            content = content.split("```json")[1].split("```")[0].strip()
        elif "```" in content:
            content = content.split("```")[1].split("```")[0].strip()
        data = json.loads(content)
        plan_steps = data.get("plan", [])
    except Exception:
        plan_steps = []

    if plan_steps:
        # Mostrar el plan en la TUI/consola
        from kogniterm.terminal.themes import ColorPalette, Icons
        plan_text = Text()
        plan_text.append(f"\n📌 Plan de investigación:\n", style=f"bold {ColorPalette.PRIMARY_LIGHT}")
        for i, step in enumerate(plan_steps, 1):
            plan_text.append(f"  {i}. {step}\n", style="dim")
        
        if is_tui and terminal_ui:
            terminal_ui.update_live(Padding(plan_text, (0, 4)))
            terminal_ui.stop_live()
        else:
            current_console.print(Padding(plan_text, (0, 4)))

        # Inyectar el plan como contexto para call_model
        plan_context = "Plan de investigación a seguir:\n" + "\n".join(f"{i}. {s}" for i, s in enumerate(plan_steps, 1))
        state.messages.append(HumanMessage(content=plan_context))

    return {"messages": state.messages}

def research_node(state: DeepResearchState, llm_service: LLMService, terminal_ui: Optional[TerminalUI] = None, interrupt_queue: Optional[queue.Queue] = None):
    """Ejecuta herramientas para el foco actual."""
    current_console = terminal_ui.console if terminal_ui else console

    state.iteration_count += 1
    focus = state.current_focus or "General"
    
    # Notificación minimalista de iteración
    try:
        from kogniterm.terminal.themes import ColorPalette
        status_text = Text()
        status_text.append("🔍 ", style=ColorPalette.SECONDARY)
        status_text.append("Foco: ", style="dim")
        status_text.append(focus, style="italic")
        status_text.append(f" ({state.iteration_count}/{state.max_iterations})", style="dim")
        current_console.print(Padding(status_text, (1, 4)))
    except ImportError:
        current_console.print(f"\n[dim]🔍 {focus} ({state.iteration_count}/{state.max_iterations})[/dim]")
    
    # Aquí llamamos al modelo para que use herramientas
    messages = [SystemMessage(content=get_deep_research_system_prompt(llm_service))] + state.messages
    messages.append(HumanMessage(content=f"Enfócate ahora en investigar: {focus}. Utiliza herramientas si es necesario."))
    
    return state

def synthesis_node(state: DeepResearchState, llm_service: LLMService, terminal_ui: Optional[TerminalUI] = None):
    """Compila todos los hallazgos en el reporte final."""
    current_console = terminal_ui.console if terminal_ui else console

    all_findings_summary = ""
    for idx, finding in enumerate(state.findings):
        all_findings_summary += f"### Hallazgo {idx+1}: {finding.get('focus')}\n{finding.get('content')}\n\n"
        
    prompt = f"""Como experto Sintetizador Técnico, utiliza toda la información recopilada para crear el informe final.
    
    HISTORIAL DE INVESTIGACIÓN:
    {all_findings_summary}
    
    CONSULTA ORIGINAL: {state.messages[0].content}
    
    REQUISITOS DEL INFORME:
    1. Estructura clara (Introducción, Arquitectura/Flujo, Detalles Técnicos, Conclusión).
    2. Cita archivos, líneas de código y URLs encontradas.
    3. Usa Mermaid para diagramas de secuencia o flujo.
    4. Proporcina recomendaciones o advertencias 'KogniInsight'.
    """
    
    # Llamada final al modelo
    response = llm_service.invoke(history=[SystemMessage(content=get_deep_research_system_prompt(llm_service)), HumanMessage(content=prompt)])
    
    # Recolectar la respuesta y mostrarla visualmente
    full_content = ""
    
    from rich.status import Status
    with Status("[dim]Sintetizando reporte final...[/dim]", spinner="dots") as status:
        for part in response:
            if isinstance(part, AIMessage):
                final_ai_message = part
            elif isinstance(part, str):
                full_content += part
            
    # Mostrar el resultado final de forma limpia
    current_console.print(Padding(Markdown(f"## 🔬 Informe de Investigación\n\n{full_content}"), (1, 4)))
            
    state.messages.append(AIMessage(content=f"## 🔬 Informe de Deep Research\n\n{full_content}"))
    return state

# --- Implementación principal conceptualmente basada en agentes previos pero con lógica mejorada ---

def call_deep_model_node(state: AgentState, llm_service: LLMService, terminal_ui: Optional[TerminalUI] = None, interrupt_queue: Optional[queue.Queue] = None):
    """Llama al LLM de Deep Research con soporte completo para TUI/CLI."""
    current_console = terminal_ui.console if terminal_ui else console
    is_tui = getattr(terminal_ui, "is_tui", False)

    # Inyectar el prompt de Deep Research al principio si no está
    messages = [SystemMessage(content=get_deep_research_system_prompt(llm_service))] + state.messages
    
    full_response_content = ""
    full_thinking_content = ""
    final_ai_message = None
    text_streamed = False

    # Importar componentes visuales
    try:
        from kogniterm.terminal.visual_components import create_processing_spinner
        from kogniterm.terminal.themes import ColorPalette, Icons
        spinner = create_processing_spinner()
    except ImportError:
        from rich.spinner import Spinner
        spinner = Spinner("dots", text="[dim]Investigando...[/dim]")
        class Icons: THINKING = "🤔"; RESEARCH = "🔍"
        class ColorPalette: PRIMARY_LIGHT = "cyan"; SECONDARY = "blue"; GRAY_800 = "#333333"; GRAY_600 = "#666666"; GRAY_900 = "#1e1e1e"

    # Iniciar KeyboardHandler para detectar ESC (solo CLI)
    kh = None
    if not is_tui:
        from kogniterm.terminal.keyboard_handler import KeyboardHandler
        kh = KeyboardHandler(interrupt_queue)
        kh.start()

    try:
        import contextlib
        if not is_tui:
            live_context = Live(spinner, console=current_console, screen=False, refresh_per_second=10)
        else:
            @contextlib.contextmanager
            def dummy_live(): 
                yield type('DummyLive', (), {'update': lambda self, x: None})()
            live_context = dummy_live()

        with live_context as live:
            TUI_BG = ColorPalette.GRAY_900 if 'ColorPalette' in globals() else "#1e1e1e"

            def update_display():
                renderables = []
                if full_thinking_content:
                    if is_tui:
                        thinking_content = Markdown(full_thinking_content)
                        
                        # Calcular ancho basado en el chat_log
                        content_width = 78
                        try:
                            if hasattr(terminal_ui, 'app') and hasattr(terminal_ui.app, 'chat_log'):
                                log_width = terminal_ui.app.chat_log.size.width
                                if log_width > 0:
                                    content_width = max(log_width - 4, 40)
                        except:
                            pass
                        
                        thought_panel = Panel(
                            thinking_content,
                            title=f"[bold {ColorPalette.PRIMARY_LIGHT}]{Icons.THINKING} KogniResearcher Pensando...[/]",
                            border_style=ColorPalette.GRAY_700,
                            style=f"dim {ColorPalette.GRAY_500} on {TUI_BG}",
                            padding=(0, 2),
                            width=content_width
                        )

                        renderables.append(thought_panel)
                    else:
                        renderables.append(Panel(
                            Markdown(full_thinking_content),
                            title=f"[bold {ColorPalette.PRIMARY_LIGHT}]{Icons.THINKING} Investigación en curso...[/bold {ColorPalette.PRIMARY_LIGHT}]",
                            border_style=ColorPalette.PRIMARY_LIGHT,
                            padding=(0, 1),
                            dim=True
                        ))
                
                if full_response_content:
                    if full_thinking_content:
                        renderables.append(Text(""))
                    renderables.append(Markdown(full_response_content))
                
                if is_tui:
                    group = Group(*renderables)
                    terminal_ui.update_live(Padding(group, (0, 4)))
                else:
                    final_renderable = Padding(Group(*renderables), (0, 4)) if renderables else spinner
                    live.update(final_renderable)

            for part in llm_service.invoke(history=messages, interrupt_queue=interrupt_queue):
                if isinstance(part, AIMessage):
                    final_ai_message = part
                elif isinstance(part, str):
                    if part.startswith("__THINKING__:") or part.startswith("THINKING:"):
                        prefix = "__THINKING__:" if part.startswith("__THINKING__:") else "THINKING:"
                        full_thinking_content += part[len(prefix):]
                        update_display()
                    else:
                        full_response_content += part
                        text_streamed = True
                        update_display()
                
                if (interrupt_queue and not interrupt_queue.empty()) or llm_service.stop_generation_flag:
                    break

            if is_tui:
                terminal_ui.stop_live()
    finally:
        if kh: kh.stop()

    if final_ai_message:
        # Asegurar que el contenido procesado se guarde en el mensaje
        if not final_ai_message.content and full_response_content:
            final_ai_message.content = full_response_content
        state.messages.append(final_ai_message)
        state.save_history(llm_service)
            
    return {"messages": state.messages}

def create_deep_researcher(llm_service: LLMService, terminal_ui: Any = None, interrupt_queue: Optional[queue.Queue] = None):
    from .code_agent import execute_tool_node, should_continue
    
    workflow = StateGraph(AgentState)

    workflow.add_node("planning", functools.partial(planning_node, llm_service=llm_service, terminal_ui=terminal_ui))
    workflow.add_node("call_model", functools.partial(call_deep_model_node, llm_service=llm_service, terminal_ui=terminal_ui, interrupt_queue=interrupt_queue))
    workflow.add_node("execute_tool", functools.partial(execute_tool_node, llm_service=llm_service, terminal_ui=terminal_ui, interrupt_queue=interrupt_queue))

    workflow.set_entry_point("planning")
    workflow.add_edge("planning", "call_model")

    workflow.add_conditional_edges(
        "call_model",
        should_continue,
        {
            "execute_tool": "execute_tool",
            END: END
        }
    )

    workflow.add_edge("execute_tool", "call_model")

    return workflow.compile()
