from __future__ import annotations

import asyncio
import json
import os
from contextlib import asynccontextmanager
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import UUID

import pytest

os.environ.setdefault(
    "DATABASE_URL",
    "postgresql+asyncpg://test:test@localhost/test",
)

from src.db import base as db_base  # noqa: E402
from src.db.awooop_models import AwoooPRunState  # noqa: E402
from src.jobs import baseline_snapshot  # noqa: E402
from src.services import (  # noqa: E402
    audit_sink,
)
from src.services import (  # noqa: E402
    paid_provider_canary_authorization as authorization,
)


class _MappingResult:
    def __init__(self, mapping: dict[str, Any]) -> None:
        self.mapping = mapping

    def mappings(self) -> _MappingResult:
        return self

    def one(self) -> dict[str, Any]:
        return self.mapping


class _ConcurrentLedger:
    def __init__(self) -> None:
        self.lock = asyncio.Lock()
        self.rows: list[dict[str, Any]] = []
        self.statements: list[str] = []


class _ConcurrentDb:
    def __init__(self, ledger: _ConcurrentLedger) -> None:
        self.ledger = ledger
        self.lock_held = False

    async def execute(
        self,
        statement: Any,
        params: dict[str, Any] | None = None,
    ) -> _MappingResult | None:
        sql = " ".join(str(statement).split())
        self.ledger.statements.append(sql)
        if "pg_advisory_xact_lock" in sql:
            await self.ledger.lock.acquire()
            self.lock_held = True
            return None
        if "SELECT COUNT(*) AS matching_count" in sql:
            return _MappingResult(
                {
                    "matching_count": len(self.ledger.rows),
                    "created_at": (
                        datetime(2026, 7, 16, 8, 0, tzinfo=UTC)
                        if self.ledger.rows
                        else None
                    ),
                }
            )
        if "INSERT INTO alert_operation_log" in sql:
            assert params is not None
            self.ledger.rows.append(json.loads(str(params["context"])))
            return None
        raise AssertionError(f"unexpected SQL: {sql}")


@pytest.mark.asyncio
async def test_required_audit_is_awaited_exact_once_without_race(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    ledger = _ConcurrentLedger()

    @asynccontextmanager
    async def fake_db_context(project_id: str):
        assert project_id == "awoooi"
        db = _ConcurrentDb(ledger)
        try:
            yield db
        finally:
            if db.lock_held:
                ledger.lock.release()

    monkeypatch.setattr(db_base, "get_db_context", fake_db_context)
    kwargs = {
        "project_id": "awoooi",
        "action": "run.approval.approve",
        "resource_type": "run",
        "resource_id": "267a3d26-3c29-470f-8fe3-388a2f408f39",
        "run_id": "267a3d26-3c29-470f-8fe3-388a2f408f39",
        "trace_id": "trace-paid-canary",
        "details": {
            "approver_id": "operator:test",
            "decision": "approve",
            "new_state": "running",
        },
        "require_durable": True,
        "idempotency_key": "paid-canary:exact-run",
    }

    first, second = await asyncio.gather(
        audit_sink.write_audit(**kwargs),
        audit_sink.write_audit(**kwargs),
    )

    assert {first["status"], second["status"]} == {
        "inserted_verified",
        "duplicate_existing",
    }
    assert first["durable_write_ack"] is True
    assert second["durable_write_ack"] is True
    assert len(ledger.rows) == 1
    context = ledger.rows[0]
    assert context["schema_version"] == "awooop_sanitized_audit_v1"
    assert context["project_id"] == "awoooi"
    assert context["run_id"] == kwargs["run_id"]
    assert context["trace_id"] == "trace-paid-canary"
    assert context["details"]["decision"] == "approve"
    sql = " ".join(ledger.statements)
    assert "pg_advisory_xact_lock" in sql
    assert "INSERT INTO alert_operation_log" in sql
    assert "INSERT INTO audit_logs" not in sql


@pytest.mark.asyncio
async def test_required_audit_rejects_preexisting_duplicate_rows(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    ledger = _ConcurrentLedger()
    ledger.rows.extend([{}, {}])

    @asynccontextmanager
    async def fake_db_context(project_id: str):
        db = _ConcurrentDb(ledger)
        try:
            yield db
        finally:
            if db.lock_held:
                ledger.lock.release()

    monkeypatch.setattr(db_base, "get_db_context", fake_db_context)

    with pytest.raises(audit_sink.AuditDurabilityError) as raised:
        await audit_sink.write_audit(
            project_id="awoooi",
            action="run.approval.approve",
            resource_type="run",
            resource_id="267a3d26-3c29-470f-8fe3-388a2f408f39",
            run_id="267a3d26-3c29-470f-8fe3-388a2f408f39",
            trace_id="trace-paid-canary",
            details={"decision": "approve"},
            require_durable=True,
            idempotency_key="paid-canary:duplicate-run",
        )

    assert raised.value.error_code == "audit_idempotency_duplicate_detected"
    assert len(ledger.rows) == 2


class _SequenceResult:
    def __init__(
        self,
        *,
        scalar: Any = None,
        row: Any = None,
        mapping: dict[str, Any] | None = None,
    ) -> None:
        self.scalar = scalar
        self.row = row
        self.mapping = mapping

    def scalar_one_or_none(self) -> Any:
        return self.scalar

    def one(self) -> Any:
        return self.row if self.row is not None else self.mapping

    def mappings(self) -> _SequenceResult:
        return self


class _SequenceDb:
    def __init__(self, results: list[_SequenceResult]) -> None:
        self.results = results
        self.calls: list[tuple[Any, dict[str, Any] | None]] = []

    async def execute(
        self,
        statement: Any,
        params: dict[str, Any] | None = None,
    ) -> _SequenceResult:
        self.calls.append((statement, params))
        return self.results.pop(0)


@pytest.mark.asyncio
async def test_paid_canary_reader_uses_exact_same_run_user_action_schema(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    run_id = UUID("267a3d26-3c29-470f-8fe3-388a2f408f39")
    trace_id = "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01"
    now = datetime(2026, 7, 16, 8, 0, tzinfo=UTC)
    run = AwoooPRunState(
        run_id=run_id,
        project_id="awoooi",
        agent_id="awoooi-paid-provider-canary",
        state="running",
        trace_id=trace_id,
        trigger_type="api",
        trigger_ref="paid-provider-canary:AIA-SRE-013",
        is_shadow=False,
        input_sha256=authorization.paid_provider_canary_authorization_input_sha256(),
        timeout_at=now + timedelta(minutes=5),
    )
    db = _SequenceDb(
        [
            _SequenceResult(scalar=run),
            _SequenceResult(row=(1, "success", False, now - timedelta(seconds=30))),
            _SequenceResult(
                mapping={
                    "approval_audit_count": 1,
                    "approval_audit_created_at": now - timedelta(seconds=20),
                }
            ),
            _SequenceResult(mapping={"prior_canary_start_count": 0}),
        ]
    )

    @asynccontextmanager
    async def fake_db_context(project_id: str):
        assert project_id == "awoooi"
        yield db

    monkeypatch.setattr(authorization, "get_db_context", fake_db_context)

    evidence = await authorization._SqlPaidProviderCanaryAuthorizationEvidenceReader().read(
        run_id=run_id,
        project_id="awoooi",
        work_item_id="AIA-SRE-013",
    )

    assert evidence is not None
    assert evidence.approval_audit_count == 1
    audit_statement, audit_params = db.calls[2]
    audit_sql = " ".join(str(audit_statement).split())
    assert "FROM alert_operation_log" in audit_sql
    assert "event_type = 'USER_ACTION'" in audit_sql
    assert "FROM audit_logs" not in audit_sql
    assert "context ->> 'project_id' = :project_id" in audit_sql
    assert "context ->> 'resource_id' = :run_id" in audit_sql
    assert "context ->> 'run_id' = :run_id" in audit_sql
    assert "context ->> 'trace_id' = :trace_id" in audit_sql
    assert "context #>> '{details,decision}' = 'approve'" in audit_sql
    assert audit_params is not None
    assert audit_params["run_id"] == str(run_id)
    assert audit_params["trace_id"] == trace_id
    expected_key = authorization.paid_provider_canary_approval_audit_idempotency_key(
        run_id
    )
    assert audit_params["idempotency_key_sha256"] == (
        audit_sink.audit_idempotency_key_sha256(expected_key)
    )


@pytest.mark.asyncio
async def test_baseline_mcp_count_uses_mcp_audit_ledger(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    statements: list[Any] = []

    class _Db:
        async def execute(self, statement: Any) -> _SequenceResult:
            statements.append(statement)
            return _SequenceResult(scalar=7)

    @asynccontextmanager
    async def fake_db_context():
        yield _Db()

    monkeypatch.setattr(baseline_snapshot, "get_db_context", fake_db_context)

    count = await baseline_snapshot._count_mcp_calls_24h(
        datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
    )

    assert count == 7
    sql = " ".join(str(statements[0]).split())
    assert "mcp_audit_log" in sql
    assert "audit_logs.action" not in sql