"""
Deterministic loop hooks for PicoAgents.

Hooks provide deterministic control points around the LLM-controlled tool loop.
Unlike middleware (which intercepts individual operations like model_call, tool_call),
hooks operate at the LOOP level - before the first LLM call and when the loop
is about to exit.

Two hook points:
  - start hooks: Run before the first LLM call. Inject instructions.
  - end hooks: Run when agent would stop (no tool calls). Can resume the loop.

The key property is that hooks are DETERMINISTIC Python code - no LLM involved
in hook logic. The LLM only sees the result (an injected UserMessage).

Example:
    from picoagents import Agent
    from picoagents._hooks import PlanningHook, CompletionCheckHook
    from picoagents.tools import create_coding_tools, create_todo_tools

    agent = Agent(
        name="planner",
        instructions="You are a software engineer.",
        model_client=client,
        tools=[*create_coding_tools(), *create_todo_tools()],
        start_hooks=[PlanningHook()],
        end_hooks=[CompletionCheckHook(max_restarts=2)],
        max_iterations=20,
    )
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from .context import AgentContext
from .messages import Message


# =============================================================================
# Loop Context
# =============================================================================


@dataclass
class LoopContext:
    """Shared state passed to hooks during a single agent run.

    Provides hooks with access to the agent's context, current messages,
    and mutable metadata for inter-hook communication.

    Attributes:
        agent_context: The agent's working context (messages, metadata)
        llm_messages: Current message list being sent to the LLM
        agent_name: Name of the agent running
        iteration: Current loop iteration (0-indexed)
        restart_count: How many times end hooks have resumed the loop
        metadata: Mutable dict for hooks to share state
    """

    agent_context: AgentContext
    llm_messages: List[Message]
    agent_name: str
    iteration: int = 0
    restart_count: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)


# =============================================================================
# Termination Conditions
# =============================================================================


class TerminationCondition(ABC):
    """Controls when end hooks should stop restarting the agent loop.

    Composable with | (OR) and & (AND) operators:
        # Stop after 3 restarts OR when all todos are done
        condition = MaxRestartsTermination(3) | AllTodosComplete()
    """

    @abstractmethod
    def should_terminate(self, context: LoopContext) -> bool:
        """Return True to allow the agent to stop (no more restarts)."""
        ...

    def reset(self) -> None:
        """Reset state for a new agent run."""
        pass

    def __or__(self, other: "TerminationCondition") -> "CompositeTermination":
        """Combine with OR - terminate if EITHER condition is met."""
        return CompositeTermination([self, other], mode="any")

    def __and__(self, other: "TerminationCondition") -> "CompositeTermination":
        """Combine with AND - terminate only if BOTH conditions are met."""
        return CompositeTermination([self, other], mode="all")


class MaxRestartsTermination(TerminationCondition):
    """Terminate after a maximum number of loop restarts.

    Args:
        max_restarts: Maximum restarts before allowing stop. Default 2.

    Example:
        hook = CompletionCheckHook(termination=MaxRestartsTermination(3))
    """

    def __init__(self, max_restarts: int = 2):
        self.max_restarts = max_restarts

    def should_terminate(self, context: LoopContext) -> bool:
        return context.restart_count >= self.max_restarts

    def __repr__(self) -> str:
        return f"MaxRestartsTermination(max_restarts={self.max_restarts})"


class CompositeTermination(TerminationCondition):
    """Combines multiple termination conditions with AND/OR logic.

    Created automatically via | and & operators on TerminationCondition.

    Args:
        conditions: List of conditions to combine
        mode: "any" (OR - terminate if any is True) or "all" (AND - all must be True)
    """

    def __init__(
        self, conditions: List[TerminationCondition], mode: str = "any"
    ):
        if mode not in ("any", "all"):
            raise ValueError("Mode must be 'any' or 'all'")
        self.conditions = conditions
        self.mode = mode

    def should_terminate(self, context: LoopContext) -> bool:
        results = [c.should_terminate(context) for c in self.conditions]
        return any(results) if self.mode == "any" else all(results)

    def reset(self) -> None:
        for c in self.conditions:
            c.reset()

    def __or__(self, other: TerminationCondition) -> "CompositeTermination":
        if self.mode == "any":
            # Flatten: (A | B) | C = A | B | C
            if isinstance(other, CompositeTermination) and other.mode == "any":
                return CompositeTermination(
                    self.conditions + other.conditions, mode="any"
                )
            return CompositeTermination(self.conditions + [other], mode="any")
        # self is AND mode - treat as a single unit: (A & B) | C
        return CompositeTermination([self, other], mode="any")

    def __and__(self, other: TerminationCondition) -> "CompositeTermination":
        if self.mode == "all":
            # Flatten: (A & B) & C = A & B & C
            if isinstance(other, CompositeTermination) and other.mode == "all":
                return CompositeTermination(
                    self.conditions + other.conditions, mode="all"
                )
            return CompositeTermination(self.conditions + [other], mode="all")
        # self is OR mode - treat as a single unit: (A | B) & C
        return CompositeTermination([self, other], mode="all")

    def __repr__(self) -> str:
        op = " | " if self.mode == "any" else " & "
        return f"({op.join(repr(c) for c in self.conditions)})"


# =============================================================================
# Base Hook Classes
# =============================================================================


class BaseStartHook(ABC):
    """Base class for hooks that run before the first LLM call.

    Return a string to inject as a UserMessage, or None to do nothing.
    Multiple start hooks run in order - all injections are applied.

    Example:
        class MyStartHook(BaseStartHook):
            async def on_start(self, context):
                return "Always think step by step before answering."
    """

    @abstractmethod
    async def on_start(self, context: LoopContext) -> Optional[str]:
        """Called before the first LLM call.

        Args:
            context: Loop context with agent state

        Returns:
            Text to inject as UserMessage, or None for no injection.
        """
        ...


class BaseEndHook(ABC):
    """Base class for hooks that run when the agent would stop.

    Return a string to inject as a UserMessage and RESUME the loop,
    or None to allow the agent to stop.

    The first hook to return a non-None value wins - subsequent hooks
    are not called for that exit attempt.

    Example:
        class MyEndHook(BaseEndHook):
            async def on_end(self, context):
                if some_condition_not_met():
                    return "You're not done yet. Continue working."
                return None  # Allow stop
    """

    @abstractmethod
    async def on_end(self, context: LoopContext) -> Optional[str]:
        """Called when the agent would stop (no tool calls returned).

        Args:
            context: Loop context with agent state and restart_count

        Returns:
            Text to inject as UserMessage to resume loop, or None to stop.
        """
        ...


# =============================================================================
# Concrete Implementations
# =============================================================================


class PlanningHook(BaseStartHook):
    """Start hook that injects a planning instruction before the first LLM call.

    Instructs the agent to analyze the task and create a structured todo
    list using the todo_write tool before starting any work.

    Requires todo tools (todo_write, todo_read) in the agent's tool list.

    Args:
        instruction: Custom instruction text. If None, uses a sensible default.

    Example:
        agent = Agent(
            ...,
            tools=[*create_coding_tools(), *create_todo_tools()],
            start_hooks=[PlanningHook()],
        )
    """

    def __init__(self, instruction: Optional[str] = None):
        self.instruction = instruction or (
            "Before starting any work, you MUST:\n"
            "1. Analyze the task and break it into clear, actionable steps\n"
            "2. Use the todo_write tool to create a structured task list\n"
            "3. Each todo needs: content (what to do), status ('pending'), "
            "activeForm (present tense)\n"
            "4. Mark the first task as 'in_progress' before starting it\n"
            "5. Only ONE task should be 'in_progress' at a time\n\n"
            "As you work, update the todo list: mark tasks 'completed' "
            "when done, set the next task to 'in_progress'."
        )

    async def on_start(self, context: LoopContext) -> Optional[str]:
        return self.instruction

    def __repr__(self) -> str:
        return "PlanningHook()"


class CompletionCheckHook(BaseEndHook):
    """End hook that checks todo list completion before allowing the agent to stop.

    When the agent would stop (no tool calls), this hook:
    1. Loads the current todo list from file
    2. Checks if all todos are marked 'completed'
    3. If incomplete items remain AND termination allows: resumes the loop
    4. If all complete OR termination says stop: allows the agent to stop

    Args:
        termination: Termination condition. Defaults to MaxRestartsTermination(2).
        max_restarts: Shorthand for MaxRestartsTermination(N). Ignored if
            termination is provided.

    Example:
        agent = Agent(
            ...,
            end_hooks=[CompletionCheckHook(max_restarts=3)],
        )
    """

    def __init__(
        self,
        termination: Optional[TerminationCondition] = None,
        max_restarts: int = 2,
    ):
        if termination is not None:
            self.termination = termination
        else:
            self.termination = MaxRestartsTermination(max_restarts)

    async def on_end(self, context: LoopContext) -> Optional[str]:
        # Check termination condition first
        if self.termination.should_terminate(context):
            return None  # Allow stop

        # Load todos from file
        from .tools._context_tools import _load_todos

        todos = _load_todos()

        if not todos:
            return None  # No todos = nothing to check

        # Check for incomplete items
        incomplete = [t for t in todos if t.get("status") != "completed"]

        if not incomplete:
            return None  # All done, allow stop

        # Build resume message
        total = len(todos)
        completed = total - len(incomplete)

        items = "\n".join(
            f"  - [{t.get('status', 'pending')}] {t.get('content', '')}"
            for t in incomplete
        )

        return (
            f"You have {len(incomplete)} incomplete tasks "
            f"({completed}/{total} completed):\n{items}\n\n"
            f"Continue working on the next pending task. "
            f"Update the todo list as you make progress. "
            f"Do not ask for user input - proceed autonomously."
        )

    def __repr__(self) -> str:
        return f"CompletionCheckHook(termination={self.termination!r})"
