"""Action repository for the MAF roundtable experiment (external writes).

Records every external write / tool action with an idempotency key so the
Action Coordinator (Phase 5) can guarantee at-most-once execution across
retries / restarts. Phase 1 ships the storage + CRUD; the Coordinator itself
is Phase 5. ``arguments_redacted_json`` / ``result_redacted_json`` must
already have credentials / tokens / cookies stripped by the writer — the
store does not re-redact.
"""

from __future__ import annotations

from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.maf_roundtable.model import MafRoundtableActionRow


def _now() -> datetime:
    return datetime.now(UTC)


def action_to_dict(row: MafRoundtableActionRow) -> dict[str, Any]:
    """Serialise an action ORM row to the internal persistence representation."""
    return {
        "id": row.id,
        "session_id": row.session_id,
        "run_id": row.run_id,
        "agent_id": row.agent_id,
        "tool_name": row.tool_name,
        "logical_action_id": row.logical_action_id,
        "turn_index": row.turn_index,
        "idempotency_key": row.idempotency_key,
        "target_key": row.target_key,
        "risk_level": row.risk_level,
        "status": row.status,
        "arguments": row.arguments_redacted_json,
        "result": row.result_redacted_json,
        "external_object_id": row.external_object_id,
        "approval_id": row.approval_id,
        "attempt_count": row.attempt_count,
        "started_at": row.started_at.isoformat() if isinstance(row.started_at, datetime) else row.started_at,
        "finished_at": row.finished_at.isoformat() if isinstance(row.finished_at, datetime) else row.finished_at,
        "created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
        "updated_at": row.updated_at.isoformat() if isinstance(row.updated_at, datetime) else row.updated_at,
    }


def action_to_public_dict(action: Mapping[str, Any]) -> dict[str, Any]:
    """Return only the stable, audit-safe Action DTO fields."""

    return {
        field: action[field]
        for field in _PUBLIC_ACTION_FIELDS
        if field in action
    }


_PUBLIC_ACTION_FIELDS: tuple[str, ...] = (
    "id",
    "session_id",
    "run_id",
    "agent_id",
    "tool_name",
    "logical_action_id",
    "turn_index",
    "risk_level",
    "status",
    "arguments",
    "result",
    "external_object_id",
    "approval_id",
    "attempt_count",
    "started_at",
    "finished_at",
    "created_at",
    "updated_at",
)


# Back-compat alias for internal callers that used the old private name.
_to_dict = action_to_dict


class MafActionStateMachineError(RuntimeError):
    """An action was claimed or transitioned from an unexpected state.

    Raised when a claim loses the ``approved → running`` CAS and the row is in
    a state that is neither ``running`` (in progress elsewhere) nor a recognized
    terminal (``succeeded`` / ``failed`` / ``unknown`` / ``rejected``) — e.g.
    still ``waiting_approval`` / ``pending`` — meaning the caller tried to
    execute an action that was never properly approved.
"""


# Terminal statuses a failed claim may reflect back to the caller.
_TERMINAL_CLAIM_STATUSES = ("succeeded", "failed", "unknown", "rejected")


class MafActionRepository:
    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    async def create_action(self, data: dict[str, Any]) -> dict[str, Any]:
        """Create an action. Raises ``IntegrityError`` if the ``idempotency_key``
        already exists (caller should treat that as "already submitted" and
        fetch the existing row via ``get_by_idempotency_key``).
        """
        now = _now()
        row = MafRoundtableActionRow(
            id=data.get("id") or uuid4().hex,
            session_id=str(data["session_id"]),
            run_id=data.get("run_id"),
            agent_id=data.get("agent_id"),
            tool_name=str(data["tool_name"]),
            logical_action_id=data.get("logical_action_id"),
            turn_index=data.get("turn_index"),
            idempotency_key=str(data["idempotency_key"]),
            target_key=data.get("target_key"),
            risk_level=str(data.get("risk_level") or "R0"),
            status=str(data.get("status") or "pending"),
            arguments_redacted_json=data.get("arguments"),
            result_redacted_json=data.get("result"),
            external_object_id=data.get("external_object_id"),
            approval_id=data.get("approval_id"),
            attempt_count=int(data.get("attempt_count", 0)),
            started_at=data.get("started_at"),
            finished_at=data.get("finished_at"),
            created_at=now,
            updated_at=now,
        )
        async with self._sf() as session:
            session.add(row)
            await session.commit()
        return await self.get_action(row.id)  # type: ignore[arg-type]

    async def get_action(self, action_id: str) -> dict[str, Any] | None:
        async with self._sf() as session:
            stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.id == action_id).limit(1)
            row = (await session.execute(stmt)).scalars().first()
            return _to_dict(row) if row is not None else None

    async def get_by_idempotency_key(self, key: str) -> dict[str, Any] | None:
        async with self._sf() as session:
            stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.idempotency_key == key).limit(1)
            row = (await session.execute(stmt)).scalars().first()
            return _to_dict(row) if row is not None else None

    async def list_actions(
        self, *, session_id: str | None = None, run_id: str | None = None, status: str | None = None
    ) -> list[dict[str, Any]]:
        stmt = select(MafRoundtableActionRow)
        if session_id is not None:
            stmt = stmt.where(MafRoundtableActionRow.session_id == session_id)
        if run_id is not None:
            stmt = stmt.where(MafRoundtableActionRow.run_id == run_id)
        if status is not None:
            stmt = stmt.where(MafRoundtableActionRow.status == status)
        stmt = stmt.order_by(MafRoundtableActionRow.created_at.asc())
        async with self._sf() as session:
            rows = (await session.execute(stmt)).scalars().all()
            return [_to_dict(r) for r in rows]

    async def update_action(self, action_id: str, fields: dict[str, Any]) -> dict[str, Any] | None:
        async with self._sf() as session:
            stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.id == action_id).limit(1)
            row = (await session.execute(stmt)).scalars().first()
            if row is None:
                return None
            col_map = {
                "status": "status",
                "risk_level": "risk_level",
                "arguments": "arguments_redacted_json",
                "result": "result_redacted_json",
                "external_object_id": "external_object_id",
                "approval_id": "approval_id",
                "attempt_count": "attempt_count",
                "started_at": "started_at",
                "finished_at": "finished_at",
            }
            for key, value in fields.items():
                col = col_map.get(key)
                if col is not None:
                    setattr(row, col, value)
            row.updated_at = _now()
            await session.commit()
            return _to_dict(row)

    async def claim_for_execution(self, action_id: str) -> str:
        """Atomically transition an action from ``approved`` to ``running``.

        Single conditional UPDATE — only the caller whose UPDATE affects
        exactly one row wins the right to execute the external write. This
        guarantees at-most-once execution even under concurrent ``resume()``
        calls or retries after a restart.

        Returns the effective claim status (callers must NOT treat a lost claim
        as a blanket "skip"):

        - ``"claimed"`` — won the CAS; the caller may execute.
        - ``"in_progress"`` — another executor already moved it to ``running``.
        - ``"succeeded"`` / ``"failed"`` / ``"unknown"`` / ``"rejected"`` — the
          action is already in that terminal state.

        Raises ``MafActionStateMachineError`` when the row is missing or in any
        other state (e.g. still ``waiting_approval`` / ``pending``) — i.e. it was
        never properly approved and must not be executed.
        """
        now = _now()
        async with self._sf() as session:
            result = await session.execute(
                update(MafRoundtableActionRow)
                .where(
                    MafRoundtableActionRow.id == action_id,
                    MafRoundtableActionRow.status == "approved",
                )
                .values(status="running", started_at=now, updated_at=now)
            )
            await session.commit()
            if result.rowcount == 1:  # type: ignore[union-attr]
                return "claimed"

        # Lost the CAS (or never approved) — report the ACTUAL state.
        current = await self.get_action(action_id)
        if current is None:
            raise MafActionStateMachineError(f"action {action_id} not found during claim")
        status = current.get("status")
        if status == "running":
            return "in_progress"
        if status in _TERMINAL_CLAIM_STATUSES:
            return str(status)
        raise MafActionStateMachineError(
            f"action {action_id} cannot be claimed from state '{status}'"
        )

    async def mark_approved(self, action_id: str) -> bool:
        """Conditionally transition ``waiting_approval → approved``.

        Returns True only if this call performed the transition (``rowcount==1``).
        A concurrent decider or an already-terminal action yields False — the
        caller must NOT overwrite a terminal state via a generic update.
        """
        now = _now()
        async with self._sf() as session:
            result = await session.execute(
                update(MafRoundtableActionRow)
                .where(
                    MafRoundtableActionRow.id == action_id,
                    MafRoundtableActionRow.status == "waiting_approval",
                )
                .values(status="approved", updated_at=now)
            )
            await session.commit()
            return result.rowcount == 1  # type: ignore[union-attr]

    async def mark_rejected(self, action_id: str) -> bool:
        """Conditionally transition ``waiting_approval → rejected`` (+ finished_at).

        Returns True only if this call performed the transition. Terminal states
        are never overwritten.
        """
        now = _now()
        async with self._sf() as session:
            result = await session.execute(
                update(MafRoundtableActionRow)
                .where(
                    MafRoundtableActionRow.id == action_id,
                    MafRoundtableActionRow.status == "waiting_approval",
                )
                .values(status="rejected", finished_at=now, updated_at=now)
            )
            await session.commit()
            return result.rowcount == 1  # type: ignore[union-attr]

    async def record_terminal_result(
        self,
        action_id: str,
        *,
        status: str,
        result: dict[str, Any] | None = None,
        external_object_id: str | None = None,
    ) -> dict[str, Any] | None:
        """Record the terminal outcome of an action execution.

        State-machine constrained: only a ``running`` action may transition
        to a terminal status (``succeeded`` / ``failed`` / ``unknown``).
        A concurrent caller that lost the race (or a replay after restart)
        finds the row already terminal and its UPDATE is a no-op
        (``rowcount == 0``) — it must NOT overwrite the real result.
        """
        if status not in ("succeeded", "failed", "unknown"):
            raise ValueError(f"invalid action result status: {status}")
        now = _now()
        async with self._sf() as session:
            result_q = await session.execute(
                update(MafRoundtableActionRow)
                .where(
                    MafRoundtableActionRow.id == action_id,
                    MafRoundtableActionRow.status == "running",
                )
                .values(
                    status=status,
                    result_redacted_json=result,
                    external_object_id=external_object_id,
                    finished_at=now,
                    updated_at=now,
                )
            )
            await session.commit()
            # Use rowcount instead of .returning() for cross-DB compatibility
            # (SQLite, PostgreSQL, MySQL).
            if result_q.rowcount != 1:  # type: ignore[union-attr]
                # Lost the race or not in running state — do not overwrite.
                return None
            return await self.get_action(action_id)