from __future__ import annotations

from typing import Any, Dict, List, Optional
from src.skills.models import SkillDefinition, SkillResult
from src.skills.lifecycle import SkillLifecycleManager, SkillLifecycleState
from src.skills.loader import LoadedSkillConfig, SkillLoader
from src.skills.registry import SkillRegistry
from src.runtime.execution_context import ExecutionContext
from src.runtime.metrics import RuntimeMetrics
import time


class AgentRuntime:
    """Generic LLM execution runtime slot with lifecycle management and context isolation."""

    def __init__(self, name: str, registry: SkillRegistry):
        self.name = name
        self.registry = registry
        self.loader = SkillLoader(registry)
        self.lifecycle = SkillLifecycleManager(slot_name=name)
        
        self.active_config: Optional[LoadedSkillConfig] = None
        self.prompt: Optional[str] = None
        self.active_tools: List[str] = []
        self.temporary_messages: List[Any] = []
        self.metrics = RuntimeMetrics()
        self.metrics.record("runtime.created", runtime=self.name)

    @property
    def active_skill(self) -> Optional[SkillDefinition]:
        return self.active_config.definition if self.active_config else None

    @property
    def state(self) -> SkillLifecycleState:
        return self.lifecycle.state

    def load_skill(self, skill_id: str, version: Optional[str] = None) -> LoadedSkillConfig:
        """Load a skill definition into this agent runtime slot."""
        started = time.perf_counter()
        if self.lifecycle.state != SkillLifecycleState.UNLOADED:
            self.flush()

        self.lifecycle.transition_to(SkillLifecycleState.LOADING)
        config = self.loader.load(skill_id, version=version)
        self.active_config = config
        self.prompt = config.instructions
        self.active_tools = list(config.allowed_tools)
        self.temporary_messages.clear()

        self.lifecycle.transition_to(SkillLifecycleState.ACTIVE, skill=config.definition)
        self.metrics.record(
            "skill.loaded",
            runtime=self.name,
            skill_id=config.definition.id,
            skill_version=config.definition.version,
            duration_ms=round((time.perf_counter() - started) * 1000, 3),
        )
        return config

    def execute(self, context: ExecutionContext, input_override: Optional[Dict[str, Any]] = None) -> SkillResult:
        """Execute the currently loaded skill within the runtime slot."""
        if self.lifecycle.state != SkillLifecycleState.ACTIVE or not self.active_config:
            raise RuntimeError(f"AgentRuntime '{self.name}' has no active skill loaded to execute.")

        self.lifecycle.transition_to(SkillLifecycleState.EXECUTING)
        skill_def = self.active_config.definition

        context.record_trace("skill_executing", {
            "runtime": self.name,
            "skill_id": skill_def.id,
            "version": skill_def.version,
        })

        # Base execution simulation/wrapper - to be connected to LLM provider
        input_query = (input_override.get("query") if input_override else None) or context.user_query
        result_output = {
            "status": "completed",
            "skill_id": skill_def.id,
            "response": f"Processed by {self.name} with skill {skill_def.id}: {input_query}",
        }

        result = SkillResult(
            skill_id=skill_def.id,
            skill_version=skill_def.version,
            output=result_output,
            metrics={"execution_slot": self.name, "context_tokens": 0, "prompt_tokens": 0},
        )

        self.metrics.record(
            "skill.executed",
            runtime=self.name,
            skill_id=skill_def.id,
            skill_version=skill_def.version,
            prompt_tokens=result.metrics.get("prompt_tokens", 0),
            context_tokens=result.metrics.get("context_tokens", 0),
        )
        context.metadata["runtime_metrics"] = self.metrics.snapshot()

        self.lifecycle.transition_to(SkillLifecycleState.ACTIVE, skill=skill_def)
        return result

    def flush(self) -> None:
        """Flush skill instructions, active tools, temporary messages, and reset to UNLOADED."""
        if self.lifecycle.state == SkillLifecycleState.UNLOADED:
            return

        if self.lifecycle.state == SkillLifecycleState.EXECUTING:
            self.lifecycle.transition_to(SkillLifecycleState.ACTIVE, skill=self.active_skill)

        started = time.perf_counter()
        previous_skill = self.active_skill.id if self.active_skill else None
        self.lifecycle.transition_to(SkillLifecycleState.FLUSHING)
        self.active_config = None
        self.prompt = None
        self.active_tools.clear()
        self.temporary_messages.clear()
        self.lifecycle.transition_to(SkillLifecycleState.UNLOADED)
        self.metrics.record(
            "skill.flushed",
            runtime=self.name,
            skill_id=previous_skill,
            duration_ms=round((time.perf_counter() - started) * 1000, 3),
        )