"""ActionSet — single-objective Action execution loop.

An ActionSet owns a queue of Actions targeting one ``subject`` (goal).
It always starts with a *thinking* Action and automatically re-inserts
*thinking* after every non-thinking Action so the agent can re-evaluate.

See design doc §3.2 for the full specification.
"""

from __future__ import annotations

import math
import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional

from oedix.core.action import Action, ActionStatus
from oedix.core.context import Context
from oedix.core.events import Event, EventType

if TYPE_CHECKING:
    from oedix.core.config import AgentConfig


class ActionSetStatus(str, Enum):
    """ActionSet lifecycle states."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class ActionSet:
    """Single-objective Action execution loop (§3.2).

    Attributes:
        id: Unique identifier (UUID).
        subject: Short goal description (set at creation time).
        approach: Approach text (set/appended by thinking Action).
        history: Completed Action list (append-only).
        queue: Pending Action list (FIFO).
        max_action_count: System limit from config.
        max_action_weight: Scaling factor (0.5–3.0), adjustable by
            the thinking Action via :class:`ActionSetHandle`.
        status: Current lifecycle state.
    """

    def __init__(
        self,
        subject: str,
        config: Optional[AgentConfig] = None,
        *,
        title: str = "",
        max_action_count: int = 15,
    ) -> None:
        self.id: str = uuid.uuid4().hex[:12]
        self.subject: str = subject
        self.title: str = title
        self.approach: str = ""
        self.history: list[Action] = []
        self.queue: list[Action] = []
        self.max_action_count: int = max_action_count
        self.max_action_weight: float = 1.0
        self.status: ActionSetStatus = ActionSetStatus.PENDING
        self.config = config
        self.created_at: datetime = datetime.now(timezone.utc)

        # Batch tracking: when thinking enqueues multiple actions,
        # auto-thinking insertion is deferred until the batch completes.
        self._batch_remaining: int = 0

        # Injected by AgentQueue before run()
        self._agent_queue: Any = None
        self._llm_client: Any = None
        self._thinking_factory: Optional[Callable[[], Action]] = None

    # ------------------------------------------------------------------
    # Effective limits
    # ------------------------------------------------------------------

    def effective_max(self) -> int:
        """Effective max actions = floor(max_action_count × weight)."""
        return math.floor(self.max_action_count * self.max_action_weight)

    def remaining_actions(self) -> int:
        """Number of Action slots still available."""
        return max(0, self.effective_max() - len(self.history))

    # ------------------------------------------------------------------
    # Queue management (exposed via ActionSetHandle)
    # ------------------------------------------------------------------

    def enqueue(self, action: Action) -> None:
        """Append an Action to the end of the queue."""
        self.queue.append(action)

    def enqueue_many(self, actions: list[Action]) -> None:
        """Append multiple Actions to the queue."""
        self.queue.extend(actions)

    def set_title(self, text: str) -> None:
        """Set the ActionSet title (brief goal, ≤20 chars)."""
        self.title = text[:20]

    def set_approach(self, text: str) -> None:
        """Append to the approach description."""
        if self.approach:
            self.approach += "\n" + text
        else:
            self.approach = text

    def set_weight(self, value: float) -> None:
        """Set ``max_action_weight``, clamped to [0.5, 3.0]."""
        self.max_action_weight = max(0.5, min(3.0, value))

    def get_history(self) -> list[Action]:
        """Return a shallow copy of completed Action history."""
        return list(self.history)

    # ------------------------------------------------------------------
    # Execution loop (§3.2 Pipeline)
    # ------------------------------------------------------------------

    async def run(
        self, context: list[Context]
    ) -> AsyncGenerator[Event, None]:
        """Execute the ActionSet loop.

        1. status → RUNNING
        2. Pop actions from queue, inject handles, run, append to history.
        3. After non-thinking actions, auto-insert thinking.
        4. status → COMPLETED (or FAILED on unhandled error).

        Args:
            context: The current AgentQueue context array.  Actions may
                append to it via their ``ActionOutput.context_entries``.

        Yields:
            :class:`Event` instances for SSE streaming.
        """
        self.status = ActionSetStatus.RUNNING

        yield self._make_event(
            EventType.ACTION_SET_START,
            {"subject": self.subject, "title": self.title, "approach": self.approach},
        )

        try:
            while self.queue and self.remaining_actions() > 0:
                action = self.queue.pop(0)

                # Inject parent references
                action._action_set = self
                action._agent_queue = self._agent_queue
                action._llm_client = self._llm_client
                action._action_loader = (
                    self._agent_queue._action_loader
                    if self._agent_queue
                    else None
                )

                # Emit ACTION_START before execution
                yield self._make_event(
                    EventType.ACTION_START,
                    {
                        "action_id": action.id,
                        "type": action.type,
                        "purpose": action.params.get("purpose", ""),
                    },
                )

                # Execute
                async for event in action.run(context):
                    yield event

                # Emit ACTION_COMPLETE after execution
                yield self._make_event(
                    EventType.ACTION_COMPLETE,
                    {
                        "action_id": action.id,
                        "type": action.type,
                        "status": action.status.value,
                        "error": (
                            str(action.output.result.get("error", ""))
                            if action.output
                            and isinstance(action.output.result, dict)
                            and "error" in action.output.result
                            else None
                        ),
                    },
                )

                # Record in history
                self.history.append(action)

                # Propagate context entries produced by the action
                _ctx_changed = False
                if action.output and action.output.context_entries:
                    if self._agent_queue:
                        for ctx_entry in action.output.context_entries:
                            self._agent_queue.append_context(ctx_entry)
                        _ctx_changed = True

                # Emit context_stats when context was modified
                if _ctx_changed and self._agent_queue:
                    yield self._make_event(
                        EventType.ACTION_COMPLETE,
                        {
                            "action_id": action.id,
                            "type": action.type,
                            "context_stats": self._agent_queue._compute_context_stats(),
                        },
                    )

                # If planner or bio_workflow completed, clear remaining queue.
                # Their newly created ActionSets will take over.
                if (
                    action.type in ("planner", "bio_workflow")
                    and action.status == ActionStatus.COMPLETED
                ):
                    self.queue.clear()
                    self._batch_remaining = 0
                    continue

                # Batch tracking: decrement after non-thinking action
                if action.type != "thinking" and self._batch_remaining > 0:
                    self._batch_remaining -= 1

                # Auto thinking insertion policy:
                # After a non-thinking Action that completed or failed,
                # if there are remaining slots, batch is done, and next
                # action isn't already thinking, insert a thinking Action
                # so the agent can re-evaluate (including error recovery).
                if (
                    action.type != "thinking"
                    and action.status in (
                        ActionStatus.COMPLETED,
                        ActionStatus.FAILED,
                    )
                    and self.remaining_actions() > 0
                    and self._batch_remaining <= 0
                    and not self._next_is_thinking()
                ):
                    thinking = self._make_thinking()
                    if thinking is not None:
                        self.queue.insert(0, thinking)

            self.status = ActionSetStatus.COMPLETED

        except Exception as e:
            self.status = ActionSetStatus.FAILED
            yield self._make_event(
                EventType.ERROR,
                {"message": f"ActionSet '{self.subject}' failed: {e}"},
            )

        # Build completion data with context_stats
        _complete_data = {
            "subject": self.subject,
            "title": self.title,
            "status": self.status.value,
            "actions_completed": len(self.history),
        }
        if self._agent_queue:
            _complete_data["context_stats"] = (
                self._agent_queue._compute_context_stats()
            )

        yield self._make_event(
            EventType.ACTION_SET_COMPLETE,
            _complete_data,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _next_is_thinking(self) -> bool:
        return bool(self.queue) and self.queue[0].type == "thinking"

    def _make_thinking(self) -> Optional[Action]:
        """Create a fresh thinking Action via the injected factory."""
        if self._thinking_factory is not None:
            try:
                return self._thinking_factory()
            except Exception:
                return None
        return None

    def _make_event(
        self,
        event_type: EventType,
        data: dict[str, Any],
    ) -> Event:
        return Event(
            type=event_type,
            data=data,
            queue_id=(
                self._agent_queue.id if self._agent_queue else None
            ),
            set_id=self.id,
        )

    # ------------------------------------------------------------------
    # Serialisation / display
    # ------------------------------------------------------------------

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "subject": self.subject,
            "title": self.title,
            "approach": self.approach,
            "status": self.status.value,
            "max_action_weight": self.max_action_weight,
            "actions_completed": len(self.history),
            "actions_remaining": self.remaining_actions(),
            "created_at": self.created_at.isoformat(),
            "actions": [a.to_dict() for a in self.history],
        }

    def __repr__(self) -> str:
        title_part = f" title='{self.title}' " if self.title else " "
        return (
            f"<ActionSet id={self.id}"
            f"{title_part}"
            f"subject='{self.subject[:30]}' "
            f"status={self.status.value} "
            f"history={len(self.history)} queue={len(self.queue)}>"
        )
