"""Agent — persistent conversational entity that owns memory, persona, and skills.

An Agent wraps an AgentDefinition (the config/YAML spec) and adds runtime
behavior: memory management, system prompt building, skill composition,
and turn execution via delegation to PathwayVM.

Architecture:
    Agent = WHO (identity + memory + persona)
    Pathway = HOW (execution logic)
    Runtime = INFRASTRUCTURE (lifecycle + plumbing)

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)
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from core.types import AgentDefinition, ConversationMessage, ThreadState

if TYPE_CHECKING:
    from core.pathways.service import PathwayService
    from core.protocols.services import PathwayVMLike, StateStoreLike
    from core.skills.models import Skill
    from core.types.execution import ExecutionBudget
    from persistence.store import StudioStore
    from server.runtime import ExecutionResult

logger = logging.getLogger(__name__)


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


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:
    """Persistent conversational entity — owns memory, persona, skills.

    Three responsibilities:
    1. Memory — load/save threads, prepare conversation history
    2. Persona — build system prompt from instructions + skills + context
    3. Turn execution — delegate to PathwayVM with the right context

    Parameters
    ----------
    definition
        The AgentDefinition config (from YAML or builder).
    state_store
        State storage backend for thread persistence.
    skills
        Scoped list of skills this agent can use.
    skills_registry
        Full registry of skills (for skill scoping).
    agent_skills
        Explicit skill IDs assigned to this agent at config level.
    """

    def __init__(
        self,
        definition: AgentDefinition,
        *,
        state_store: StateStoreLike,
        skills: list[Skill] | None = None,
        skills_registry: dict[str, Skill] | None = None,
        agent_skills: list[str] | None = None,
    ) -> None:
        self._definition = definition
        self._state_store = state_store
        self._skills_registry: dict[str, Skill] = skills_registry or {}
        self._agent_skills: list[str] = agent_skills or []
        self._scoped_skills: list[Skill] | None = skills

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

    @property
    def id(self) -> str:
        """Agent identifier."""
        return self._definition.id

    @property
    def name(self) -> str:
        """Human-readable name."""
        return self._definition.display_name

    @property
    def definition(self) -> AgentDefinition:
        """The underlying config spec (for serialization and introspection)."""
        return self._definition

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

    async def load_thread(self, thread_id: str) -> ThreadState:
        """Load conversation thread, creating a new one if it doesn't exist."""
        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 (which is sent separately as the 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 context.

        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
        agent_skill_ids = self._agent_skills
        if agent_skill_ids:
            for skill_id in agent_skill_ids:
                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: {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,
        attachments: list[object] | None = 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.
        attachments
            Optional message attachments.

        Returns
        -------
        ExecutionResult
            The turn result with response, metrics, and status.
        """
        from core.pathways.run_state import RunState  # noqa: PLC0415 (circular)
        from core.types.execution import ExecutionBudget as _EB  # noqa: PLC0415 (circular)
        from server.runtime import ExecutionResult  # noqa: PLC0415 (circular)

        budget = budget or _EB()

        # 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,
        )

        # 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 = ""
        if record.outputs:
            response = (
                record.outputs.get("response", "")
                or record.outputs.get("output", {}).get("response", "")
                or str(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, and listing.
    Each agent is constructed with a shared state store for memory persistence.

    Parameters
    ----------
    state_store
        The shared state store used by all agents for thread persistence.
    store
        Optional StudioStore for loading persisted agent definitions.
    """

    def __init__(
        self,
        state_store: StateStoreLike,
        *,
        store: StudioStore | None = None,
    ) -> None:
        self._state_store = state_store
        self._store = 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,
            state_store=self._state_store,
            skills_registry=skills_registry,
            agent_skills=agent_skills,
        )
        self._agents[definition.id] = agent
        logger.info("Registered agent: %s (%s)", definition.id, agent.name)
        return agent

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

        Parameters
        ----------
        host
            The host agent definition.
        **kwargs
            Passed to Agent constructor (skills_registry, etc.).
        """
        if host.id != "host":
            logger.warning("Host agent should have id='host', got '%s'", host.id)
        agent = Agent(
            host,
            state_store=self._state_store,
            **kwargs,
        )
        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, **kwargs: Any) -> Agent:
        """Save an agent definition — creates an Agent instance and persists to StudioStore.

        When a StudioStore is configured, the agent definition is persisted as a
        document with a revision. This ensures agents survive restarts and can
        be loaded by ``load_all()``.

        Parameters
        ----------
        definition
            Agent config spec.
        **kwargs
            Passed to Agent constructor (skills_registry, etc.).
        """
        agent = Agent(
            definition,
            state_store=self._state_store,
            **kwargs,
        )
        self._agents[definition.id] = agent

        # Persist to StudioStore when available
        if self._store:
            try:
                self._persist_definition(definition)
            except Exception as e:
                logger.warning("Failed to persist agent %s: %s", definition.id, e)

        logger.info("Saved agent: %s", definition.id)
        return agent

    def _persist_definition(self, definition: AgentDefinition) -> None:
        """Persist an agent definition to StudioStore as document + revision."""
        from persistence.models import DocType, Document  # noqa: PLC0415 (circular)

        if not self._store:
            return

        import uuid

        doc_id = f"agent:{definition.id}"
        rev_id = uuid.uuid4().hex[:12]

        # Upsert document record
        doc = Document(
            id=doc_id,
            type=DocType.AGENT,
            name=definition.display_name,
            workspace_id=definition.workspace_id,
            head_rev=rev_id,
        )
        self._store.upsert_document(doc)

        # Write revision with the full definition content
        self._store.write_revision(
            doc_id=doc_id,
            rev_id=rev_id,
            content=definition.model_dump(mode="json"),
        )

    async def list(self, *, workspace_id: str | None = None) -> 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

    async def load_all(self) -> None:
        """Load agents from persistence (StudioStore)."""
        if not self._store:
            logger.debug("No store configured, skipping load")
            return

        try:
            docs = self._store.list_documents()
            for doc in docs:
                if doc.type != "agent":
                    continue
                head_rev = doc.head_rev
                if not head_rev:
                    continue
                rev = self._store.get_revision(doc_id=doc.id, rev_id=head_rev)
                if not rev or not rev.content:
                    continue
                try:
                    defn = AgentDefinition.model_validate(rev.content)
                    if defn.id != "host":
                        await self.save(defn)
                        logger.debug("Loaded agent: %s", defn.id)
                except Exception as e:
                    logger.warning("Failed to load agent %s: %s", doc.id, e)
            logger.info("Loaded %d agents from persistence", len(self._agents))
        except Exception as e:
            logger.exception("Failed to load agents: %s", e)


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