"""OutputSession context manager for the output rendering framework.

The session manages the lifecycle of all output elements for a single
command invocation.  It bridges format-agnostic producer code and the
format-specific materialisation strategy.

Based on the Output Rendering Framework specification (v3_spec.md §Output
Rendering Framework / OutputSession).
"""

from __future__ import annotations

import threading
import time
from collections import OrderedDict
from collections.abc import Callable
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, Field

from cleveragents.cli.output._ids import generate_handle_id, generate_session_id
from cleveragents.cli.output.handles import (
    MAX_ELEMENTS_PER_SESSION,
    ActionHint,
    ActionHintHandle,
    CodeBlock,
    CodeHandle,
    ColumnDef,
    DiffBlock,
    DiffHandle,
    ElementClosed,
    ElementCreated,
    ElementEvent,
    ElementHandle,
    ElementSnapshot,
    ElementUpdated,
    Panel,
    PanelHandle,
    ProgressHandle,
    ProgressIndicator,
    ProgressStep,
    Separator,
    SeparatorHandle,
    SessionEnd,
    StatusHandle,
    StatusMessage,
    Table,
    TableHandle,
    TextBlock,
    TextHandle,
    Tree,
    TreeHandle,
    TreeNode,
)
from cleveragents.cli.output.selection import select_materializer

if TYPE_CHECKING:
    from cleveragents.cli.output.materializers import MaterializationStrategy


# ---------------------------------------------------------------------------
# StructuredOutput — static snapshot of a complete command output
# ---------------------------------------------------------------------------


class StructuredOutput(BaseModel):
    """Static snapshot of a complete command output.

    Produced by ``OutputSession.snapshot()`` and ``OutputSession.close()``.
    Used by accumulate-mode strategies (json/yaml) for final serialisation.
    """

    command: str
    session_id: str
    status: str = "ok"
    elements: list[ElementSnapshot] = Field(default_factory=list)
    exit_code: int = 0
    timing: dict[str, Any] | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


# ---------------------------------------------------------------------------
# OutputSession
# ---------------------------------------------------------------------------


class OutputSession:
    """A live output document that coordinates element production and
    materialisation.

    The session owns the set of active handles, tracks their lifecycle,
    emits ``ElementEvent`` objects to the active materialisation strategy,
    and provides a ``snapshot()`` method for static inspection.

    Thread Safety
    -------------
    The session is thread-safe.  Multiple producers may create and write
    to handles concurrently.

    Usage
    -----
    ::

        with OutputSession(format="rich") as session:
            panel = session.panel("Title")
            panel.set_entry("key", "value")
    """

    def __init__(
        self,
        *,
        format: str = "rich",
        command: str = "",
        strategy: MaterializationStrategy | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self.command = command
        self.session_id = generate_session_id()
        self.created_at = time.time()
        self._format = format
        self._strategy: MaterializationStrategy = (
            strategy if strategy is not None else select_materializer(format)
        )
        self._handles: OrderedDict[str, ElementHandle[ElementSnapshot]] = OrderedDict()
        self._next_index: int = 0
        self._state: str = "open"
        self._lock: threading.Lock = threading.Lock()
        self._metadata: dict[str, Any] = metadata or {}
        self._exit_code: int = 0
        self._timing: dict[str, Any] | None = None
        self._strategy.on_session_begin(self)

    # --- Element handle factories ---

    def panel(
        self,
        title: str,
        *,
        border_style: str = "rounded",
        priority: str = "normal",
        collapse_hint: str = "auto",
        metadata: dict[str, Any] | None = None,
    ) -> PanelHandle:
        """Create a panel element handle for key-value pair output."""
        if not title:
            raise ValueError("title cannot be empty")
        element = Panel(
            title=title,
            border_style=border_style,
            priority=priority,
            collapse_hint=collapse_hint,
            metadata=metadata or {},
        )
        return self._register_handle(PanelHandle, element, "panel")

    def table(
        self,
        title: str | None = None,
        *,
        columns: list[str | ColumnDef] | None = None,
        summary: dict[str, Any] | None = None,
        max_rows_hint: int | None = None,
        sort_key: str | None = None,
        priority: str = "normal",
        collapse_hint: str = "auto",
        metadata: dict[str, Any] | None = None,
    ) -> TableHandle:
        """Create a table element handle for tabular data."""
        col_defs: list[ColumnDef] = []
        for col in columns or []:
            if isinstance(col, str):
                col_defs.append(ColumnDef(name=col))
            else:
                col_defs.append(col)

        element = Table(
            title=title,
            columns=col_defs,
            summary=summary,
            max_rows_hint=max_rows_hint,
            sort_key=sort_key,
            priority=priority,
            collapse_hint=collapse_hint,
            metadata=metadata or {},
        )
        return self._register_handle(TableHandle, element, "table")

    def status(
        self,
        message: str,
        *,
        level: str = "info",
        detail: str | None = None,
        priority: str = "normal",
        metadata: dict[str, Any] | None = None,
    ) -> StatusHandle:
        """Create a status message handle."""
        if not message:
            raise ValueError("message cannot be empty")
        valid_levels = {"ok", "warn", "error", "info"}
        if level not in valid_levels:
            raise ValueError(f"level must be one of {valid_levels!r}, got {level!r}")
        element = StatusMessage(
            message=message,
            level=level,
            detail=detail,
            priority=priority,
            metadata=metadata or {},
        )
        return self._register_handle(StatusHandle, element, "status")

    def progress(
        self,
        label: str,
        *,
        total: int | None = None,
        indeterminate: bool = False,
        steps: list[str] | None = None,
        priority: str = "normal",
        metadata: dict[str, Any] | None = None,
    ) -> ProgressHandle:
        """Create a progress indicator handle.

        P1-1 fix (Luis review): validates *total* is non-negative at
        the factory level, consistent with ``set_progress()``.
        """
        if not label:
            raise ValueError("label cannot be empty")
        if total is not None and total < 0:
            raise ValueError("total must be non-negative")
        step_objs: list[ProgressStep] | None = None
        if steps is not None:
            step_objs = [ProgressStep(label=s) for s in steps]

        element = ProgressIndicator(
            label=label,
            total=total,
            indeterminate=indeterminate,
            steps=step_objs,
            priority=priority,
            metadata=metadata or {},
        )
        return self._register_handle(ProgressHandle, element, "progress")

    def tree(
        self,
        root_label: str,
        *,
        root_style: str | None = None,
        max_depth_hint: int | None = None,
        show_guides: bool = True,
        priority: str = "normal",
        collapse_hint: str = "auto",
        metadata: dict[str, Any] | None = None,
    ) -> TreeHandle:
        """Create a tree element handle for hierarchical data."""
        if not root_label:
            raise ValueError("root_label cannot be empty")
        root_node = TreeNode(label=root_label, style_hint=root_style)
        element = Tree(
            root=root_node,
            max_depth_hint=max_depth_hint,
            show_guides=show_guides,
            priority=priority,
            collapse_hint=collapse_hint,
            metadata=metadata or {},
        )
        return self._register_handle(TreeHandle, element, "tree")

    def text(
        self,
        content: str = "",
        *,
        wrap: bool = True,
        indent: int = 0,
        priority: str = "normal",
        metadata: dict[str, Any] | None = None,
    ) -> TextHandle:
        """Create a text block handle for free-form text."""
        element = TextBlock(
            content=content,
            wrap=wrap,
            indent=indent,
            priority=priority,
            metadata=metadata or {},
        )
        return self._register_handle(TextHandle, element, "text")

    def code(
        self,
        content: str = "",
        *,
        language: str | None = None,
        line_numbers: bool = False,
        highlight_lines: list[int] | None = None,
        priority: str = "normal",
        metadata: dict[str, Any] | None = None,
    ) -> CodeHandle:
        """Create a code block handle for source code display."""
        element = CodeBlock(
            content=content,
            language=language,
            line_numbers=line_numbers,
            highlight_lines=highlight_lines or [],
            priority=priority,
            metadata=metadata or {},
        )
        return self._register_handle(CodeHandle, element, "code")

    def diff(
        self,
        *,
        file_a: str | None = None,
        file_b: str | None = None,
        priority: str = "normal",
        metadata: dict[str, Any] | None = None,
    ) -> DiffHandle:
        """Create a diff block handle for unified-diff display."""
        element = DiffBlock(
            file_a=file_a,
            file_b=file_b,
            priority=priority,
            metadata=metadata or {},
        )
        return self._register_handle(DiffHandle, element, "diff")

    def separator(self, style: str = "line") -> SeparatorHandle:
        """Create a separator element (auto-closed on creation)."""
        valid_styles = {"line", "blank", "double"}
        if style not in valid_styles:
            raise ValueError(f"style must be one of {valid_styles!r}, got {style!r}")
        element = Separator(style=style)
        handle: SeparatorHandle = self._register_handle(
            SeparatorHandle, element, "separator"
        )
        handle.close()
        return handle

    def action_hint(
        self,
        commands: list[str],
        description: str = "",
    ) -> ActionHintHandle:
        """Create an action-hint element (auto-closed on creation)."""
        if not commands:
            raise ValueError("commands cannot be empty")
        element = ActionHint(commands=list(commands), description=description)
        handle: ActionHintHandle = self._register_handle(
            ActionHintHandle, element, "action_hint"
        )
        handle.close()
        return handle

    # --- Internal registration ---

    def _register_handle[H: ElementHandle[Any]](
        self,
        handle_cls: Callable[..., H],
        element: ElementSnapshot,
        element_kind: str,
    ) -> H:
        """Create a handle, register it, and emit a created event.

        Uses a single lock scope to prevent TOCTOU races between state
        check, index allocation, handle registration, and event dispatch.
        The ``ElementCreated`` event is dispatched inside the lock to
        prevent a race where ``session.close()`` could force-close the
        handle before the created event reaches the strategy (N1 fix).

        Enforces ``MAX_ELEMENTS_PER_SESSION`` soft limit (F3 fix).
        """
        with self._lock:
            if self._state != "open":
                raise RuntimeError("Cannot create handles on a closed session")
            if len(self._handles) >= MAX_ELEMENTS_PER_SESSION:
                limit = MAX_ELEMENTS_PER_SESSION
                raise ValueError(f"Maximum elements per session ({limit}) exceeded")
            handle_id = generate_handle_id()
            idx = self._next_index
            self._next_index += 1

            handle: H = handle_cls(
                handle_id=handle_id,
                declaration_index=idx,
                session=self,
                element=element,
            )
            self._handles[handle_id] = handle

            # Dispatch inside the lock to ensure ElementCreated arrives
            # before any ElementClosed from session.close() (N1 fix).
            event = ElementCreated(
                event_type="created",
                handle_id=handle_id,
                element_kind=element_kind,
                declaration_index=idx,
                initial_state=element,
            )
            self._dispatch_event(event)

        return handle

    # --- Event dispatch ---

    def _dispatch_event(self, event: ElementEvent) -> None:
        """Route an event to the active strategy."""
        if isinstance(event, ElementCreated):
            self._strategy.on_element_created(event)
        elif isinstance(event, ElementUpdated):
            self._strategy.on_element_updated(event)
        elif isinstance(event, ElementClosed):
            self._strategy.on_element_closed(event)
        elif isinstance(event, SessionEnd):
            self._strategy.on_session_end(event)

    # --- Session operations ---

    def snapshot(self) -> StructuredOutput:
        """Return a static snapshot of all elements accumulated so far.

        F6 fix: uses ``element_copy()`` for deep copies to prevent
        mutation of live session state.

        N2 fix: includes ``exit_code`` from the session.

        P1-2 fix (Luis review): includes ``timing`` when available so
        that accumulate strategies (JSON/YAML) emit the timing field.

        M4 fix: derives ``status`` from ``exit_code`` — ``"ok"`` when
        zero, ``"error"`` otherwise.
        """
        with self._lock:
            elements = [h.element_copy() for h in self._handles.values()]
        status = "error" if self._exit_code != 0 else "ok"
        snap = StructuredOutput(
            command=self.command,
            session_id=self.session_id,
            status=status,
            elements=list(elements),
            exit_code=self._exit_code,
            metadata=dict(self._metadata),
        )
        if self._timing is not None:
            snap.timing = self._timing
        return snap

    def close(self, *, exit_code: int = 0) -> StructuredOutput:
        """Close the session and finalise all output."""
        with self._lock:
            if self._state == "closed":
                # Already closed — return snapshot with cached timing
                # (R2-C2 fix: previously timing was None on double-close).
                elements = [h.element_copy() for h in self._handles.values()]
                status = "error" if self._exit_code != 0 else "ok"
                snap = StructuredOutput(
                    command=self.command,
                    session_id=self.session_id,
                    status=status,
                    elements=list(elements),
                    exit_code=self._exit_code,
                    metadata=dict(self._metadata),
                )
                snap.timing = self._timing
                return snap
            self._state = "closed"
            self._exit_code = exit_code

        # Force-close any open handles
        for handle in list(self._handles.values()):
            if handle.is_open:
                handle.close()

        snap = self.snapshot()
        now = time.time()
        self._timing = {
            "start_time": self.created_at,
            "end_time": now,
            "duration": now - self.created_at,
        }
        snap.timing = self._timing

        event = SessionEnd(
            event_type="session_end",
            handle_id="",
            element_kind="session",
            exit_code=exit_code,
        )
        self._dispatch_event(event)
        return snap

    # --- Context manager ---

    def __enter__(self) -> OutputSession:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,
    ) -> None:
        if self._state != "closed":
            code = 1 if exc_type is not None else 0
            self.close(exit_code=code)
