"""Native admission honors external reservations until their explicit release."""
from datetime import UTC, datetime, timedelta
from uuid import uuid4

import pytest
from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session

from modelforge_api.persistence.models import (
    Base,
    ExternalCreativeAllocation,
    ProjectBinding,
    ServiceClient,
)
from modelforge_api.services.manifest_registry import ManifestRegistry
from modelforge_api.services.serving import ServingService
from modelforge_api.settings import Settings
from tests.test_selected_comparison_postgres import _valid_selection
from tests.test_serving import MemoryPayloadStore


@pytest.fixture
def budget_fixture():
    engine = create_engine("sqlite://")
    event.listen(engine, "connect", lambda connection, _: connection.execute("PRAGMA foreign_keys=ON"))
    Base.metadata.create_all(engine)
    with Session(engine) as session:
        selection, client_id = _valid_selection(session)
        service = ServingService(session, Settings(_env_file=None, env="test"),
                                 ManifestRegistry(), MemoryPayloadStore())
        deployment = service.repo.deployment(selection.current_deployment_id)
        client = session.get(ServiceClient, client_id)
        assert deployment is not None and client is not None
        binding = session.get(ProjectBinding, client.project_binding_id)
        assert binding is not None
        yield session, service, deployment, (client.id, binding.project_id)
    engine.dispose()


def allocation(deployment, client, *, state, required):
    return ExternalCreativeAllocation(
        service_client_id=client[0], project_id=client[1],
        accelerator_id=deployment.accelerator_id, request_id=str(uuid4()),
        idempotency_key=str(uuid4()), request_sha256="a" * 64, workload="image",
        resource_ref="fixture", provider_resource_id="fixture", required_vram_bytes=required,
        priority="background", ttl_seconds=30, state=state, decision_code="FIXTURE",
        decision_evidence={}, required_actions=[],
        expires_at=datetime.now(UTC) - timedelta(seconds=1),
    )


@pytest.mark.parametrize(("state", "reserved"), [
    ("granted", True), ("uncertain", True), ("queued", False), ("released", False),
])
def test_actual_native_budget_retains_expired_grants_and_uncertain_reservations(
    budget_fixture, state, reserved,
):
    session, service, deployment, client = budget_fixture
    before = service._budget(deployment)
    required = 512 * 1024**2
    session.add(allocation(deployment, client, state=state, required=required))
    session.flush()
    after = service._budget(deployment)
    delta = required if reserved else 0
    assert after.leased_vram_bytes == before.leased_vram_bytes + delta
    assert after.schedulable_free_vram_bytes == before.schedulable_free_vram_bytes - delta
    assert after.invariant_delta_bytes == before.invariant_delta_bytes == 0


def test_native_placement_cannot_spend_external_reservation(budget_fixture):
    session, service, deployment, client = budget_fixture
    before = service._budget(deployment)
    session.add(allocation(deployment, client, state="uncertain",
                           required=before.schedulable_free_vram_bytes - 1))
    session.flush()
    plan = service._placement_plan(deployment, priority="production", dry_run=True, commit=False)
    assert plan.headroom_before_bytes == 1
    assert plan.verdict not in {"ADMIT", "ADMIT_AFTER_EVICTION"}


def test_overview_does_not_commit_owning_admission_transaction(budget_fixture):
    session, service, _deployment, _client = budget_fixture
    commits = []
    event.listen(session, "after_commit", lambda _: commits.append(True))
    assert service.scheduler_overview(commit=False)
    assert commits == []
    assert session.in_transaction()
    service.scheduler_overview()
    assert commits == [True]