"""Managed-PG qualification: 0027 upgrades, recovery quarantine and identity races."""

from __future__ import annotations

import importlib
import os
import re
import threading
import uuid
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from types import SimpleNamespace

import pytest
from alembic.config import Config
from sqlalchemy import create_engine, select, text
from sqlalchemy.engine import URL, Engine, make_url
from sqlalchemy.exc import DBAPIError
from sqlalchemy.orm import Session

from modelforge_api.api.authorization import Principal, PrincipalRole
from modelforge_api.domain.managed_identity import (
    ManagedIdentityChange,
    ManagedIdentityCreate,
    ManagedIdentityRevoke,
    ManagedIdentityRotate,
)
from modelforge_api.domain.release import TARGET_SCHEMA_REVISION
from modelforge_api.persistence.models import (
    AuditEvent,
    ManagedOperatorCredentialActivation,
    ManagedOperatorState,
)
from modelforge_api.services.audit import AuditWriter
from modelforge_api.services.recovery_fingerprint import (
    audit_fingerprint_differences_are_compatible,
    fingerprint_compatibility,
    semantic_fingerprint,
)
from modelforge_api.services.recovery_postgres import CommandResult, PostgresEngine, PostgresTarget
from modelforge_api.settings import Settings
from tests.test_configuration_postgres import _database_tool

PRIOR = "20260908_0026"
CURRENT = TARGET_SCHEMA_REVISION
EXTERNAL_ADMIN = Principal("legacy-operator", PrincipalRole.ADMIN, "legacy_admin_token")
REASON = "Isolated managed PostgreSQL qualification"


@dataclass
class IdentityDatabase:
    runtime: Engine
    restored_runtime: Engine
    owner: URL
    restored_owner: URL

    def migration(self, *, restored: bool = False) -> Config:
        config = Config()
        config.set_main_option("script_location", str(Path(__file__).resolve().parents[1] / "alembic"))
        url = self.restored_owner if restored else self.owner
        config.set_main_option(
            "sqlalchemy.url", url.render_as_string(hide_password=False).replace("%", "%%")
        )
        return config

    def archive(self, path: Path, *, restore: bool = False) -> None:
        owner = self.restored_owner if restore else self.owner
        host = (
            "127.0.0.1" if os.getenv("MODELFORGE_TEST_POSTGRES_CONTAINER")
            else owner.host or "localhost"
        )
        args = ["--host", host, "--port", str(owner.port or 5432), "--username",
                owner.username or "modelforge", "--dbname", owner.database or ""]
        _database_tool(
            "pg_restore" if restore else "pg_dump",
            [*args, "--exit-on-error" if restore else "--format=custom"],
            owner.password, path,
        )
        assert path.stat().st_size > 0


@pytest.fixture
def identity_database() -> Iterator[IdentityDatabase]:
    admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
    if not admin_url:
        pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
    owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
    runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
    assert owner_url and runtime_url, "managed qualification requires owner and runtime roles"
    owner, runtime = make_url(owner_url), make_url(runtime_url)
    assert owner.username == "modelforge"
    assert runtime.username == "modelforge_runtime"
    names = ["test_identities_" + uuid.uuid4().hex,
             "test_identities_restore_" + uuid.uuid4().hex]
    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
    runtime_engine = create_engine(runtime.set(database=names[0]), connect_args={
        "connect_timeout": 15, "options": "-c statement_timeout=30000 -c lock_timeout=15000",
    })
    restored_engine = create_engine(runtime.set(database=names[1]))
    created: list[str] = []
    try:
        with admin.connect() as connection:
            for name in names:
                connection.execute(text(f'CREATE DATABASE "{name}" OWNER modelforge'))
                created.append(name)
        yield IdentityDatabase(runtime_engine, restored_engine,
                               owner.set(database=names[0]), owner.set(database=names[1]))
    finally:
        runtime_engine.dispose()
        restored_engine.dispose()
        with admin.connect() as connection:
            for name in created:
                connection.execute(text(f'DROP DATABASE "{name}" WITH (FORCE)'))
        admin.dispose()


def _service(session: Session, principal: Principal = EXTERNAL_ADMIN):
    from modelforge_api.services.operator_identities import ManagedIdentityService

    return ManagedIdentityService(session, principal, "managed-pg-qualification")


def _create(engine: Engine, slug: str):
    with Session(engine) as session:
        revision = session.scalar(select(ManagedOperatorState.revision))
        assert revision is not None
        return _service(session).create(ManagedIdentityCreate(
            slug=slug, role="admin", expires_in_days=14,
            expected_revision=revision, reason=REASON,
        ))


def test_managed_identity_upgrade_prior_dump_restore_and_least_privilege(
    identity_database: IdentityDatabase, tmp_path: Path,
) -> None:
    database = identity_database
    alembic = importlib.import_module("alembic.command")
    config = database.migration()
    alembic.upgrade(config, PRIOR)
    with Session(database.runtime) as session:
        AuditWriter(session, "operator", "prior-release").write(
            "qualification.baseline", "managed_identity", None, {"source_schema": PRIOR},
        )
        session.commit()
        source = semantic_fingerprint(session.connection())
    archive = tmp_path / "pre-managed-identities.dump"
    database.archive(archive)
    alembic.upgrade(config, "head")
    with Session(database.runtime) as session:
        assert session.scalar(text("SELECT version_num FROM alembic_version")) == CURRENT
        assert session.scalar(select(ManagedOperatorState.revision)) == 0
        current = semantic_fingerprint(session.connection())
        comparison = fingerprint_compatibility(source, current, source_schema_revision=PRIOR)
        assert comparison.compatible
        assert audit_fingerprint_differences_are_compatible(source, current, comparison)
    alembic.downgrade(config, PRIOR)
    alembic.upgrade(config, "head")
    _create(database.runtime, "qualified-admin")
    with database.runtime.connect() as connection:
        role = connection.execute(text(
            "SELECT rolsuper, rolcreatedb, rolcreaterole FROM pg_roles WHERE rolname=current_user"
        )).one()
        assert tuple(role) == (False, False, False)
        for statement in (
            "ALTER TABLE managed_operator_identities ADD COLUMN forbidden integer",
            "TRUNCATE managed_operator_credentials CASCADE",
            "DELETE FROM managed_operator_identities",
            "DELETE FROM audit_events",
        ):
            with pytest.raises(DBAPIError):
                connection.execute(text(statement))
            connection.rollback()
    with pytest.raises(RuntimeError, match="history exists"):
        alembic.downgrade(config, PRIOR)
    with database.runtime.connect() as connection:
        assert connection.scalar(text("SELECT version_num FROM alembic_version")) == CURRENT
    database.archive(archive, restore=True)
    with database.restored_runtime.connect() as connection:
        restored = semantic_fingerprint(connection)
        assert restored["tables"] == source["tables"]
        assert connection.scalar(text("SELECT version_num FROM alembic_version")) == PRIOR
    alembic.upgrade(database.migration(restored=True), "head")
    with database.restored_runtime.connect() as connection:
        restored = semantic_fingerprint(connection)
        comparison = fingerprint_compatibility(source, restored, source_schema_revision=PRIOR)
        assert comparison.compatible
        assert audit_fingerprint_differences_are_compatible(source, restored, comparison)


@pytest.mark.parametrize("operation", ["revoke", "demote", "disable"])
def test_concurrent_cross_admin_changes_preserve_a_usable_admin(
    identity_database: IdentityDatabase, operation: str,
) -> None:
    from modelforge_api.services.operator_identities import (
        ManagedIdentityError,
        authenticate_managed_operator,
    )

    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    alpha, beta = _create(database.runtime, "alpha"), _create(database.runtime, "beta")
    barrier = threading.Barrier(2)

    def contender(pair) -> str:
        initiator, target = pair
        with Session(database.runtime) as session:
            principal = authenticate_managed_operator(session, initiator.one_time_key)
            session.commit()
            barrier.wait(timeout=30)
            try:
                service = _service(session, principal)
                if operation == "revoke":
                    service.revoke(target.identity.slug, ManagedIdentityRevoke(
                        expected_version=target.identity.version,
                        confirm_subject=target.identity.subject, reason=REASON,
                    ))
                else:
                    service.change(target.identity.slug, ManagedIdentityChange(
                        expected_version=target.identity.version,
                        role="viewer" if operation == "demote" else "admin",
                        enabled=operation != "disable", reason=REASON,
                    ))
                return "applied"
            except ManagedIdentityError as error:
                assert error.code == "MANAGED_ADMIN_NO_LONGER_AUTHORIZED"
                return "rejected"

    with ThreadPoolExecutor(max_workers=2) as pool:
        assert sorted(pool.map(contender, [(alpha, beta), (beta, alpha)])) == ["applied", "rejected"]
    usable_admins = 0
    for issued in (alpha, beta):
        with Session(database.runtime) as session:
            try:
                principal = authenticate_managed_operator(session, issued.one_time_key)
                usable_admins += principal.role is PrincipalRole.ADMIN
            except ManagedIdentityError:
                pass
    assert usable_admins == 1
    with Session(database.runtime) as session:
        events = list(session.scalars(select(AuditEvent).where(
            AuditEvent.action.in_(["managed_identity.revoked", "managed_identity.changed"]),
        )))
        assert len(events) == 1
        assert events[0].actor_id in {"managed:alpha", "managed:beta"}


def test_simultaneous_rotation_uses_version_cas(identity_database: IdentityDatabase) -> None:
    from modelforge_api.services.operator_identities import ManagedIdentityError

    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    identity = _create(database.runtime, "rotate-admin")
    barrier = threading.Barrier(2)

    def contender(_index: int) -> str:
        with Session(database.runtime) as session:
            barrier.wait(timeout=30)
            try:
                _service(session).rotate("rotate-admin", ManagedIdentityRotate(
                    expected_version=identity.identity.version, expires_in_days=14,
                    overlap_seconds=0, reason=REASON,
                ))
                return "applied"
            except ManagedIdentityError as error:
                assert error.status_code == 409
                assert error.code == "MANAGED_IDENTITY_VERSION_CONFLICT"
                return "conflict"

    with ThreadPoolExecutor(max_workers=2) as pool:
        assert sorted(pool.map(contender, [0, 1])) == ["applied", "conflict"]
    with Session(database.runtime) as session:
        assert len(list(session.scalars(select(AuditEvent).where(
            AuditEvent.action == "managed_identity.rotated",
        )))) == 1


def test_revoked_initiator_cannot_commit_a_previously_authorized_change(
    identity_database: IdentityDatabase,
) -> None:
    from modelforge_api.services.operator_identities import (
        ManagedIdentityError,
        authenticate_managed_operator,
    )

    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    alpha, beta = _create(database.runtime, "alpha"), _create(database.runtime, "beta")
    authenticated = threading.Event()
    revoked = threading.Event()

    def stale_request() -> str:
        with Session(database.runtime) as session:
            principal = authenticate_managed_operator(session, alpha.one_time_key)
            session.commit()
            authenticated.set()
            assert revoked.wait(timeout=30)
            with pytest.raises(ManagedIdentityError) as rejected:
                _service(session, principal).rotate("beta", ManagedIdentityRotate(
                    expected_version=beta.identity.version, expires_in_days=14,
                    overlap_seconds=0, reason=REASON,
                ))
            assert rejected.value.status_code in {401, 403}
            return "rejected"

    with ThreadPoolExecutor(max_workers=1) as pool:
        pending = pool.submit(stale_request)
        assert authenticated.wait(timeout=30)
        try:
            with Session(database.runtime) as session:
                _service(session).revoke("alpha", ManagedIdentityRevoke(
                    expected_version=alpha.identity.version,
                    confirm_subject=alpha.identity.subject, reason=REASON,
                ))
        finally:
            revoked.set()
        assert pending.result(timeout=30) == "rejected"
    with Session(database.runtime) as session:
        assert authenticate_managed_operator(session, beta.one_time_key).role is PrincipalRole.ADMIN


class ManagedPostgresTools(PostgresEngine):
    """Use real psql in the managed DB container, retaining the production SQL adapter."""

    def _run(
        self, argv: list[str], target: PostgresTarget, *, stdin: str | None = None,
    ) -> CommandResult:
        container = os.getenv("MODELFORGE_TEST_POSTGRES_CONTAINER")
        if container:
            assert re.fullmatch(r"[0-9a-f]{12,64}", container)
            argv = list(argv)
            argv[argv.index("--host") + 1] = "127.0.0.1"
            argv = ["docker", "exec", "-i", "-e", "PGPASSWORD", container, *argv]
        return super()._run(argv, target, stdin=stdin)


def test_supported_restore_quarantines_old_managed_keys_until_fresh_issuance(
    identity_database: IdentityDatabase, tmp_path: Path,
) -> None:
    from modelforge_api.services.operator_identities import (
        ManagedIdentityError,
        authenticate_managed_operator,
    )
    from modelforge_api.services.recovery import RecoveryError, RecoveryService

    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    issued = _create(database.runtime, "restore-admin")
    with Session(database.runtime) as session:
        assert authenticate_managed_operator(session, issued.one_time_key).role is PrincipalRole.ADMIN
        session.commit()
        source = semantic_fingerprint(session.connection())
    archive = tmp_path / "managed-identity.dump"
    database.archive(archive)
    database.archive(archive, restore=True)
    with Session(database.restored_runtime) as session:
        assert len(list(session.scalars(select(ManagedOperatorCredentialActivation)))) == 1
        assert semantic_fingerprint(session.connection())["tables"] == source["tables"]

    # Exercise the supported recovery phase itself, with the actual restored database,
    # audit-chain checks and marker CAS; only the tool's execution location is adapted.
    restore_id = uuid.uuid4()
    owner = database.restored_owner
    with Session(database.runtime) as session:
        recovery = RecoveryService(session, Settings(
            _env_file=None, env="test",
            alembic_directory=Path(__file__).resolve().parents[1] / "alembic",
        ))
        recovery.engine = ManagedPostgresTools(base_environment=dict(os.environ), timeout_seconds=60)
        record = SimpleNamespace(id=restore_id, source_fingerprint=source)
        plan = SimpleNamespace(
            database_destination=owner.render_as_string(hide_password=False),
            backup_set=SimpleNamespace(backup_id="managed-identity-qualification",
                                       schema_revision=CURRENT, completed_at=datetime.now(UTC)),
            validation_requirements={},
        )
        with pytest.raises(RecoveryError) as unready:
            recovery._validate_restored(record, plan)
        assert unready.value.details["gate"]["current_truth_reconciled"] is False
        session.rollback()
        result = recovery._reconcile_restored(record, plan)
        session.commit()
        assert result["cleared_current_truth"]["managed_operator_credential_activations"] == 1
        ready = recovery._validate_restored(record, plan)
        assert ready["gate"]["current_truth_reconciled"] is True
        assert ready["residual_current_truth_rows"]["managed_operator_credential_activations"] == 0
        session.commit()
    with Session(database.restored_runtime) as session:
        assert list(session.scalars(select(ManagedOperatorCredentialActivation))) == []
        restored = semantic_fingerprint(session.connection())
        for table in ("managed_operator_state", "managed_operator_identities",
                      "managed_operator_credentials"):
            assert restored["tables"][table] == source["tables"][table]
        with pytest.raises(ManagedIdentityError):
            authenticate_managed_operator(session, issued.one_time_key)
        session.rollback()
        fresh = _service(session).rotate("restore-admin", ManagedIdentityRotate(
            expected_version=issued.identity.version, expires_in_days=14,
            overlap_seconds=0, reason=REASON,
        ))
    with Session(database.restored_runtime) as session:
        assert authenticate_managed_operator(session, fresh.one_time_key).role is PrincipalRole.ADMIN
        with pytest.raises(ManagedIdentityError):
            authenticate_managed_operator(session, issued.one_time_key)