import unittest
from unittest.mock import AsyncMock, patch
from uuid import uuid4

from src_agent.agent import AgentRunner
from src_agent.repositories.runtime_state import RuntimeStateConflict
from src_shared.contracts import (
    AgentInboundRequest,
    WorkflowRunResponse,
    WorkflowRuntimeState,
)


class _Transaction:
    def __init__(self, connection):
        self.connection = connection
        self.exception = None

    async def __aenter__(self):
        return self.connection

    async def __aexit__(self, exc_type, exc, traceback):
        self.exception = exc
        return False


class _Database:
    def __init__(self):
        self.connection = object()
        self.transactions = []

    def transaction(self):
        transaction = _Transaction(self.connection)
        self.transactions.append(transaction)
        return transaction


class AgentPersistenceTest(unittest.IsolatedAsyncioTestCase):
    @staticmethod
    def _request() -> AgentInboundRequest:
        return AgentInboundRequest(
            trace_id=uuid4(),
            request_id=uuid4(),
            session_id=None,
            ts_ms=1,
            text="hello",
        )

    @staticmethod
    def _response(request: AgentInboundRequest) -> WorkflowRunResponse:
        return WorkflowRunResponse(
            trace_id=request.trace_id,
            request_id=request.request_id,
            session_id=request.session_id,
            ts_ms=request.ts_ms,
            status="DONE",
            result="world",
            next_runtime=WorkflowRuntimeState(version=2),
        )

    def _runner(self, database: _Database, response: WorkflowRunResponse) -> AgentRunner:
        runner = AgentRunner(
            nc=object(),
            db=database,
            workflow_subject="workflow.run",
            user_id="test-user",
            workflow_timeout_s=1,
        )
        runner._call_workflow = AsyncMock(return_value=response)
        return runner

    async def test_runtime_state_and_history_share_one_transaction(self) -> None:
        request = self._request()
        response = self._response(request)
        database = _Database()
        runner = self._runner(database, response)

        save_state = AsyncMock(return_value=2)
        persist_turn = AsyncMock()
        with (
            patch("src_agent.agent.save_runtime_state", save_state),
            patch("src_agent.agent.persist_chat_turn", persist_turn),
        ):
            result = await runner._run_with_optimistic_lock(
                req=request,
                session_id=uuid4(),
                runtime=WorkflowRuntimeState(version=1),
            )

        self.assertEqual("DONE", result.status)
        self.assertEqual(1, len(database.transactions))
        self.assertIs(database.connection, save_state.await_args.kwargs["connection"])
        self.assertIs(database.connection, persist_turn.await_args.kwargs["connection"])

    async def test_persistence_failure_does_not_repeat_workflow_side_effects(self) -> None:
        request = self._request()
        response = self._response(request)
        database = _Database()
        runner = self._runner(database, response)

        with (
            patch("src_agent.agent.save_runtime_state", AsyncMock(return_value=2)),
            patch(
                "src_agent.agent.persist_chat_turn",
                AsyncMock(side_effect=RuntimeError("database unavailable")),
            ),
            self.assertRaisesRegex(RuntimeError, "database unavailable"),
        ):
            await runner._run_with_optimistic_lock(
                req=request,
                session_id=uuid4(),
                runtime=WorkflowRuntimeState(version=1),
            )

        self.assertEqual(1, runner._call_workflow.await_count)
        self.assertIsInstance(database.transactions[0].exception, RuntimeError)

    async def test_only_version_conflicts_repeat_the_workflow(self) -> None:
        request = self._request()
        response = self._response(request)
        database = _Database()
        runner = self._runner(database, response)
        refreshed = WorkflowRuntimeState(version=2)

        save_state = AsyncMock(
            side_effect=[RuntimeStateConflict("conflict"), 3],
        )
        with (
            patch("src_agent.agent.save_runtime_state", save_state),
            patch("src_agent.agent.persist_chat_turn", AsyncMock()),
            patch(
                "src_agent.agent.load_runtime_state",
                AsyncMock(return_value=refreshed),
            ) as load_state,
        ):
            result = await runner._run_with_optimistic_lock(
                req=request,
                session_id=uuid4(),
                runtime=WorkflowRuntimeState(version=1),
            )

        self.assertEqual("DONE", result.status)
        self.assertEqual(2, runner._call_workflow.await_count)
        self.assertEqual(1, load_state.await_count)


if __name__ == "__main__":
    unittest.main()