"""Agent — persistent conversational entity (memory, persona, skills).

Wraps AgentDefinition (config/YAML) and adds runtime behavior: memory,
system prompt building, skill composition, turn execution via PathwayVM.

Architecture:
    Agent = configured stateful pathway (config + memory + skills)
    Pathway = execution graph template

Usage:
    from core.agent import Agent, AgentRepository

    repo = AgentRepository(state_store=store)
    await repo.register(agent_definition, skills_registry=skills)
    agent = await repo.get("host")
    result = await agent.turn(
        "Hello", "thread-1", pathway_vm=vm, pathway_service=svc
    )
"""

import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from core.types.execution import ExecutionResult

from pydantic import BaseModel, ConfigDict, Field, computed_field

from core.pathways.service import PathwayService
from core.protocols.services import PathwayVMLike, StateStoreLike
from core.skills.models import Skill
from core.types import AgentDefinition, ThreadState
from core.types.execution import ExecutionBudget

logger = logging.getLogger(__name__)


# =============================================================================
# HELPERS
# =============================================================================


_RESPONSE_KEYS = ("response", "content", "text", "message")
"""Keys tried when extracting text from a node output dict."""


def _extract_agent_response(outputs: dict[str, Any]) -> str:
    """Extract clean response text from pathway outputs.

    Walks through possible response locations in priority order:

    1. Top-level ``response`` / ``content`` key (promoted from OutputNode).
    2. ``output`` node result's ``response`` / ``content`` field.
    3. Any node's ``output.response`` or ``output.content`` (LLM output).
    4. Last resort: stringify the outputs dict.

    Parameters
    ----------
    outputs
        The ``PathwayRecord.outputs`` dict from the DAG scheduler.
    """
    if not outputs:
        return ""

    # 1. Top-level keys — set when OutputNode promotes its mapping.
    for key in _RESPONSE_KEYS:
        r = outputs.get(key)
        if isinstance(r, str) and r.strip():
            return r

    # 2. OutputNode result stored under "output" node key.
    out = outputs.get("output")
    if isinstance(out, dict):
        for key in _RESPONSE_KEYS:
            r = out.get(key)
            if isinstance(r, str) and r.strip():
                return r

    # 3. Walk node results looking for LLM output.response / content.
    #    Iterate in reverse insertion order so the *last* LLM node wins.
    for _node_id in reversed(list(outputs)):
        node_data = outputs[_node_id]
        if not isinstance(node_data, dict):
            continue
        llm_out = node_data.get("output")
        if not isinstance(llm_out, dict):
            continue
        for key in _RESPONSE_KEYS:
            r = llm_out.get(key)
            if isinstance(r, str) and r.strip():
                return r

    # 4. Last resort.
    return str(outputs)


def _extract_skill_attrs(skill: Skill | dict[str, Any], *attrs: str) -> tuple[Any, ...]:
    """Extract attributes from a skill object or dict.

    Works with both Skill objects (hasattr) and dicts (get).
    Returns a tuple of values in the same order as attrs.
    """
    values = []
    for attr in attrs:
        if hasattr(skill, attr):
            values.append(getattr(skill, attr, ""))
        elif isinstance(skill, dict):
            values.append(skill.get(attr, ""))
        else:
            values.append("")
    return tuple(values) if len(values) > 1 else (values[0],)


def _tool_patterns_to_skill_ids(
    tool_patterns: list[str],
) -> tuple[set[str], bool]:
    """Derive skill IDs from tool patterns.

    Examples
    --------
    - ["web.*", "memory.search"] -> ({"web", "memory"}, False)
    - ["*"] -> (set(), True)
    """
    include_all = False
    skill_ids: set[str] = set()
    for pattern in tool_patterns or []:
        if pattern == "*":
            include_all = True
            break
        if "." in pattern:
            skill_ids.add(pattern.split(".", 1)[0])
        else:
            skill_ids.add(pattern)
    return skill_ids, include_all


# =============================================================================
# AGENT
# =============================================================================


class Agent(BaseModel):
    """Persistent conversational entity — owns memory, persona, skills.

    A Pydantic model that wraps an ``AgentDefinition`` (the config spec)
    and adds runtime state: memory management, system prompt building,
    skill composition, and turn execution via PathwayVM delegation.

    Attributes
    ----------
    definition
        The AgentDefinition config (from YAML or builder).
    state_store
        State storage backend for thread persistence.
    skills_registry
        Full registry of skills (for skill scoping).
    agent_skills
        Explicit skill IDs assigned to this agent at config level.
    scoped_skills
        Pre-scoped list of skills (overrides registry-based scoping).
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    definition: AgentDefinition
    state_store: Any = Field(exclude=True)  # StateStoreLike protocol
    skills_registry: dict[str, Any] = Field(default_factory=dict, exclude=True)
    agent_skills: list[str] = Field(default_factory=list)
    scoped_skills: list[Any] | None = Field(default=None, exclude=True)

    # -------------------------------------------------------------------------
    # Identity (computed from definition)
    # -------------------------------------------------------------------------

    @computed_field  # type: ignore[prop-decorator]
    @property
    def id(self) -> str:
        """Agent identifier."""
        return self.definition.id

    @computed_field  # type: ignore[prop-decorator]
    @property
    def name(self) -> str:
        """Human-readable name."""
        return self.definition.display_name

    # -------------------------------------------------------------------------
    # Memory (agent owns this)
    # -------------------------------------------------------------------------

    async def load_thread(self, thread_id: str) -> ThreadState:
        """Load conversation thread, creating new if absent."""
        return await self.state_store.load_thread(self.id, thread_id)

    async def save_thread(self, state: ThreadState) -> None:
        """Persist thread state."""
        await self.state_store.save_thread(state)

    def prepare_history(
        self,
        state: ThreadState,
        max_messages: int,
    ) -> list[dict[str, str]]:
        """Prepare conversation history with truncation.

        Converts ThreadState messages to LLM-compatible dicts,
        excluding the latest user message (sent separately as goal).

        Parameters
        ----------
        state
            Thread state containing conversation messages.
        max_messages
            Maximum number of history messages to include.

        Returns
        -------
        list[dict[str, str]]
            LLM-compatible message dicts with ``role`` and ``content``.
        """
        history_messages = state.messages[:-1] if len(state.messages) > 1 else []
        if len(history_messages) > max_messages:
            history_messages = history_messages[-max_messages:]
        return [msg.to_llm_message() for msg in history_messages]

    # -------------------------------------------------------------------------
    # Persona (agent builds this)
    # -------------------------------------------------------------------------

    def scope_skills(self) -> tuple[list[Skill], dict[str, Skill]]:
        """Scope global skills to what this agent can use.

        Returns
        -------
        tuple[list[Skill], dict[str, Skill]]
            Scoped skills list and registry dict.
        """
        if self.scoped_skills is not None:
            scoped_registry = {
                getattr(s, "id", ""): s for s in self.scoped_skills if getattr(s, "id", "")
            }
            return list(self.scoped_skills), scoped_registry

        if not self.skills_registry:
            return [], {}

        assigned_ids: set[str] = set(self.agent_skills)
        assigned_ids.update(self.definition.skills or [])

        implied_ids, include_all = _tool_patterns_to_skill_ids(self.definition.tools)
        if include_all:
            all_skills = list(self.skills_registry.values())
            return all_skills, dict(self.skills_registry)

        allowed_ids = assigned_ids.union(implied_ids)
        if not allowed_ids:
            return [], {}

        scoped_registry = {sid: s for sid, s in self.skills_registry.items() if sid in allowed_ids}
        scoped_list = list(scoped_registry.values())
        return scoped_list, scoped_registry

    def build_system_prompt(self, state: ThreadState) -> str:
        """Build the system prompt from persona + skills + memory.

        Parameters
        ----------
        state
            Current thread state (for working memory and metadata).

        Returns
        -------
        str
            The assembled system prompt.
        """
        parts: list[str] = []
        defn = self.definition

        if defn.instructions:
            parts.append(f"## Instructions\n{defn.instructions}")

        if state.working_memory:
            wm_lines = [f"- **{key}**: {value}" for key, value in state.working_memory.items()]
            if wm_lines:
                parts.append(f"## Current Context\n{'\n'.join(wm_lines)}")

        # Scoped skills overview
        scoped_skills, scoped_registry = self.scope_skills()
        if scoped_skills:
            skill_lines = []
            for skill in scoped_skills:
                s_name, sid, desc = _extract_skill_attrs(skill, "name", "id", "description")
                if sid and desc:
                    skill_lines.append(f"- **{s_name}** (`{sid}.*`): {desc}")
            if skill_lines:
                parts.append(
                    "## Your Skills\n"
                    "You have these capabilities. Each skill provides "
                    "tools as `{skill}.{tool}`.\n"
                    "Use `skill.list` for details, "
                    "`skill.activate(name)` to load instructions.\n\n"
                    f"{'\n'.join(skill_lines)}",
                )

        # Activated skills (full instructions loaded during conversation)
        activated_skills = state.metadata.get("activated_skills", []) if state.metadata else []
        for skill_data in activated_skills:
            skill_name, instructions = _extract_skill_attrs(skill_data, "name", "instructions")
            if not skill_name:
                skill_name, _ = _extract_skill_attrs(skill_data, "id")
                skill_name = skill_name or "Skill"
            if instructions:
                parts.append(f"## Active Skill: {skill_name}\n{instructions}")

        # Agent-level skills with full knowledge injection
        if self.agent_skills:
            for skill_id in self.agent_skills:
                skill = scoped_registry.get(skill_id) or self.skills_registry.get(skill_id)
                if skill:
                    skill_name, instructions = _extract_skill_attrs(skill, "name", "instructions")
                    if instructions:
                        parts.append(
                            f"## Skill Knowledge: "
                            f"{skill_name or skill_id}\n\n"
                            f"{instructions}\n\n"
                            "**IMPORTANT**: Use the information above "
                            "to answer questions. This is authoritative "
                            "knowledge about AlbusOS.",
                        )

        # Active delegations
        delegated_tasks = state.metadata.get("delegated_tasks", []) if state.metadata else []
        active_tasks = [
            t for t in delegated_tasks if isinstance(t, dict) and t.get("status") == "in_progress"
        ]
        if active_tasks:
            tasks_text = "\n".join(
                f"- [{t.get('agent_id', '?')}] {t.get('description', '')}" for t in active_tasks
            )
            parts.append(f"## Active Delegations\n{tasks_text}")

        parts.append(
            "## Operating Rules\n"
            "- Use only the tools you have access to\n"
            "- If you need detailed guidance for a skill, "
            "use `skill.activate(name)`\n"
            "- Follow the loop's required response format",
        )

        return "\n\n".join(parts)

    # -------------------------------------------------------------------------
    # Turn execution (agent delegates to pathway)
    # -------------------------------------------------------------------------

    async def turn(
        self,
        message: str,
        thread_id: str,
        *,
        pathway_vm: PathwayVMLike,
        pathway_service: PathwayService,
        budget: ExecutionBudget | None = None,
        pathway_override: str | None = None,
        stream_callback: Any = None,
    ) -> "ExecutionResult":
        """Execute one conversational turn.

        Flow:
        1. Load thread (agent owns memory)
        2. Add user message
        3. Build system prompt (agent owns persona)
        4. Prepare conversation history
        5. Load pathway
        6. Create RunState with execution defaults
        7. Execute via PathwayVM
        8. Add assistant response to memory
        9. Save thread

        Parameters
        ----------
        message
            User's message.
        thread_id
            Conversation thread identifier.
        pathway_vm
            Pathway execution engine.
        pathway_service
            Pathway loading service.
        budget
            Execution budget (tokens, cost, duration).
        pathway_override
            Override the agent's default pathway for this turn.
        stream_callback
            Optional callback invoked with ``StreamEvent`` objects as
            LLM tokens arrive and tools execute.

        Returns
        -------
        ExecutionResult
            The turn result with response, metrics, and status.
        """
        from core.pathways.run_state import RunState
        from core.types.execution import ExecutionBudget as _Budget
        from core.types.execution import ExecutionResult

        budget = budget or _Budget()

        # 1. Load thread (agent owns memory)
        state = await self.load_thread(thread_id)
        state.add_message("user", message)
        state.turn_count += 1

        # 2. Build system prompt (agent owns persona)
        system = self.build_system_prompt(state)

        # 3. Prepare conversation history
        conversation_history = self.prepare_history(state, budget.max_context_messages)

        # 4. Load pathway
        pathway_id = pathway_override or self.definition.pathway
        pathway = pathway_service.load(pathway_id)
        if pathway is None:
            await self.save_thread(state)
            return ExecutionResult(
                success=False,
                response=f"Pathway not found: {pathway_id}. "
                "Load pathways from directories in config.",
                agent_id=self.id,
                thread_id=thread_id,
                pathway=pathway_id,
            )

        # 5. Create RunState with execution defaults from this agent
        defn = self.definition
        run_state = RunState(
            agent_id=self.id,
            thread_id=thread_id,
            state=state,
            system_prompt=system,
            default_model=defn.model,
            default_tools=list(defn.tools),
            default_temperature=defn.temperature,
            max_steps=defn.max_steps,
            model_routing=defn.model_routing,
        )
        if stream_callback is not None:
            run_state.set_stream_callback(stream_callback)

        # 6. Execute via pathway VM
        try:
            record = await pathway_vm.execute(
                pathway,
                inputs={
                    "goal": (state.messages[-1].content if state.messages else ""),
                    "conversation_history": conversation_history,
                    "budget_max_tokens": budget.max_tokens,
                    "budget_max_cost": budget.max_cost_usd,
                },
                timeout=float(budget.max_duration_seconds),
                run_state=run_state,
            )
        except Exception:
            await self.save_thread(state)
            raise

        # 7. Extract response
        response = _extract_agent_response(record.outputs)

        tokens_used = record.metadata.get("total_tokens", 0)
        if not tokens_used:
            diag = record.metadata.get("_diagnostics", {})
            tokens_used = diag.get("total_tokens", 0)

        # 8. Update memory
        state.add_message("assistant", response)
        state.total_tokens_used += tokens_used
        await self.save_thread(state)

        return ExecutionResult(
            success=record.success,
            response=response,
            agent_id=self.id,
            thread_id=thread_id,
            steps_taken=record.metadata.get("_diagnostics", {}).get("nodes_executed", 0),
            tokens_used=tokens_used,
            tool_calls=record.metadata.get("tool_calls", []),
            pathway=pathway_id,
        )


# =============================================================================
# AGENT REPOSITORY
# =============================================================================


class AgentRepository:
    """Registry of Agent instances.

    Manages the lifecycle of agents: registration, lookup, listing.
    Each agent is constructed with a shared state store for memory.

    Parameters
    ----------
    state_store
        Shared state store used by all agents for thread persistence.
    """

    def __init__(
        self,
        state_store: StateStoreLike,
    ) -> None:
        self._state_store = state_store
        self._agents: dict[str, Agent] = {}

    async def register(
        self,
        definition: AgentDefinition,
        *,
        skills_registry: dict[str, Skill] | None = None,
        agent_skills: list[str] | None = None,
    ) -> Agent:
        """Register an agent from its definition.

        Parameters
        ----------
        definition
            Agent config spec.
        skills_registry
            Global skills registry for skill scoping.
        agent_skills
            Explicit skill IDs assigned at config level.

        Returns
        -------
        Agent
            The created and registered agent.
        """
        agent = Agent(
            definition=definition,
            state_store=self._state_store,
            skills_registry=skills_registry or {},
            agent_skills=agent_skills or [],
        )
        self._agents[definition.id] = agent
        logger.info("Registered agent: %s (%s)", definition.id, agent.name)
        return agent

    def register_host(self, host: AgentDefinition) -> Agent:
        """Register the host agent (synchronous convenience for boot).

        Parameters
        ----------
        host
            The host agent definition.
        """
        if host.id != "host":
            logger.warning("Host agent should have id='host', got '%s'", host.id)
        agent = Agent(
            definition=host,
            state_store=self._state_store,
        )
        self._agents["host"] = agent
        logger.info("Registered Host agent: %s", agent.name)
        return agent

    async def get(self, agent_id: str) -> Agent | None:
        """Get an agent by ID."""
        return self._agents.get(agent_id)

    async def save(self, definition: AgentDefinition) -> Agent:
        """Save definition: create Agent and register it.

        Parameters
        ----------
        definition
            Agent config spec.
        """
        agent = Agent(
            definition=definition,
            state_store=self._state_store,
        )
        self._agents[definition.id] = agent
        logger.info("Saved agent: %s", definition.id)
        return agent

    async def list(self) -> list[Agent]:
        """List all registered agents."""
        return list(self._agents.values())

    async def exists(self, agent_id: str) -> bool:
        """Check if an agent exists."""
        return agent_id in self._agents


# =============================================================================
# GLOBAL SINGLETON
# =============================================================================

_global_agent_repo: AgentRepository | None = None


def get_agent_repository() -> AgentRepository:
    """Get the global agent repository.

    Lazily creates a default instance backed by :class:`FileStateStore`
    if none has been installed via :func:`set_agent_repository`.
    """
    global _global_agent_repo
    if _global_agent_repo is None:
        from core.state import FileStateStore

        _global_agent_repo = AgentRepository(state_store=FileStateStore())
    return _global_agent_repo


def set_agent_repository(repo: AgentRepository) -> None:
    """Install *repo* as the global agent repository.

    Called at bootstrap or by consumer code that wants to provide a
    custom ``StateStore``.
    """
    global _global_agent_repo
    _global_agent_repo = repo


__all__ = [
    "Agent",
    "AgentRepository",
    "get_agent_repository",
    "set_agent_repository",
]
