from __future__ import annotations

import asyncio
from collections import defaultdict

from nexus.core.events.bus import EventBus
from nexus.core.events.types import ReflectionCompleteEvent, SessionEndEvent, TaskCompleteEvent
from nexus.persona.evolver import PromptEvolver
from nexus.persona.reflector import PostTaskReflector, ReflectionRecord
from nexus.persona.role_assignor import RoleAssignor
from nexus.persona.skill_registry import SkillRegistry


class PersonaEngine:
    """Coordinates reflection, prompt evolution, and skill capture."""

    def __init__(
        self,
        bus: EventBus,
        reflector: PostTaskReflector,
        evolver: PromptEvolver,
        skill_registry: SkillRegistry,
        role_assignor: RoleAssignor,
        evolution_every_n_tasks: int = 5,
    ) -> None:
        self.bus = bus
        self.reflector = reflector
        self.evolver = evolver
        self.skill_registry = skill_registry
        self.role_assignor = role_assignor
        self.evolution_every_n_tasks = evolution_every_n_tasks
        self._reflection_history: dict[str, list[ReflectionRecord]] = defaultdict(list)
        self._evolved_prompts: dict[str, str] = {}

    async def start(self) -> None:
        """Subscribe to events and ensure bus is running."""
        self.bus.subscribe(TaskCompleteEvent, self._on_task_complete)
        self.bus.subscribe(SessionEndEvent, self._on_session_end)
        await self.bus.start()

    async def _on_task_complete(self, event: TaskCompleteEvent) -> None:
        """Schedule reflection for completed tasks and update performance."""
        result = event.result if isinstance(event.result, dict) else {"value": event.result}
        if "quality_score" in result:
            self.role_assignor.update_performance(
                event.agent_id,
                float(result.get("quality_score", 0.5)),
                float(result.get("token_efficiency", 1.0)),
            )
        asyncio.create_task(
            self._run_reflection(
                agent_id=event.agent_id,
                node_id=event.node_id,
                task_description=str(result.get("task_description", "")),
                result=str(result),
                tokens_used=event.tokens_used,
            )
        )

    async def _run_reflection(self, agent_id: str, node_id: str, task_description: str, result: str, tokens_used: int) -> None:
        """Run reflection and evolve prompt once threshold is reached."""
        rec = await self.reflector.reflect(
            agent_id=agent_id,
            node_id=node_id,
            task_description=task_description,
            task_result=result,
            tokens_used=tokens_used,
            errors_encountered=[],
        )
        self._reflection_history[agent_id].append(rec)

        history = self._reflection_history[agent_id]
        if len(history) >= self.evolution_every_n_tasks:
            evolved = await self.evolver.evolve(agent_id, "", history)
            self._evolved_prompts[agent_id] = evolved

        await self.bus.publish_and_wait(
            ReflectionCompleteEvent(
                source_id=agent_id,
                agent_id=agent_id,
                reflection=rec.model_dump(),
                quality_score=rec.quality_score,
            )
        )

    async def _on_session_end(self, event: SessionEndEvent) -> None:
        """Finalize evolutions and persist high-quality skills at session end."""
        for agent_id, history in self._reflection_history.items():
            if len(history) >= self.evolution_every_n_tasks:
                evolved = await self.evolver.evolve(agent_id, self._evolved_prompts.get(agent_id, ""), history)
                self._evolved_prompts[agent_id] = evolved

            for rec in history:
                await self.skill_registry.store(
                    agent_id=agent_id,
                    task_type="generic",
                    description=rec.task_description,
                    execution_trace="\n".join(rec.what_went_well + rec.suggestions),
                    quality_score=rec.quality_score,
                    tags=["reflection-derived"],
                )

    async def get_evolved_prompt(self, agent_id: str, base_prompt: str) -> str:
        """Return evolved prompt for agent if available."""
        return self._evolved_prompts.get(agent_id, base_prompt)
