"""AgentTemplate — a starting point for building custom Agents.

Copy this file, rename the class, and add your custom logic.
The main customization points are the pre/post hooks.

This template extends SimpleAgent so you get out-of-the-box:
- LLM initialization (self.llm)
- Tool service (self.tool_service)
- Tool/handoff/agent/skill resolution (self._prepare)
- Monitoring (self._collector)
- System and user prompt templates

Customize forward() to control the full execution flow, or override
the individual hooks (pre_forward / post_forward) for lighter changes.
"""

from __future__ import annotations

from typing import Any

from fluxibly.agent.base import AgentConfig, AgentResponse
from fluxibly.agent.simple_agent import SimpleAgent


class AgentTemplate(SimpleAgent):
    """Custom agent template — rename and modify to suit your use case.

    Inherits all SimpleAgent infrastructure (LLM, tools, monitoring,
    handoffs, skills). Override forward() for full control, or just
    override pre_forward() / post_forward() for lighter customization.
    """

    def __init__(self, config: AgentConfig) -> None:
        super().__init__(config)
        # Add any custom initialization here
        # e.g. self.my_state = {}

    # ══════════════════════════════════════════════════════════════
    # Hooks — override these for light customization
    # ══════════════════════════════════════════════════════════════

    async def pre_forward(
        self,
        messages: list[dict[str, Any]],
        context: dict[str, Any],
    ) -> tuple[list[dict[str, Any]], dict[str, Any]]:
        """Called before the main pipeline. Modify messages/context as needed.

        Args:
            messages: Conversation history (mutable copy is fine).
            context: Runtime context dict.

        Returns:
            (messages, context) — potentially modified.
        """
        return messages, context

    async def post_forward(
        self,
        response: AgentResponse,
        messages: list[dict[str, Any]],
        context: dict[str, Any],
    ) -> AgentResponse:
        """Called after the main pipeline. Modify or inspect the response.

        Args:
            response: The AgentResponse produced by the pipeline.
            messages: Original conversation history.
            context: Runtime context dict.

        Returns:
            AgentResponse — potentially modified.
        """
        return response

    # ══════════════════════════════════════════════════════════════
    # forward() — full pipeline with pre/post hooks
    # ══════════════════════════════════════════════════════════════

    async def forward(
        self,
        messages: list[dict[str, Any]],
        context: dict[str, Any] | None = None,
    ) -> AgentResponse:
        context = context or {}

        # ── Hook: pre-processing ──────────────────────────────────
        parsed_messages, parsed_context = await self.pre_forward(
            messages, context
        )

        # ── Delegate to SimpleAgent.forward() ─────────────────────
        response = await super().forward(parsed_messages, parsed_context)

        # ── Hook: post-processing ─────────────────────────────────
        return await self.post_forward(response, messages, context)
