"""Real PostgreSQL admission serialization with deterministic offline capacity."""

import importlib
import threading
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import Mock
from uuid import uuid4

from sqlalchemy import select
from sqlalchemy.orm import Session

from modelforge_api.domain.dispatch_context import DispatchContext
from modelforge_api.domain.external_creative import (
    ExternalAllocationCreate,
    ExternalResourceApproval,
)
from modelforge_api.persistence.models import (
    Accelerator,
    ExternalCreativeAllocation,
    GatewayRequest,
    ServiceClient,
    ServingGpuLease,
)
from modelforge_api.services import external_creative as external_creative_module
from modelforge_api.services.external_creative import ExternalCreativeService
from modelforge_api.services.manifest_registry import ManifestRegistry
from modelforge_api.services.serving import ServingError, ServingService
from modelforge_api.settings import Settings
from tests.test_managed_identities_postgres import IdentityDatabase, identity_database
from tests.test_selected_comparison_postgres import _valid_selection
from tests.test_serving import MemoryPayloadStore

__all__ = ["identity_database"]


def test_competing_external_admissions_retain_one_capacity_reservation(
    identity_database: IdentityDatabase,
) -> None:
    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    with Session(database.runtime) as session:
        _request, client_id = _valid_selection(session)
        client = session.get(ServiceClient, client_id)
        client.allowed_capabilities = ["creative.resource@1"]
        session.commit()
        accelerator_id = session.scalar(select(Accelerator.id))
        service = ExternalCreativeService(session, Mock())
        project_id = service._project(client)
        service.upsert_resource(
            ExternalResourceApproval(
                project_id=project_id,
                resource_ref="race-fixture",
                workload="image",
                provider_resource_id="offline-provider",
                configured_vram_bytes=60,
            ),
            "qualification",
        )

    barrier = threading.Barrier(2)

    def contender(index: int) -> str:
        with Session(database.runtime) as session:

            def overview(*, commit):
                assert commit is False
                retained = sum(
                    session.scalars(
                        select(ExternalCreativeAllocation.required_vram_bytes).where(
                            ExternalCreativeAllocation.state.in_(["granted", "uncertain"])
                        )
                    )
                )
                return [
                    SimpleNamespace(
                        accelerator_id=accelerator_id,
                        attribution_confidence="HIGH",
                        pressure_state="LOW",
                        schedulable_free_vram_bytes=100 - retained,
                        policy_revision="offline-fixture",
                        observed_at=None,
                    )
                ]

            serving = Mock()
            serving.scheduler_overview.side_effect = overview
            service = ExternalCreativeService(session, serving)
            client = session.get(ServiceClient, client_id)
            session.commit()
            barrier.wait(timeout=30)
            return service.acquire(
                ExternalAllocationCreate(
                    request_id=f"race-{index}",
                    idempotency_key=f"race-{index}",
                    installation_id="fixture",
                    workload="image",
                    resource_ref="race-fixture",
                    provider_resource_id="offline-provider",
                    required_vram_bytes=60,
                ),
                client,
            ).state

    with ThreadPoolExecutor(max_workers=2) as pool:
        assert sorted(pool.map(contender, [1, 2])) == ["granted", "queued"]
    with Session(database.runtime) as session:
        rows = list(session.scalars(select(ExternalCreativeAllocation)))
        assert len(rows) == 2
        assert sum(row.required_vram_bytes for row in rows if row.state == "granted") == 60


def test_native_and_external_admission_share_actual_postgres_capacity(
    identity_database: IdentityDatabase,
) -> None:
    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    settings = Settings(
        _env_file=None,
        env="test",
        gateway_queue_timeout_seconds=1,
        gateway_request_timeout_seconds=60,
    )

    def serving(session):
        return ServingService(session, settings, ManifestRegistry(), MemoryPayloadStore())

    with Session(database.runtime) as session:
        selection, client_id = _valid_selection(session)
        client = session.get(ServiceClient, client_id)
        client.allowed_capabilities = ["rag.embedding@1", "creative.resource@1"]
        session.commit()
        service = serving(session)
        deployment_id = selection.current_deployment_id
        deployment = service.repo.deployment(deployment_id)
        baseline = service._budget(deployment)
        plan = service._placement_plan(
            deployment, priority=client.workload_priority, dry_run=True, persist=False, commit=False
        )
        assert plan.verdict == "ADMIT"
        available, native_required = baseline.schedulable_free_vram_bytes, plan.required_vram_bytes
        assert 0 < native_required < available
        external_required = available - native_required // 2
        assert external_required <= available < external_required + native_required
        external = ExternalCreativeService(session, service)
        external.upsert_resource(
            ExternalResourceApproval(
                project_id=external._project(client),
                resource_ref="mixed-race",
                workload="image",
                provider_resource_id="offline-provider",
                configured_vram_bytes=external_required,
            ),
            "qualification",
        )
        gateway = GatewayRequest(
            request_id=uuid4(),
            service_client_id=client.id,
            capability_key="rag.embedding",
            capability_version=1,
            capability_deployment_id=deployment.id,
            compute_node_id=deployment.compute_node_id,
            status="queued",
            priority=client.workload_priority,
            input_sha256="a" * 64,
            input_count=1,
        )
        session.add(gateway)
        session.commit()
        gateway_id = gateway.id

    barrier = threading.Barrier(2)

    def contender(native):
        with Session(database.runtime) as session:
            service = serving(session)
            client = session.get(ServiceClient, client_id)
            deployment = service.repo.deployment(deployment_id)
            gateway = session.get(GatewayRequest, gateway_id)
            context = DispatchContext.for_service_client(
                client.id, client.name, client.workload_priority
            )
            session.commit()
            barrier.wait(timeout=30)
            if native:
                try:
                    service._acquire_request_lease(deployment, gateway, context)
                    return "native_granted"
                except ServingError as error:
                    assert error.code in {
                        "REQUEST_TIMEOUT",
                        "INSUFFICIENT_SCHEDULABLE_VRAM",
                        "INSUFFICIENT_CAPACITY",
                    }, error.code
                    return "native_blocked"
            result = ExternalCreativeService(session, service).acquire(
                ExternalAllocationCreate(
                    request_id="mixed-race",
                    idempotency_key="mixed-race",
                    installation_id="fixture",
                    resource_ref="mixed-race",
                    workload="image",
                    provider_resource_id="offline-provider",
                    required_vram_bytes=external_required,
                    ttl_seconds=60,
                ),
                client,
            )
            return "external_" + result.state

    with ThreadPoolExecutor(max_workers=2) as pool:
        outcomes = list(pool.map(contender, [True, False]))
    assert outcomes in (
        ["native_granted", "external_queued"],
        ["native_blocked", "external_granted"],
    )
    with Session(database.runtime) as session:
        native = list(
            session.scalars(
                select(ServingGpuLease).where(ServingGpuLease.state.in_(["active", "granted"]))
            )
        )
        external = list(
            session.scalars(
                select(ExternalCreativeAllocation).where(
                    ExternalCreativeAllocation.state.in_(["granted", "uncertain"])
                )
            )
        )
        assert len(native) + len(external) == 1
        assert (
            sum(row.reserved_vram_bytes for row in native)
            + sum(row.required_vram_bytes for row in external)
            <= available
        )


def test_expired_external_listing_keeps_postgres_count_and_rows_in_sync(
    identity_database: IdentityDatabase, monkeypatch
) -> None:
    database = identity_database
    importlib.import_module("alembic.command").upgrade(database.migration(), "head")
    with Session(database.runtime) as session:
        _request, client_id = _valid_selection(session)
        client = session.get(ServiceClient, client_id)
        client.allowed_capabilities = ["creative.resource@1"]
        session.commit()

        accelerator_id = session.scalar(select(Accelerator.id))
        serving = Mock()
        serving.scheduler_overview.return_value = [
            SimpleNamespace(
                accelerator_id=accelerator_id,
                attribution_confidence="HIGH",
                pressure_state="LOW",
                schedulable_free_vram_bytes=100,
                policy_revision="postgres-boundary-fixture",
                observed_at=None,
            )
        ]
        service = ExternalCreativeService(session, serving)
        project_id = service._project(client)
        service.upsert_resource(
            ExternalResourceApproval(
                project_id=project_id,
                resource_ref="expiry-boundary-fixture",
                workload="image",
                provider_resource_id="offline-provider",
                configured_vram_bytes=60,
            ),
            "qualification",
        )
        acquired = service.acquire(
            ExternalAllocationCreate(
                request_id="expiry-boundary",
                idempotency_key="expiry-boundary",
                installation_id="fixture",
                workload="image",
                resource_ref="expiry-boundary-fixture",
                provider_resource_id="offline-provider",
                required_vram_bytes=60,
            ),
            client,
        )
        row = session.get(ExternalCreativeAllocation, acquired.allocation_id)
        boundary = datetime(2030, 1, 1, tzinfo=UTC)
        row.expires_at = boundary
        session.commit()
        monkeypatch.setattr(external_creative_module, "_now", lambda: boundary)

        uncertain = service.list_allocations("uncertain", 20, 0)
        granted = service.list_allocations("granted", 20, 0)

        assert uncertain.total == len(uncertain.items) == 1
        assert uncertain.items[0].allocation_id == acquired.allocation_id
        assert uncertain.items[0].state == "uncertain"
        assert uncertain.items[0].decision_code == "ALLOCATION_EXPIRED_UNCERTAIN"
        assert uncertain.items[0].uncertain_at is None
        assert granted.total == len(granted.items) == 0
        session.refresh(row)
        assert row.state == "granted" and row.expires_at == boundary