"""DeepAgent — self-learning agent execution for netanel-core.

Wraps the deepagents library with our hook-based learning system.
deepagents handles: orchestration, ReAct loop, tools, planning, summarization.
Our layer adds: hook-based learning, knowledge preservation, safety bounds.

Flow: CREATE -> INVOKE -> (middleware fires hooks per step) -> RETURN
NOTE: on_complete is NOT fired inside spawn(). LearningLLM._agent_path() fires it
after evaluation passes — ensuring only quality-verified output reaches the hook.

Architecture ref: ARCHITECTURE.md Sections 4 (Component 7), 6 (Flow), 13 (Step 7).
Memory decision: deepagents = StateBackend only (ephemeral). MemoryStore = persistent.
"""

from __future__ import annotations

import contextlib
import json
import logging
from collections.abc import Callable, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable

if TYPE_CHECKING:
    # BUG-064: BaseTool is used only in type annotations. Guarding it behind
    # TYPE_CHECKING prevents langchain_core from loading 200+ modules at import
    # time. With `from __future__ import annotations`, all annotations are
    # evaluated lazily as strings — the runtime never needs the actual class.
    from langchain_core.tools import BaseTool

_LANGCHAIN_MIDDLEWARE_AVAILABLE = False
with contextlib.suppress(ImportError):
    from langchain.agents.middleware.types import (  # type: ignore[import-not-found]
        AgentMiddleware,
        ModelRequest,
        ModelResponse,
    )

    _LANGCHAIN_MIDDLEWARE_AVAILABLE = True

if not _LANGCHAIN_MIDDLEWARE_AVAILABLE:
    # langchain.agents.middleware.types does not exist in all langchain versions.
    # Provide minimal stubs so `import netanel_core` never fails at import time.
    # NathanMiddleware will still fail at *runtime* when deepagents is not installed
    # (same as before), but that's acceptable — package import itself is always safe.

    class AgentMiddleware:  # type: ignore[no-redef]
        """Stub: replaces AgentMiddleware when the langchain path is unavailable."""

        tools: tuple = ()

        def before_agent(self, state: object, runtime: object) -> object:
            """Called before agent execution."""

        def after_model(self, state: object, runtime: object) -> object:
            """Called after each model call."""

        def after_agent(self, state: object, runtime: object) -> object:
            """Called after agent execution."""

        def wrap_model_call(self, request: object, handler: object) -> object:
            """Wraps model call — stub does nothing."""

    class ModelRequest:  # type: ignore[no-redef]
        """Stub: replaces ModelRequest when the langchain path is unavailable."""

        def __init__(self) -> None:
            self.messages: list = []

    class ModelResponse:  # type: ignore[no-redef]
        """Stub: replaces ModelResponse when the langchain path is unavailable."""

        def __init__(self) -> None:
            self.result: list = []


from netanel_core.config import NathanConfig  # noqa: E402
from netanel_core.hooks import HookManager  # noqa: E402

logger = logging.getLogger(__name__)


def _is_tool_validation_error(msg: Any) -> bool:
    """Return True if msg is a ToolMessage containing a validation error.

    Handles both LangChain message objects (with .type attribute) and
    plain dicts (with "type" key). Case-insensitive match on error patterns.
    """
    msg_type = getattr(msg, "type", None) or (msg.get("type") if isinstance(msg, dict) else None)
    if msg_type != "tool":
        return False
    content = getattr(msg, "content", None) or (
        msg.get("content", "") if isinstance(msg, dict) else ""
    )
    content_lower = str(content).lower()
    return "validationerror" in content_lower or "validation error" in content_lower


_DEEPAGENTS_MISSING_MSG = (
    "deepagents package not installed. Install with: pip install 'deepagents>=0.4.1,<1.0'"
)


# ────────────────────────────────────────────
# Protocols
# ────────────────────────────────────────────


@runtime_checkable
class AgentModel(Protocol):
    """Protocol for the underlying agent — enables mock injection for testing.

    Matches the CompiledStateGraph interface from LangGraph/deepagents.
    Real agents: created via create_deep_agent().
    Tests: inject a mock with configurable responses.
    """

    def invoke(self, input: dict, config: dict | None = None) -> dict:  # noqa: A002
        """Execute the agent synchronously."""
        ...


# ────────────────────────────────────────────
# Data Types
# ────────────────────────────────────────────


@dataclass
class SpawnResult:
    """Result of a DeepAgent.spawn() invocation.

    Returned to the caller with the agent output plus metadata about
    execution steps, compactions, and completion status.
    """

    result: str
    """The final agent output text."""

    steps_taken: int
    """Number of LLM calls made during execution."""

    compactions: int
    """Number of compaction cycles detected."""

    completed: bool = False
    """True = natural completion. False = max iterations or error."""

    metadata: dict = field(default_factory=dict)
    """Extra info: error, max_iterations_reached, etc."""


# ────────────────────────────────────────────
# NathanMiddleware
# ────────────────────────────────────────────


class NathanMiddleware(AgentMiddleware):
    """Maps our HookManager lifecycle to deepagents middleware hooks.

    Inherits from AgentMiddleware so deepagents wires it into the pipeline.

    Three HookManager hooks are mapped here:
    - post_step      -> after_model (fires after every LLM response)
    - pre_compaction  -> wrap_model_call (before handler when approaching limit)
    - post_compaction -> wrap_model_call (after handler when summarization detected)

    NOTE: on_complete is NOT fired inside spawn(). LearningLLM._agent_path()
    fires it after evaluation passes — so only quality-verified output reaches
    the hook. See BUG-031 fix.
    """

    tools: Sequence[BaseTool] = ()
    """No extra tools registered by this middleware (required by AgentMiddleware)."""

    COMPACTION_THRESHOLD: float = 0.8
    """Trigger pre_compaction when estimated tokens exceed 80% of max_tokens_per_task."""

    def __init__(self, hooks: HookManager, config: NathanConfig, task_prompt: str) -> None:
        self._hooks = hooks
        self._config = config
        self._task_prompt = task_prompt
        self._step_count = 0
        self._compaction_count = 0
        self._total_tokens = 0
        self._pre_compaction_fired = False
        self._last_state: dict[str, Any] | None = None

    def before_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None:
        """Reset state at agent start."""
        self._step_count = 0
        self._compaction_count = 0
        self._total_tokens = 0
        self._pre_compaction_fired = False
        self._last_state = None
        return None

    def after_model(self, state: Any, runtime: Any) -> dict[str, Any] | None:
        """Post-step hook — micro-store facts after each LLM response."""
        self._step_count += 1
        try:
            messages = state.get("messages", [])
            if not messages:
                return None
            last_message = messages[-1]
            if hasattr(last_message, "content"):
                step_output = str(last_message.content)
            elif isinstance(last_message, dict):
                step_output = str(last_message.get("content", last_message))
            else:
                step_output = str(last_message)
            self._total_tokens += self._estimate_tokens(step_output)
            self._hooks.post_step(
                task=self._task_prompt,
                step_output=step_output,
            )
        except Exception:
            logger.warning("after_model hook failed", exc_info=True)
        return None

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse:
        """Wraps each LLM call to detect and handle compaction.

        BEFORE call: if approaching context limit, fire pre_compaction
        to save decisions/artifacts/patterns to disk.

        AFTER call: if message count dropped (summarization happened),
        fire post_compaction to re-inject memories.

        Args:
            request: ModelRequest with messages and context.
            handler: Callable that executes the actual LLM call.

        Returns:
            ModelResponse from the handler.
        """
        msg_count_before = len(request.messages)

        # PRE-COMPACTION: approaching limit? Save everything before it's lost
        if self._should_compact() and not self._pre_compaction_fired:
            try:
                context = "\n".join(
                    str(m.content) if hasattr(m, "content") else str(m) for m in request.messages
                )
                self._hooks.pre_compaction(
                    context=context,
                    decisions=[],
                    artifacts=[],
                )
                self._pre_compaction_fired = True
            except Exception:
                logger.warning("pre_compaction in wrap_model_call failed", exc_info=True)

        # Let deepagents + SummarizationMiddleware handle the call
        result = handler(request)

        # POST-COMPACTION: detect if summarization happened (message count dropped)
        # BUG-017 fix: result.result may be str/object, not list — use safe type check
        if hasattr(result, "messages") and isinstance(result.messages, list):
            msg_count_after = len(result.messages)
        elif hasattr(result, "result") and isinstance(result.result, list):
            msg_count_after = len(result.result)
        else:
            msg_count_after = msg_count_before  # Can't detect — skip compaction
        if msg_count_after < msg_count_before:
            try:
                self._hooks.post_compaction()
                self._compaction_count += 1
                self._pre_compaction_fired = False  # Reset for next cycle
            except Exception:
                logger.warning("post_compaction in wrap_model_call failed", exc_info=True)

        return result

    def after_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None:
        """Capture final state for result extraction in spawn()."""
        self._last_state = dict(state) if state else None
        return None

    def _should_compact(self) -> bool:
        """Check if we're approaching the context limit."""
        return (
            self._total_tokens > self.COMPACTION_THRESHOLD * self._config.safety.max_tokens_per_task
        )

    @staticmethod
    def _estimate_tokens(text: str) -> int:
        """Estimate token count from text. Uses chars/4 — same as ContextBuilder."""
        return len(text) // 4


# ────────────────────────────────────────────
# DeepAgent
# ────────────────────────────────────────────


class DeepAgent:
    """Wraps create_deep_agent() with self-learning hooks.

    deepagents handles: orchestration, ReAct loop, tools, planning, summarization.
    Our layer adds: hook-based learning, knowledge preservation, safety bounds.

    Usage:
        agent = DeepAgent(config, hooks)
        result = agent.spawn("Analyze this code", tools=[my_tool])
        print(result.result, result.steps_taken)

    Testing:
        mock_model = MockAgentModel(responses=[...])
        agent = DeepAgent(config, hooks, model=mock_model)
    """

    def __init__(
        self,
        config: NathanConfig,
        hooks: HookManager,
        model: AgentModel | None = None,
    ) -> None:
        """Initialize DeepAgent.

        Args:
            config: System configuration (safety bounds, model names).
            hooks: HookManager for lifecycle hooks.
            model: Optional agent model. None = create real agent via deepagents.
                   Explicit = use as-is (for testing with mocks).
        """
        self._config = config
        self._hooks = hooks
        self._model = model

    def spawn(
        self,
        prompt: str,
        tools: list | None = None,
        system_prompt: str | None = None,
        expect_json: bool = False,
    ) -> SpawnResult:
        """Execute a deep agent task with self-learning hooks.

        Args:
            prompt: The task description — used as user message and for hook logging.
            tools: Additional tools (added alongside deepagents built-ins).
                   Defaults to empty list.
            system_prompt: Optional full system prompt (role + memories). When
                           provided, the agent sees accumulated knowledge from
                           ContextBuilder. When None, falls back to prompt for
                           backward compatibility.
            expect_json: When True, the caller expects the final result to be a
                         JSON object and the BUG-040 tool-validation-error detector
                         is bypassed.  Use this when the task explicitly asks for
                         structured JSON output so legitimate responses are not
                         misidentified as failed tool calls.

        Returns:
            SpawnResult with execution details. Always returns — never raises.
        """
        tools = tools or []

        # BUG-012: use system_prompt when provided; fall back to prompt for backward compat.
        effective_system_prompt = system_prompt if system_prompt is not None else prompt

        # Create middleware for this spawn call (fresh state each time).
        # Always uses `prompt` (the task) for hook logging — not the full system prompt.
        middleware = NathanMiddleware(self._hooks, self._config, prompt)

        try:
            # Get agent (real or mock) — inside try so ImportError is caught
            # deepagents calls before_agent automatically when middleware is wired
            agent = (
                self._model
                if self._model is not None
                else self._create_agent(effective_system_prompt, tools, middleware)
            )

            # Invoke agent with safety bounds
            result_state = agent.invoke(
                {"messages": [{"role": "user", "content": prompt}]},
                config={"recursion_limit": self._config.safety.max_iterations},
            )

            # Extract result from final state
            result_str = self._extract_result(result_state)

            # BUG-025: Detect silent recursion limit (LangGraph issue #2165).
            # When the last message is a ToolMessage, the agent stopped mid-loop
            # without completing — completed=True would be wrong.
            messages = result_state.get("messages", [])
            if messages:
                last_msg = messages[-1]
                # Duck-typing: handles LangChain message objects and dicts equally
                msg_type = getattr(last_msg, "type", None) or (
                    last_msg.get("type") if isinstance(last_msg, dict) else None
                )
                if msg_type == "tool":
                    logger.warning(
                        "Silent recursion limit detected: last message is ToolMessage "
                        "(steps=%d, limit=%d). Agent stopped mid-tool-loop. "
                        "See LangGraph issue #2165.",
                        middleware._step_count,
                        self._config.safety.max_iterations,
                    )
                    return SpawnResult(
                        result=result_str,
                        steps_taken=middleware._step_count,
                        compactions=middleware._compaction_count,
                        completed=False,
                        metadata={
                            "error": "silent_recursion_limit",
                            "max_iterations_reached": True,
                        },
                    )

            # BUG-040: Detect tool validation error (deepagents issue #947).
            # When the agent encounters a tool validation error (wrong param name/type),
            # deepagents passes the error to the LLM which responds with corrected params
            # as raw JSON text. _extract_result() picks this up as a "completed" result —
            # but it's not. Detect by two factors: (1) result is pure JSON object,
            # (2) message history has at least one validation error ToolMessage.
            # Recovery exception: if a successful ToolMessage follows the last validation
            # error, the agent recovered and the JSON result may be legitimate.
            result_stripped = result_str.strip()
            if (
                not expect_json
                and result_stripped.startswith("{")
                and result_stripped.endswith("}")
            ):
                with contextlib.suppress(json.JSONDecodeError, ValueError):
                    json.loads(result_str)
                    # Find last validation error ToolMessage in history
                    last_validation_idx = next(
                        (
                            i
                            for i in range(len(messages) - 1, -1, -1)
                            if _is_tool_validation_error(messages[i])
                        ),
                        -1,
                    )
                    if last_validation_idx != -1:
                        # Recovery check: successful ToolMessage after last validation error
                        has_recovery = any(
                            not _is_tool_validation_error(m)
                            and (
                                getattr(m, "type", None) == "tool"
                                or (isinstance(m, dict) and m.get("type") == "tool")
                            )
                            for m in messages[last_validation_idx + 1 :]
                        )
                        if not has_recovery:
                            logger.warning(
                                "Tool validation error detected: agent returned raw JSON "
                                "params as text instead of completing task (steps=%d). "
                                "deepagents issue #947.",
                                middleware._step_count,
                            )
                            return SpawnResult(
                                result=result_str,
                                steps_taken=middleware._step_count,
                                compactions=middleware._compaction_count,
                                completed=False,
                                metadata={"error": "tool_validation_error"},
                            )

            # NOTE: on_complete is NOT fired here (BUG-031 fix).
            # LearningLLM._agent_path() fires it after evaluation passes —
            # ensuring only quality-verified output reaches the hook.

            return SpawnResult(
                result=result_str,
                steps_taken=middleware._step_count,
                compactions=middleware._compaction_count,
                completed=True,
                metadata={"max_iterations_reached": False},
            )

        except RecursionError:
            # GraphRecursionError is a subclass of RecursionError in LangGraph
            logger.warning("Agent hit recursion limit", exc_info=True)
            return SpawnResult(
                result="",
                steps_taken=middleware._step_count,
                compactions=middleware._compaction_count,
                completed=False,
                metadata={"error": "max_iterations", "max_iterations_reached": True},
            )

        except Exception as e:
            logger.warning("Agent execution failed", exc_info=True)
            return SpawnResult(
                result="",
                steps_taken=middleware._step_count,
                compactions=middleware._compaction_count,
                completed=False,
                metadata={"error": str(e), "max_iterations_reached": False},
            )

    def _create_agent(
        self,
        prompt: str,
        tools: Sequence[BaseTool | Callable | dict],
        middleware: NathanMiddleware,
    ) -> AgentModel:
        """Create a real deepagents agent (auto-wiring path).

        Lazy imports create_deep_agent so the function is only resolved at
        runtime, keeping mocks injectable for tests.

        Args:
            prompt: System prompt for the agent.
            tools: Tool objects to provide to the agent.
            middleware: NathanMiddleware instance wired into the pipeline.

        Returns:
            A CompiledStateGraph from deepagents that satisfies AgentModel.
        """
        try:
            from deepagents import create_deep_agent  # type: ignore[import-untyped]
        except ImportError as e:
            raise ImportError(_DEEPAGENTS_MISSING_MSG) from e

        # BUG-042 fix: guard against non-string prompt — deepagents may reformat
        # a non-string as [{"type":"text","text":...}] which OpenAI API rejects
        if not isinstance(prompt, str):
            raise TypeError(f"system_prompt must be a plain string, got {type(prompt).__name__}")

        return create_deep_agent(  # type: ignore[return-value]
            model=self._config.models.primary_model,
            tools=tools,
            system_prompt=prompt,
            middleware=[middleware],
            backend=self._create_backend(),
        )

    def _create_backend(self) -> Any:
        """Instantiate the configured deepagents backend.

        Returns the correct backend based on config.safety.backend_type:
        - "ephemeral"    → None (deepagents StateBackend default, in-memory only)
        - "filesystem"   → FilesystemBackend (real files, no shell execution)
        - "local_shell"  → LocalShellBackend (real files + bash + search, all three pillars)

        Lazy imports: backend classes resolved at runtime so mocks work in tests.

        Returns:
            Backend instance for deepagents, or None for StateBackend (ephemeral).
        """
        backend_type = self._config.safety.backend_type
        root_dir = self._config.safety.backend_root_dir

        if backend_type == "ephemeral":
            return None  # deepagents defaults to StateBackend when backend=None

        if backend_type == "filesystem":
            try:
                from deepagents.backends import FilesystemBackend  # type: ignore[import-untyped]
            except ImportError as e:
                raise ImportError(_DEEPAGENTS_MISSING_MSG) from e
            return FilesystemBackend(
                root_dir=root_dir,
                virtual_mode=False,
            )

        # "local_shell" — all three pillars: bash + files + search
        try:
            from deepagents.backends import LocalShellBackend  # type: ignore[import-untyped]
        except ImportError as e:
            raise ImportError(_DEEPAGENTS_MISSING_MSG) from e
        return LocalShellBackend(
            root_dir=root_dir,
            virtual_mode=False,
            timeout=self._config.safety.max_iterations * 10.0,  # Scale with safety bounds
            inherit_env=True,  # Allow access to project environment
        )

    @staticmethod
    def _extract_result(state: dict[str, Any]) -> str:
        """Extract the final result string from agent state.

        Args:
            state: The final state dict from agent.invoke().

        Returns:
            The content of the last message, or empty string.
        """
        messages = state.get("messages", [])
        if not messages:
            return ""
        last_message = messages[-1]
        if hasattr(last_message, "content"):
            return str(last_message.content)
        if isinstance(last_message, dict):
            return str(last_message.get("content", ""))
        return str(last_message)
