# app/xulcan/kernel/runtime.py

"""Kernel Engine implementation for Xulcan Agents.

This module provides the ProtoKernel, the central execution engine that drives
the agent's cognitive loop. It acts as a State Machine Driver, evaluating the
current FSM state, executing the associated side-effects (LLM calls, tool execution),
and computing the next valid state.

Architecture Philosophy:
    - The FSM (KernelState) defines the *Static Graph* of valid operations.
    - The Ledger (LedgerRepository) records the *Dynamic DAG* of historical events.
    - The Kernel (ProtoKernel) is the *Driver* that unrolls the loop.
"""

import uuid
import logging
import json
import asyncio
from typing import List, Optional, Any, Tuple, Dict

# Core & Primitives
from xulcan.core.primitives import MachineID
from xulcan.core.economics import UsageStats

# FSM & Interfaces
from xulcan.kernel.states import KernelState, validate_transition, is_terminal_state
from xulcan.kernel.interfaces import LedgerRepository, ToolExecutor
# Nota: SystemRegistry tipado se asume compatible con LLMRegistry
from xulcan.system.registry import ProviderRegistry 
from xulcan.kernel.interfaces import LLMProvider

from xulcan.kernel.environment import SystemEnvironment
from xulcan.utils.template import render_template  # <--- MAGIA TRANSVERSAL

# Blueprint
from xulcan.blueprint.schema import AgentBlueprint

# Ledger Events (Santuario de Datos)
from xulcan.history.events import (
    RunCreated, 
    StepStarted, 
    ModelRequest, 
    ModelResponse, 
    ToolExecution,
    ToolOutput,
    RunCompleted, 
    RunFailed,
    StepType
)

# Protocol Messages
from xulcan.protocol.message import (
    UnifiedMessage, 
    SystemMessage,
    UserMessage, 
    AssistantMessage, 
    ToolMessage
)

logger = logging.getLogger("xulcan.kernel")


class ProtoKernel:
    """The FSM Driver (Engine) for Xulcan Agents."""

    def __init__(
        self, 
        repository: LedgerRepository, 
        llm_registry: ProviderRegistry[LLMProvider], 
        tool_executor: ToolExecutor,
        environment: Optional[SystemEnvironment] = None,
        llm_credentials: Optional[Dict[str, str]] = None  # <--- CREDENCIALES
    ):
        self.repo = repository
        self.registry = llm_registry
        self.tool_executor = tool_executor
        self.environment = environment
        self.llm_credentials = llm_credentials or {}

    def _log(self, emoji: str, msg: str, **kwargs: Any):
        extra = f" | {json.dumps(kwargs, default=str)}" if kwargs else ""
        logger.info(f"{emoji} {msg}{extra}")


    async def execute_run(
        self, 
        blueprint: AgentBlueprint, 
        user_input: str, 
        parent_id: Optional[str] = None
    ) -> Tuple[str, UnifiedMessage]:
    
        run_id = str(uuid.uuid4())
        context: List[UnifiedMessage] =[UserMessage(content=user_input)]
        
        loop_counter = 0
        MAX_LOOPS = 20
        final_response: Optional[UnifiedMessage] = None
        cumulative_usage = UsageStats.zero()
        
        current_state = KernelState.CREATED
        sequence_number = 1
        
        def next_seq() -> int:
            nonlocal sequence_number
            val = sequence_number
            sequence_number += 1
            return val

        def transition(new_state: KernelState):
            nonlocal current_state
            self._log("⚡", f"[FSM] {current_state.name} -> {new_state.name}")
            validate_transition(current_state, new_state)
            current_state = new_state

        self._log("🎬", f"STARTING RUN: {run_id}", agent=blueprint.id)

        try:
            # 1. INITIALIZATION
            meta = {"parent_run_id": parent_id} if parent_id else {}
            await self.repo.append(RunCreated(
                run_id=run_id,
                sequence_number=next_seq(),
                step_index=0,
                agent_id=blueprint.id,
                agent_version=blueprint.version,
                user_input=list(context),
                metadata=meta
            ))
            
            # 2. THE FSM DRIVER LOOP
            while not is_terminal_state(current_state):
                
                if current_state == KernelState.PREPARING_CONTEXT:
                    loop_counter += 1
                    self._log("🔄", f"--- REASONING LOOP {loop_counter} ---")

                # --- STATE HANDLERS ---
                
                if current_state == KernelState.CREATED:
                    transition(KernelState.HYDRATING)

                elif current_state == KernelState.HYDRATING:
                    # (Nota de Arquitectura: Pronto delegaremos esto al ContextEngine)
                    if parent_id:
                        self._log("💧", f"Hydrating context from Ledger Session (parent: {parent_id})")
                        all_runs_events =[]
                        current_ancestor_id = parent_id

                        while current_ancestor_id:
                            past_events = await self.repo.get_events(current_ancestor_id)
                            if not past_events:
                                break
                            all_runs_events = past_events + all_runs_events
                            run_created = next((e for e in past_events if isinstance(e, RunCreated)), None)
                            current_ancestor_id = run_created.metadata.get("parent_run_id") if run_created else None

                        historical_messages =[]
                        for event in all_runs_events:
                            if isinstance(event, RunCreated):
                                historical_messages.extend(event.user_input)
                            elif isinstance(event, ModelResponse):
                                historical_messages.append(event.message)
                            elif isinstance(event, ToolOutput):
                                historical_messages.append(ToolMessage(
                                    tool_call_id=event.tool_call_id,
                                    name=event.tool_name,
                                    content=event.output
                                ))

                        context = historical_messages + context
                        self._log("🧠", f"Context hydrated with {len(historical_messages)} past messages.")

                    transition(KernelState.HYDRATED)

                elif current_state == KernelState.HYDRATED:
                    transition(KernelState.CHECKING_BUDGET)

                elif current_state == KernelState.CHECKING_BUDGET:
                    transition(KernelState.PREPARING_CONTEXT)

                elif current_state == KernelState.PREPARING_CONTEXT:
                    await self.repo.append(StepStarted(
                        run_id=run_id,
                        sequence_number=next_seq(),
                        step_index=loop_counter,
                        step_type=StepType.INFERENCE
                    ))
                    transition(KernelState.CALLING_MODEL)

                elif current_state == KernelState.CALLING_MODEL:
                    # 1. Preparar el Prompt del Sistema
                    if not any(isinstance(m, SystemMessage) for m in context):
                        raw_prompt = blueprint.system_prompt
                        rendered_prompt = raw_prompt

                        # 🪄 MAGIA JINJA2 PURIFICADA
                        if self.environment and self.environment.state_store:
                            keys = await self.environment.state_store.keys(run_id)
                            memory_dict = {k: await self.environment.state_store.get(run_id, k) for k in keys}
                            rendered_prompt = render_template(raw_prompt, memory_dict)
                            self._log("🪄", "System prompt rendered with memory context.")

                        sys_msg = SystemMessage(content=rendered_prompt)
                        context.insert(0, sys_msg)

                    # 2. Obtener Schemas de Herramientas
                    active_tool_names = blueprint.llm_tools
                    tool_defs = None
                    if active_tool_names and hasattr(self.tool_executor, 'get_definitions'):
                        tool_defs = self.tool_executor.get_definitions(active_tool_names)
                    
                    # 3. 🚀 LAZY LOADING: INSTANCIAR EL LLM AL VUELO
                    AdapterClass = self.registry.get_class(blueprint.model_provider)
                    llm_secret = self.llm_credentials.get(blueprint.model_provider)
                    
                    # (Fallback temporal hasta que pasemos todo a ConfigSchema)
                    if blueprint.model_provider == "ollama":
                        adapter = AdapterClass(host=llm_secret)
                    else:
                        adapter = AdapterClass(api_key=llm_secret)
                        
                    self._log("🧠", "Thinking...", model=blueprint.model_name)
                    
                    # 4. Guardar Request en Ledger
                    await self.repo.append(ModelRequest(
                        run_id=run_id,
                        sequence_number=next_seq(),
                        step_index=loop_counter,
                        provider=blueprint.model_provider,
                        model=blueprint.model_name,
                        prompt_messages=list(context),
                        parameters={"temperature": blueprint.temperature}
                    ))

                    # 5. LLAMADA FINAL DINÁMICA
                    llm_response = await adapter.generate(
                        model=blueprint.model_name,
                        messages=context,
                        tools=tool_defs,
                        temperature=blueprint.temperature,
                        max_tokens=blueprint.max_tokens
                    )
                    
                    cumulative_usage = cumulative_usage + llm_response.usage

                    assistant_msg = AssistantMessage(
                        content=llm_response.content,
                        tool_calls=llm_response.tool_calls
                    )
                    
                    if assistant_msg.tool_calls:
                        tools_str = ", ".join([t.name for t in assistant_msg.tool_calls])
                        self._log("💡", f"Model decided to act: {tools_str}")
                    else:
                        self._log("🗣️", f"Model response: {str(assistant_msg.content)[:50]}...")

                    await self.repo.append(ModelResponse(
                        run_id=run_id,
                        sequence_number=next_seq(),
                        step_index=loop_counter,
                        message=assistant_msg,
                        usage=llm_response.usage,
                        latency_ms=llm_response.usage.latency_ms
                    ))
                    context.append(assistant_msg)

                    transition(KernelState.PROCESSING_RESPONSE)

                elif current_state == KernelState.PROCESSING_RESPONSE:
                    last_msg = context[-1]
                    if isinstance(last_msg, AssistantMessage) and last_msg.tool_calls:
                        transition(KernelState.PARSING_TOOL_ARGS)
                    else:
                        final_response = last_msg
                        transition(KernelState.COMPLETED)

                elif current_state == KernelState.PARSING_TOOL_ARGS:
                    transition(KernelState.CHECKING_POLICY)

                elif current_state == KernelState.CHECKING_POLICY:
                    transition(KernelState.EXECUTING_TOOL)

                elif current_state == KernelState.EXECUTING_TOOL:
                    last_msg = context[-1]
                    
                    # 🔒 INYECCIÓN INMUTABLE
                    updated_tool_calls =[]
                    for t_call in last_msg.tool_calls:
                        new_context = {"run_id": run_id, "parent_run_id": parent_id}
                        updated_call = t_call.model_copy(update={"context": new_context})
                        updated_tool_calls.append(updated_call)

                        await self.repo.append(ToolExecution(
                            run_id=run_id,
                            sequence_number=next_seq(),
                            step_index=loop_counter,
                            tool_call=updated_call,
                            tool_input=updated_call.arguments  
                        ))
                    
                    last_msg = last_msg.model_copy(update={"tool_calls": updated_tool_calls})
                    context[-1] = last_msg
                    
                    self._log("🛠️", "Delegating to Tool Executor...")
                    tool_results = await self.tool_executor.execute_batch(last_msg.tool_calls)
                    
                    for result in tool_results:
                        await self.repo.append(ToolOutput(
                            run_id=run_id,
                            sequence_number=next_seq(),
                            step_index=loop_counter,
                            tool_call_id=result.tool_call_id,
                            tool_name=result.name or "unknown",
                            output=result.content
                        ))
                        context.append(result)

                    # 🔄 TRANSICIÓN CORREGIDA: Va a presupuesto, NO al disco duro (Ledger)
                    transition(KernelState.CHECKING_BUDGET)
                    
                    if loop_counter >= MAX_LOOPS:
                         raise RuntimeError(f"Max reasoning loops ({MAX_LOOPS}) exceeded.")
                    
                else:
                    raise RuntimeError(f"Kernel stalled: No handler implemented for state {current_state}")

            # 3. COMPLETION & TEARDOWN
            if current_state == KernelState.COMPLETED:
                self._log("🏁", "RUN COMPLETED", tokens=cumulative_usage.total_tokens)
                await self.repo.append(RunCompleted(
                    run_id=run_id,
                    sequence_number=next_seq(),
                    step_index=loop_counter + 1,
                    final_response=final_response,
                    total_usage=cumulative_usage
                ))
                
                if self.environment and self.environment.state_store:
                    self._log("🧹", "Clearing Volatile Memory (StateStore)")
                    await self.environment.state_store.clear(run_id)

                return run_id, final_response

        except Exception as e:
            logger.error(f"💥 Kernel Panic: {str(e)}", exc_info=True)
            transition(KernelState.FAILED)
            
            try:
                await self.repo.append(RunFailed(
                    run_id=run_id,
                    sequence_number=next_seq(),
                    step_index=999, 
                    error_type="kernel_panic",
                    error_message=str(e)
                ))
            except Exception as repo_error:
                logger.critical(f"Doble fallo Ledger: {repo_error}")
            
            if self.environment and self.environment.state_store:
                await self.environment.state_store.clear(run_id)
            
            raise e