"""High-level Agent facade for the oedix framework.

Provides a single entry-point that wires up config, LLM, loaders,
tools, and the AgentQueue pipeline — so callers only need::

    import oedix

    agent = oedix.Agent(config="/path/to/agent.json", data_dir="/data")
    agent.set_event_handler(my_handler)

    async for event in agent.run("안녕하세요"):
        ...
"""

from __future__ import annotations

import os
from typing import Any, AsyncGenerator, Callable, Optional

from oedix.core.agent_queue import AgentQueue
from oedix.core.config import AgentConfig
from oedix.core.context import Context
from oedix.core.events import Event
from oedix.llm.client import LLMClient
from oedix.loader.action_loader import ActionLoader
from oedix.actions.thinking import ThinkingAction
from oedix.actions.response import ResponseAction
from oedix.tools.registry import ToolRegistry


class Agent:
    """High-level facade that builds and runs :class:`AgentQueue`.

    Args:
        config: Path to ``agent.json`` **or** an :class:`AgentConfig`
            instance.  When a path string is given it is loaded via
            :meth:`AgentConfig.from_file`.
        data_dir: Override ``data_dir`` in the config (optional).
        provider: Default LLM provider key (e.g. ``"openai"``).
        model: Default LLM model key (e.g. ``"gpt-5.4"``).
        instruction: System instruction text.  When *not* supplied the
            agent loads ``{data_dir}/instruction/README.md``
            automatically (if the file exists).
        **kwargs: Forwarded to :class:`AgentConfig` attribute overrides
            (``max_queue_count``, ``max_action_count``, …).
    """

    def __init__(
        self,
        config: str | AgentConfig | None = None,
        *,
        data_dir: str | None = None,
        provider: str | None = None,
        model: str | None = None,
        instruction: str | None = None,
        **kwargs: Any,
    ) -> None:
        # ── Config ────────────────────────────────────────────────
        if config is None:
            self._config = AgentConfig()
        elif isinstance(config, str):
            self._config = AgentConfig.from_file(config, data_dir=data_dir)
        else:
            self._config = config

        if data_dir and not isinstance(config, str):
            self._config.data_dir = data_dir
        if provider:
            self._config.default_provider = provider
        if model:
            self._config.default_model = model

        for key, value in kwargs.items():
            if hasattr(self._config, key):
                setattr(self._config, key, value)

        # ── Instruction ───────────────────────────────────────────
        if instruction is not None:
            self._instruction = instruction
        else:
            self._instruction = self._load_instruction()

        # ── Dependencies (lazy-init on first run) ─────────────────
        self._llm_client: LLMClient | None = None
        self._action_loader: ActionLoader | None = None
        self._tool_registry: ToolRegistry | None = None

        # ── Tool context (passed to ToolRegistry) ────────────────
        self._tool_context: dict[str, Any] = {}

        # ── Event handler ─────────────────────────────────────────
        self._event_handler: Callable[[dict], None] | None = None

        # ── Session state (client-managed) ────────────────────────
        self._context: list[dict] | None = None    # previous context (dict list)
        self._history: dict | None = None           # previous history
        self._events: list[dict] = []               # accumulated events from last run
        self._last_queue: AgentQueue | None = None   # reference to last executed queue

    # ==================================================================
    # Configuration helpers
    # ==================================================================

    @property
    def config(self) -> AgentConfig:
        return self._config

    def set_provider(self, provider: str) -> "Agent":
        """Change the default LLM provider."""
        self._config.default_provider = provider
        return self

    def set_model(self, model: str) -> "Agent":
        """Change the default LLM model."""
        self._config.default_model = model
        return self

    def set_instruction(self, instruction: str) -> "Agent":
        """Replace the system instruction text."""
        self._instruction = instruction
        return self

    def set_event_handler(self, handler: Callable[[dict], None]) -> "Agent":
        """Register a callback invoked for every event during ``run()``.

        The *handler* receives an event dict (same shape as
        :meth:`Event.to_dict`).  This is where you would implement
        ``print_event``-style rendering.
        """
        self._event_handler = handler
        return self

    def set(self, **kwargs: Any) -> "Agent":
        """Bulk-set config attributes.

        Example::

            agent.set(max_queue_count=10, max_action_count=5)
        """
        for key, value in kwargs.items():
            if hasattr(self._config, key):
                setattr(self._config, key, value)
        return self

    def set_tool_context(self, context: dict[str, Any]) -> "Agent":
        """Set the execution context passed to tools.

        This context dict is forwarded to
        :meth:`ToolRegistry.set_context` and made available to
        each tool instance via ``self.ctx``.
        """
        self._tool_context = dict(context)
        # If registry is already initialized, update it immediately
        if self._tool_registry is not None:
            self._tool_registry.set_context(self._tool_context)
        return self

    # ==================================================================
    # Session state — context get/set
    # ==================================================================

    def get_context(self) -> list[dict]:
        """Return the context from the last ``run()`` as a list of dicts.

        If no run has been executed yet, returns the previously set
        context (via :meth:`set_context`) or an empty list.
        """
        if self._last_queue is not None:
            return [c.to_dict() for c in self._last_queue.context]
        return list(self._context) if self._context else []

    def set_context(self, context: list[dict]) -> "Agent":
        """Set initial context for the next ``run()`` call.

        Accepts a list of dicts (as returned by :meth:`get_context`).
        The context will be restored from dicts via ``Context.from_dict``
        and merged into the new AgentQueue's context.
        """
        self._context = list(context)
        return self

    # ==================================================================
    # Session state — history get/set
    # ==================================================================

    def get_history(self) -> dict:
        """Return the full execution history from the last ``run()``.

        Returns a dict containing:
            - ``action_sets``: list of ActionSet dicts (with nested
              Action dicts including output/result).
            - ``events``: list of all Event dicts emitted during the run.
            - ``context``: final context state (list of Context dicts).
        """
        if self._last_queue is None:
            return self._history or {"action_sets": [], "events": [], "context": []}

        return {
            "action_sets": [s.to_dict() for s in self._last_queue.history],
            "events": list(self._events),
            "context": [c.to_dict() for c in self._last_queue.context],
        }

    def set_history(self, history: dict) -> "Agent":
        """Restore history from a previous session.

        Accepts a dict (as returned by :meth:`get_history`).
        The history is stored for reference; context within the
        history is also set via :meth:`set_context` if present.
        """
        self._history = dict(history)
        if "context" in history and history["context"]:
            self._context = list(history["context"])
        return self

    # ==================================================================
    # Session state — events
    # ==================================================================

    def get_events(self) -> list[dict]:
        """Return all events from the last ``run()`` as a list of dicts."""
        return list(self._events)

    # ==================================================================
    # Execution
    # ==================================================================

    async def run(self, query: str) -> AsyncGenerator[Event, None]:
        """Build an :class:`AgentQueue` and yield events.

        If an event handler was registered via :meth:`set_event_handler`,
        it is called for every event *before* yielding.

        After completion, the queue's context and history are retained
        for retrieval via :meth:`get_context` and :meth:`get_history`.
        """
        queue = self._build_queue(query)
        self._events = []
        self._last_queue = queue
        async for event in queue.run():
            event_dict = event.to_dict()
            self._events.append(event_dict)
            if self._event_handler is not None:
                self._event_handler(event_dict)
            yield event

    # ==================================================================
    # Internal wiring
    # ==================================================================

    def _ensure_llm(self) -> LLMClient:
        if self._llm_client is None:
            self._llm_client = LLMClient(self._config)
        return self._llm_client

    def _ensure_loader(self) -> ActionLoader | None:
        if self._action_loader is None:
            actions_dir = self._config.actions_dir()
            if os.path.isdir(actions_dir):
                self._action_loader = ActionLoader(actions_dir)
                self._action_loader.discover()
        return self._action_loader

    def _ensure_tools(self) -> ToolRegistry | None:
        if self._tool_registry is None:
            tools_dir = self._config.tools_dir()
            if os.path.isdir(tools_dir):
                self._tool_registry = ToolRegistry(tools_dir)
                if self._tool_context:
                    self._tool_registry.set_context(self._tool_context)
                self._tool_registry.scan()
        return self._tool_registry

    def _load_instruction(self) -> str:
        try:
            readme = os.path.join(self._config.instruction_dir(), "README.md")
            if os.path.isfile(readme):
                with open(readme, "r", encoding="utf-8") as f:
                    return f.read()
        except Exception:
            pass
        return ""

    def _build_queue(self, query: str) -> AgentQueue:
        """Create a fully wired AgentQueue for one query."""
        queue = AgentQueue(
            query=query,
            config=self._config,
            instruction=self._instruction,
        )

        # Merge previous context if set
        if self._context:
            prev_contexts = [Context.from_dict(c) for c in self._context]
            queue.context.extend(prev_contexts)
            self._context = None  # consumed

        queue._llm_client = self._ensure_llm()
        queue._action_loader = self._ensure_loader()
        queue._tool_registry = self._ensure_tools()
        queue._thinking_factory = lambda: ThinkingAction()
        queue._response_factory = lambda: ResponseAction()
        return queue
