"""Concrete AgentRecorder backed by ExecutionJournal + ArtifactStore.

Enforces basic invariants (paired enter/exit, no duplicate unit_ids)
and provides the ``unit()`` context manager for guaranteed cleanup.
All writes are no-ops when the journal is ``None``.
"""

from __future__ import annotations

import logging
from collections.abc import Iterator
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Literal

from easycat.integrations.agents.base import (
    CancellationMode,
    ExecutionCursor,
    RecorderContext,
    RecorderInvariantError,
)
from easycat.runtime.record_contracts import validate_builtin_record
from easycat.runtime.records import ErrorInfo, JournalRecordKind

if TYPE_CHECKING:
    from easycat.runtime.artifacts import ArtifactStore
    from easycat.runtime.journal import ExecutionJournal

logger = logging.getLogger(__name__)

_TOKEN_COUNT_KEYS = frozenset(
    {
        "input_tokens",
        "output_tokens",
        "cached_input_tokens",
    }
)


class JournalAgentRecorder:
    """Concrete ``AgentRecorder`` that writes to the ExecutionJournal.

    Bridges receive a fresh instance per ``invoke()`` call, bound to
    the current run/session/turn context.  All writes pass through the
    journal ``apply_write_filter`` hook before append.
    """

    def __init__(
        self,
        journal: ExecutionJournal | None,
        artifact_store: ArtifactStore | None,
        context: RecorderContext,
    ) -> None:
        self._journal = journal
        self._artifact_store = artifact_store
        self._context = context
        # Stack of open cursors for invariant enforcement.
        self._open_cursors: list[ExecutionCursor] = []
        self._seen_unit_ids: set[str] = set()

    @property
    def context(self) -> RecorderContext:
        return self._context

    # ── Unit lifecycle ───────────────────────────────────────────

    def record_unit_entered(self, cursor: ExecutionCursor) -> None:
        if cursor.unit_id in self._seen_unit_ids:
            raise RecorderInvariantError(f"Duplicate unit_id {cursor.unit_id!r} within this turn")
        self._seen_unit_ids.add(cursor.unit_id)
        self._open_cursors.append(cursor)
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="unit_entered",
            data={
                "unit_id": cursor.unit_id,
                "unit_kind": (
                    cursor.unit_kind.value
                    if hasattr(cursor.unit_kind, "value")
                    else str(cursor.unit_kind)
                ),
                "display_name": cursor.display_name,
                "parent_unit_id": cursor.parent_unit_id,
                "committable": cursor.committable,
                "direction": "enter",
            },
        )

    def record_unit_exited(self, cursor: ExecutionCursor, reason: str | None = None) -> None:
        if not self._open_cursors:
            raise RecorderInvariantError(
                f"record_unit_exited({cursor.unit_id!r}) without a matching record_unit_entered"
            )
        top = self._open_cursors[-1]
        if top.unit_id != cursor.unit_id:
            raise RecorderInvariantError(
                f"record_unit_exited({cursor.unit_id!r}) but the top of the "
                f"cursor stack is {top.unit_id!r}"
            )
        self._open_cursors.pop()
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="unit_exited",
            data={
                "unit_id": cursor.unit_id,
                "unit_kind": (
                    cursor.unit_kind.value
                    if hasattr(cursor.unit_kind, "value")
                    else str(cursor.unit_kind)
                ),
                "display_name": cursor.display_name,
                "committable": cursor.committable,
                "exit_reason": reason,
                "direction": "exit",
            },
        )

    def safe_exit_cursor(self, cursor: ExecutionCursor, reason: str | None = "error") -> None:
        """Close ``cursor`` defensively, swallowing recorder errors.

        Bridges call this from their error / cancellation cleanup arms so a
        recorder fault can't mask the original exception while still keeping
        the strict enter/exit stack invariant intact for the postmortem
        journal.  Logs and continues on failure; never raises.
        """
        try:
            self.record_unit_exited(cursor, reason=reason)
        except Exception:
            logger.debug("Failed to close cursor %r during cleanup", cursor.unit_id, exc_info=True)

    @contextmanager
    def unit(
        self, cursor: ExecutionCursor, *, commit_on_exit: bool = True
    ) -> Iterator[ExecutionCursor]:
        """Context manager wrapping enter/exit with guaranteed cleanup."""
        self.record_unit_entered(cursor)
        try:
            yield cursor
        except BaseException as exc:
            self.record_unit_exited(cursor, reason=f"exception:{type(exc).__name__}")
            raise
        else:
            final = cursor.with_committable(True) if commit_on_exit else cursor
            self.record_unit_exited(final, reason=None)

    @contextmanager
    def turn_cursor(self, cursor: ExecutionCursor) -> Iterator[ExecutionCursor]:
        """Per-turn cursor lifecycle: enter, error/cancel cleanup, clean commit.

        Centralizes the enter → error → BaseException → clean-exit ordering that
        the shipped bridges used to hand-roll: a plain ``Exception`` records a
        framework error plus an ``error`` unit-exit, a ``BaseException``
        (CancelledError / GeneratorExit from an AgentRunner timeout or barge-in
        ``aclose``) defensively closes the cursor via ``safe_exit_cursor``, and a
        clean exit commits the cursor.
        """
        self.record_unit_entered(cursor)
        try:
            yield cursor
        except Exception as exc:
            self.record_framework_error(ErrorInfo.from_exception(exc))
            self.record_unit_exited(cursor, reason="error")
            raise
        except BaseException:
            self.safe_exit_cursor(cursor)
            raise
        else:
            self.record_unit_exited(cursor.with_committable(True), reason=None)

    # ── Tool calls ───────────────────────────────────────────────

    def record_tool_call(
        self,
        phase: Literal["start", "delta", "result", "error"],
        name: str,
        args_ref: str | None = None,
        result_ref: str | None = None,
        call_id: str | None = None,
    ) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="tool_phase_changed",
            data={
                "phase": phase,
                "tool_name": name,
                "args_ref": args_ref,
                "result_ref": result_ref,
                "call_id": call_id,
            },
        )

    # ── State snapshots ──────────────────────────────────────────

    def record_state_snapshot(self, ref: str, *, payload: bytes | None = None) -> str:
        stored_ref: str | None = ref
        if payload is not None and self._artifact_store is not None:
            result = self._artifact_store.put(payload)
            stored_ref = result if result else None
        else:
            if payload is not None and self._artifact_store is None:
                logger.warning(
                    "record_state_snapshot: payload dropped (no artifact store); "
                    "ref %r will not resolve to stored bytes",
                    ref,
                )
            stored_ref = None
        actual_ref = stored_ref or ref
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="state_snapshot",
            data={"state_ref": actual_ref},
            output_ref=stored_ref,
        )
        return actual_ref

    # ── Framework handoffs ───────────────────────────────────────

    def record_framework_handoff(
        self,
        from_unit: str | None,
        to_unit: str,
        reason: str | None = None,
    ) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="framework_handoff",
            data={
                "from_unit": from_unit,
                "to_unit": to_unit,
                "handoff_reason": reason,
            },
        )

    # ── Cancellation boundaries ──────────────────────────────────

    def record_cancellation_boundary(
        self,
        mode: CancellationMode,
        reason: str | None = None,
        caused_by_signal_id: str | None = None,
    ) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="cancellation_boundary",
            data={
                "cancellation_mode": mode.value,
                "boundary_reason": reason,
                "caused_by_signal_id": caused_by_signal_id,
            },
        )

    # ── Framework errors ─────────────────────────────────────────

    def record_framework_error(self, error: ErrorInfo) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="framework_error",
            error=error,
        )

    def record_usage(
        self,
        *,
        provider: str | None = None,
        model: str | None = None,
        input_tokens: int | None = None,
        output_tokens: int | None = None,
        cached_input_tokens: int | None = None,
    ) -> None:
        """Record provider-reported token counts without pricing assumptions."""
        raw_counts = {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cached_input_tokens": cached_input_tokens,
        }
        counts = {
            key: value
            for key, value in raw_counts.items()
            if isinstance(value, int) and not isinstance(value, bool) and value >= 0
        }
        if not counts or not any(counts.values()):
            return
        if (
            "cached_input_tokens" in counts
            and "input_tokens" in counts
            and counts["cached_input_tokens"] > counts["input_tokens"]
        ):
            counts.pop("cached_input_tokens")
        data: dict[str, Any] = dict(counts)
        if provider:
            data["provider"] = provider
        if model:
            data["model"] = model
        self._append(
            kind=JournalRecordKind.EVENT,
            name="agent_usage",
            data=data,
        )

    # ── Atomic interruption records ─────────────────────────────

    def record_state_committed(
        self,
        mutation_kind: str,
        pre_state_ref: str | None = None,
        post_state_ref: str | None = None,
    ) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="state_committed",
            data={
                "mutation_kind": mutation_kind,
                "pre_state_ref": pre_state_ref,
                "post_state_ref": post_state_ref,
                "direction": "enter",
            },
        )

    def record_interruption_apply_failed(
        self,
        mutation_kind: str,
        pre_state_ref: str | None = None,
        post_state_ref: str | None = None,
        failure_error: ErrorInfo | None = None,
    ) -> None:
        self._append(
            kind=JournalRecordKind.FRAMEWORK_TRANSITION,
            name="interruption_apply_failed",
            data={
                "mutation_kind": mutation_kind,
                "pre_state_ref": pre_state_ref,
                "post_state_ref": post_state_ref,
            },
            error=failure_error,
        )

    # ── Internal journal write ───────────────────────────────────

    def _append(
        self,
        kind: JournalRecordKind,
        name: str,
        data: dict[str, Any] | None = None,
        error: ErrorInfo | None = None,
        output_ref: str | None = None,
    ) -> None:
        if self._journal is None:
            return
        record_data = {"run_id": self._context.run_id}
        if data:
            record_data.update(data)
        validate_builtin_record(name=name, kind=kind, data=record_data)
        self._journal.append(
            kind=kind,
            name=name,
            session_id=self._context.session_id,
            turn_id=self._context.turn_id,
            data=_scrub_secrets(record_data),
            error=error,
            output_ref=output_ref,
        )


def _scrub_secrets(data: dict[str, Any]) -> dict[str, Any]:
    """Remove data keys whose names match secret-adjacent fragments.

    Keeps bridge metadata from carrying raw API keys, auth headers, or
    credentials into the journal. The scrub recurses into nested dicts and
    lists so secrets buried below the top level (e.g.
    ``{"headers": {"Authorization": "…"}}``) are dropped too.
    """
    from easycat.runtime.safe_defaults import _is_secret_name

    def _scrub(value: Any) -> Any:
        if isinstance(value, dict):
            return {
                k: _scrub(v)
                for k, v in value.items()
                if not _is_secret_name(str(k))
                or (
                    k in _TOKEN_COUNT_KEYS
                    and isinstance(v, int)
                    and not isinstance(v, bool)
                    and v >= 0
                )
            }
        if isinstance(value, list):
            return [_scrub(item) for item in value]
        if isinstance(value, tuple):
            return tuple(_scrub(item) for item in value)
        return value

    return _scrub(data)
