import hashlib
import json
from collections.abc import Iterator
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import create_engine, func, select
from sqlalchemy.orm import Session

from modelforge_api.domain.lab import LabExecutionStartRequest
from modelforge_api.domain.serving import (
    AgentServingJobCancellationAck,
    AgentServingJobCancellationCheck,
    AgentServingJobComplete,
)
from modelforge_api.persistence.models import (
    Base,
    GatewayRequest,
    LabExecutionRun,
    ServingJob,
)
from modelforge_api.services.lab_runs import LabRunService
from modelforge_api.services.serving import ServingError
from tests.test_dispatch_context import operator_fixture


@pytest.fixture
def session() -> Iterator[Session]:
    engine = create_engine("sqlite+pysqlite:///:memory:")
    Base.metadata.create_all(engine)
    with Session(engine) as value:
        yield value
    engine.dispose()


def make_start_request(request, selected) -> LabExecutionStartRequest:
    return LabExecutionStartRequest(
        target=request.target,
        input=request.input,
        limits=request.limits,
        deployment_id=selected.capability_deployment_id,
        experiment_route_id=selected.id,
    )


def test_same_actor_key_and_fingerprint_returns_one_durable_run(session: Session) -> None:
    serving, preflight, selected = operator_fixture(session)
    service = LabRunService(serving)
    request = make_start_request(preflight, selected)

    first = service.start(request, actor_subject="operator-test", idempotency_key="run-001")
    retry = service.start(request, actor_subject="operator-test", idempotency_key="run-001")

    assert retry.id == first.id
    assert retry.status == "accepted"
    assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 1
    assert "private input content" not in first.model_dump_json()
    assert "payload_reference" not in first.model_dump()


def test_reused_actor_key_with_different_input_is_a_conflict(session: Session) -> None:
    serving, preflight, selected = operator_fixture(session)
    service = LabRunService(serving)
    request = make_start_request(preflight, selected)
    changed = LabExecutionStartRequest.model_validate(
        {
            **request.model_dump(mode="json"),
            "input": {"capability": "rag.embedding", "payload": {"input": "different"}},
        }
    )
    service.start(request, actor_subject="operator-test", idempotency_key="run-002")

    with pytest.raises(ServingError) as raised:
        service.start(changed, actor_subject="operator-test", idempotency_key="run-002")

    assert raised.value.status_code == 409
    assert raised.value.code == "LAB_IDEMPOTENCY_CONFLICT"
    assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 1


def test_run_status_and_cancel_are_scoped_to_authenticated_actor(session: Session) -> None:
    serving, preflight, selected = operator_fixture(session)
    service = LabRunService(serving)
    run = service.start(
        make_start_request(preflight, selected),
        actor_subject="operator-owner",
        idempotency_key="run-003",
    )

    for action in (
        lambda: service.get(run.id, actor_subject="operator-other"),
        lambda: service.cancel(run.id, actor_subject="operator-other"),
    ):
        with pytest.raises(ServingError) as raised:
            action()
        assert raised.value.status_code == 404
        assert raised.value.code == "LAB_RUN_NOT_FOUND"

    assert service.get(run.id, actor_subject="operator-owner").id == run.id
    cancelled = service.cancel(run.id, actor_subject="operator-owner")
    assert cancelled.status == "cancelled"
    assert service.cancel(run.id, actor_subject="operator-owner").status == "cancelled"


def _claimed_run(session: Session, actor: str, key: str):
    serving, preflight, selected = operator_fixture(session)
    run_service = LabRunService(serving)
    run_response = run_service.start(
        make_start_request(preflight, selected), actor_subject=actor, idempotency_key=key
    )
    run = session.get(LabExecutionRun, run_response.id)
    assert run is not None
    gateway = session.get(GatewayRequest, run.gateway_request_id)
    assert gateway is not None
    token = "t" * 32
    job = ServingJob(
        capability_deployment_id=selected.capability_deployment_id,
        compute_node_id=selected.deployment.compute_node_id,
        gateway_request_id=gateway.id,
        lab_execution_run_id=run.id,
        operation="invoke",
        status="running",
        priority="lab",
        idempotency_key=f"test-{key}",
        payload_reference=run.payload_reference,
        attempt_count=1,
        lease_token_hash=hashlib.sha256(token.encode()).hexdigest(),
        lease_expires_at=datetime.now(UTC) + timedelta(minutes=2),
        result_summary={},
        started_at=datetime.now(UTC),
    )
    session.add(job)
    run.status = "running"
    run.started_at = datetime.now(UTC)
    run.heartbeat_at = datetime.now(UTC)
    gateway.status = "running"
    session.commit()
    return serving, run_service, run, job, selected, token


def _queued_run_with_job(session: Session, actor: str, key: str):
    serving, preflight, selected = operator_fixture(session)
    run_service = LabRunService(serving)
    response = run_service.start(
        make_start_request(preflight, selected), actor_subject=actor, idempotency_key=key
    )
    run = session.get(LabExecutionRun, response.id)
    assert run is not None
    gateway = session.get(GatewayRequest, run.gateway_request_id)
    assert gateway is not None
    job = ServingJob(
        capability_deployment_id=selected.capability_deployment_id,
        compute_node_id=selected.deployment.compute_node_id,
        gateway_request_id=gateway.id,
        lab_execution_run_id=run.id,
        operation="invoke",
        status="queued",
        priority="lab",
        idempotency_key=f"queued-{key}",
        payload_reference=run.payload_reference,
        attempt_count=0,
        result_summary={},
    )
    run.status = "queued"
    gateway.status = "queued"
    session.add(job)
    session.commit()
    return serving, run_service, run, job, selected, preflight


def test_cancel_racing_with_worker_completion_discards_vectors_and_acknowledges(
    session: Session,
) -> None:
    serving, run_service, run, job, selected, token = _claimed_run(
        session, "operator-race", "run-race"
    )
    requested = run_service.cancel(run.id, actor_subject="operator-race")
    assert requested.status == "cancel_requested"

    deployment = serving.repo.deployment(selected.capability_deployment_id)
    assert deployment is not None
    node = serving.repo.node(deployment.compute_node_id)
    space = serving.repo.embedding_space(deployment.embedding_space_id)
    assert node is not None and space is not None
    vector = [1.0] + [0.0] * (space.dimension - 1)
    serving.complete_job(
        job.id,
        node,
        AgentServingJobComplete(
            lease_token=token,
            worker_instance_id="worker-cancel-race",
            state="warm",
            result={"vectors": [vector], "count": 1, "normalized": True},
            metrics={"inference_time_ms": 12.0, "private_metric": "must-not-persist"},
            health={},
            runtime_facts={"private_runtime_fact": "must-not-persist"},
        ),
    )

    session.refresh(run)
    session.refresh(job)
    assert run.status == "cancelled"
    assert run.cancellation_acknowledged_at is not None
    assert run.output_metadata is None
    assert job.status == "cancelled"
    assert job.result_reference is None
    persisted = f"{run.worker_observations!r} {job.result_summary!r}"
    assert "private_metric" not in persisted
    assert "private_runtime_fact" not in persisted
    assert "must-not-persist" not in persisted
    assert serving.payload_store.get(run.payload_reference) is None


def test_successful_lab_completion_persists_hash_and_allowlisted_metadata_only(
    session: Session,
) -> None:
    serving, _run_service, run, job, selected, token = _claimed_run(
        session, "operator-success", "run-success"
    )
    deployment = serving.repo.deployment(selected.capability_deployment_id)
    assert deployment is not None
    node = serving.repo.node(deployment.compute_node_id)
    space = serving.repo.embedding_space(deployment.embedding_space_id)
    assert node is not None and space is not None
    vector = [1.0] + [0.0] * (space.dimension - 1)
    serving.complete_job(
        job.id,
        node,
        AgentServingJobComplete(
            lease_token=token,
            worker_instance_id="worker-success",
            state="warm",
            result={"vectors": [vector], "count": 1, "normalized": True},
            metrics={
                "inference_time_ms": 12.0,
                "gpu_uuid": "GPU-test",
                "private_metric": "must-not-persist",
            },
            health={},
            runtime_facts={"private_runtime_fact": "must-not-persist"},
        ),
    )

    session.refresh(run)
    session.refresh(job)
    assert run.status == "completed"
    assert run.output_metadata is not None
    assert run.output_metadata["dimension"] == space.dimension
    assert run.output_metadata["item_count"] == 1
    assert len(run.output_metadata["sha256"]) == 64
    assert (
        run.output_metadata["sha256"]
        == hashlib.sha256(
            json.dumps([vector], sort_keys=True, separators=(",", ":")).encode()
        ).hexdigest()
    )
    assert run.output_metadata["content_logged"] is False
    assert run.worker_observations is not None
    assert run.worker_observations["source"] == "runtime-worker"
    assert "private_metric" not in str(run.worker_observations)
    assert "private_runtime_fact" not in str(run.worker_observations)
    assert job.status == "completed"
    assert job.result_reference is None
    assert "vectors" not in str(job.result_summary)
    gateway = session.get(GatewayRequest, run.gateway_request_id)
    assert gateway is not None
    persisted_evidence = gateway.decision_evidence["lab_run_evidence"]
    assert persisted_evidence["status"] == "completed"
    assert persisted_evidence["output"] == run.output_metadata
    assert "private_metric" not in str(gateway.decision_evidence)
    assert "private_runtime_fact" not in str(gateway.decision_evidence)
    assert "vectors" not in str(gateway.decision_evidence)
    assert "private input content" not in str(gateway.decision_evidence)
    assert serving.payload_store.get(run.payload_reference) is None


def test_reconciler_dispatches_exact_cold_target_through_load_and_invoke_jobs(
    session: Session,
) -> None:
    serving, preflight, selected = operator_fixture(session)
    run_service = LabRunService(serving)
    created = run_service.start(
        make_start_request(preflight, selected),
        actor_subject="operator-reconcile",
        idempotency_key="run-reconcile",
    )

    serving.reconcile()
    run = session.get(LabExecutionRun, created.id)
    assert run is not None
    assert run.status == "preparing"
    load_job = session.scalar(
        select(ServingJob).where(
            ServingJob.lab_execution_run_id == run.id,
            ServingJob.operation == "load",
        )
    )
    assert load_job is not None and load_job.status == "queued"
    node = serving.repo.node(selected.deployment.compute_node_id)
    assert node is not None
    load_lease = serving.claim_next(node)
    assert load_lease is not None and load_lease.job_id == load_job.id
    assert load_lease.lab_run_id == run.id
    serving.complete_job(
        load_job.id,
        node,
        AgentServingJobComplete(
            lease_token=load_lease.lease_token,
            worker_instance_id="synthetic-worker",
            state="warm",
            result={},
            metrics={
                "baseline_vram_bytes": 1_000,
                "loaded_vram_bytes": 3_000,
            },
            health={"model": "healthy"},
            runtime_facts={"offline_local_only": True},
        ),
    )

    serving.reconcile()
    session.refresh(run)
    assert run.status == "queued"
    invoke_job = session.scalar(
        select(ServingJob).where(
            ServingJob.lab_execution_run_id == run.id,
            ServingJob.operation == "invoke",
        )
    )
    assert invoke_job is not None and invoke_job.status == "queued"
    invoke_lease = serving.claim_next(node)
    assert invoke_lease is not None and invoke_lease.job_id == invoke_job.id
    assert invoke_lease.deployment_id == selected.capability_deployment_id
    assert invoke_lease.lab_run_id == run.id
    assert invoke_lease.input == ["private input content"]
    assert invoke_lease.expected_dimension is not None
    deployment = serving.repo.deployment(selected.capability_deployment_id)
    assert deployment is not None
    accelerator = serving.repo.accelerator(deployment.accelerator_id)
    assert accelerator is not None
    vector = [1.0] + [0.0] * (invoke_lease.expected_dimension - 1)
    serving.complete_job(
        invoke_job.id,
        node,
        AgentServingJobComplete(
            lease_token=invoke_lease.lease_token,
            worker_instance_id="synthetic-worker",
            state="warm",
            result={"vectors": [vector], "count": 1, "normalized": True},
            metrics={
                "inference_time_ms": 12.0,
                "gpu_uuid": accelerator.device_uuid,
                "used_vram_bytes": 3_000,
                "process_used_vram_bytes": 2_000,
                "free_vram_bytes": 7_000,
                "gpu_utilization_percent": 35,
                "temperature_c": 60,
                "power_draw_w": 90,
                "private_metric": "must-not-persist",
            },
            health={"model": "healthy"},
            runtime_facts={"private_runtime_fact": "must-not-persist"},
        ),
    )

    session.refresh(run)
    session.refresh(invoke_job)
    assert run.status == "completed"
    assert run.output_metadata is not None
    assert run.output_metadata["item_count"] == 1
    assert len(run.output_metadata["sha256"]) == 64
    assert run.worker_observations is not None
    assert run.worker_observations["device_matches_target"] is True
    assert "private_metric" not in str(run.worker_observations)
    assert "private_runtime_fact" not in str(run.worker_observations)
    assert invoke_job.result_reference is None
    assert "vectors" not in str(invoke_job.result_summary)
    assert (
        "private input content"
        not in run_service.get(run.id, actor_subject="operator-reconcile").model_dump_json()
    )


def test_expired_claim_is_interrupted_and_never_reclaimed(session: Session) -> None:
    serving, _run_service, run, job, selected, _token = _claimed_run(
        session, "operator-stale", "run-stale"
    )
    job.lease_expires_at = datetime.now(UTC) - timedelta(seconds=1)
    session.commit()

    serving.reconcile()

    session.refresh(run)
    session.refresh(job)
    assert run.status == "interrupted"
    assert run.failure_code == "LAB_WORKER_LEASE_STALE"
    assert job.status == "failed"
    assert job.failure_code == "LAB_WORKER_LEASE_STALE"
    node = serving.repo.node(selected.deployment.compute_node_id)
    assert node is not None
    assert serving.claim_next(node) is None


@pytest.mark.parametrize("expiry_mode", ["redis_key_missing", "durable_expiry_elapsed"])
def test_expired_or_missing_redis_input_fails_closed_without_dispatch(
    session: Session, expiry_mode: str
) -> None:
    serving, preflight, selected = operator_fixture(session)
    run_service = LabRunService(serving)
    created = run_service.start(
        make_start_request(preflight, selected),
        actor_subject=f"operator-expired-{expiry_mode}",
        idempotency_key=f"expired-{expiry_mode}",
    )
    run = session.get(LabExecutionRun, created.id)
    assert run is not None and run.payload_reference is not None
    if expiry_mode == "redis_key_missing":
        # Redis TTL expiry is represented by the payload store returning no key.
        serving.payload_store.delete(run.payload_reference)
    else:
        run.input_payload_expires_at = datetime.now(UTC) - timedelta(seconds=1)
        session.commit()

    serving.reconcile()

    session.refresh(run)
    gateway = session.get(GatewayRequest, run.gateway_request_id)
    assert gateway is not None
    assert run.status == "failed"
    assert run.failure_code == "REQUEST_PAYLOAD_EXPIRED"
    assert run.payload_reference is None
    assert gateway.status == "failed"
    assert session.scalar(
        select(func.count()).select_from(ServingJob).where(
            ServingJob.lab_execution_run_id == run.id
        )
    ) == 0


def test_claim_with_missing_redis_key_fails_run_instead_of_leaving_it_queued(
    session: Session,
) -> None:
    serving, _run_service, run, job, selected, _preflight = _queued_run_with_job(
        session, "operator-claim-expired", "claim-expired"
    )
    assert run.payload_reference is not None
    serving.payload_store.delete(run.payload_reference)
    node = serving.repo.node(selected.deployment.compute_node_id)
    assert node is not None

    lease = serving.claim_next(node)

    session.refresh(run)
    session.refresh(job)
    assert lease is None
    assert run.status == "failed"
    assert run.failure_code == "REQUEST_PAYLOAD_EXPIRED"
    assert job.status == "failed"
    assert job.failure_code == "REQUEST_PAYLOAD_EXPIRED"


def test_lab_claim_rejects_another_runs_live_redis_payload(session: Session) -> None:
    serving, run_service, run, job, selected, preflight = _queued_run_with_job(
        session, "operator-payload-isolation", "payload-isolation"
    )
    changed_request = LabExecutionStartRequest.model_validate(
        {
            **make_start_request(preflight, selected).model_dump(mode="json"),
            "input": {
                "capability": "rag.embedding",
                "payload": {"input": "another run's private input"},
            },
        }
    )
    other = run_service.start(
        changed_request,
        actor_subject="operator-payload-isolation",
        idempotency_key="payload-isolation-other",
    )
    assert other.id != run.id
    assert other.input_summary.sha256 != run.input_sha256
    other_run_payload = session.get(LabExecutionRun, other.id)
    assert other_run_payload is not None and other_run_payload.payload_reference is not None
    job.payload_reference = other_run_payload.payload_reference
    session.commit()
    node = serving.repo.node(selected.deployment.compute_node_id)
    assert node is not None

    lease = serving.claim_next(node)

    session.refresh(run)
    session.refresh(job)
    session.refresh(other_run_payload)
    assert lease is None
    assert run.status == "failed"
    assert run.failure_code == "LAB_PAYLOAD_REFERENCE_MISMATCH"
    assert job.failure_code == "LAB_PAYLOAD_REFERENCE_MISMATCH"
    assert other_run_payload.status == "accepted"


def test_worker_poll_renews_lease_and_ack_completes_cancel_transition(session: Session) -> None:
    serving, run_service, run, job, selected, token = _claimed_run(
        session, "operator-ack", "run-ack"
    )
    run.timeout_seconds = 300
    job.lease_expires_at = datetime.now(UTC) + timedelta(seconds=30)
    session.commit()
    node = serving.repo.node(selected.deployment.compute_node_id)
    assert node is not None
    before = job.lease_expires_at
    requested = run_service.cancel(run.id, actor_subject="operator-ack")
    assert requested.status == "cancel_requested"

    state = serving.check_lab_cancellation(
        job.id, node, AgentServingJobCancellationCheck(lease_token=token)
    )
    assert state.cancel_requested is True
    session.refresh(job)
    assert job.lease_expires_at is not None and job.lease_expires_at >= before
    serving.acknowledge_lab_cancellation(
        job.id,
        node,
        AgentServingJobCancellationAck(
            lease_token=token,
            worker_instance_id="worker-ack",
            state="warm",
        ),
    )

    session.refresh(run)
    session.refresh(job)
    assert run.status == "cancelled"
    assert run.cancellation_acknowledged_at is not None
    assert run.worker_observations == {
        "source": "runtime-worker",
        "phase": "after_inference",
        "availability": "cancelled",
        "safe_boundary_reached": True,
        "worker_instance_id": "worker-ack",
        "worker_state": "warm",
    }
    assert job.status == "cancelled"
    assert job.result_reference is None


def test_queued_run_cancel_prevents_worker_claim_and_deletes_ttl_input(session: Session) -> None:
    serving, preflight, selected = operator_fixture(session)
    run_service = LabRunService(serving)
    response = run_service.start(
        make_start_request(preflight, selected),
        actor_subject="operator-queued-cancel",
        idempotency_key="queued-cancel",
    )
    run = session.get(LabExecutionRun, response.id)
    assert run is not None
    payload_reference = run.payload_reference
    gateway = session.get(GatewayRequest, run.gateway_request_id)
    assert gateway is not None
    job = ServingJob(
        capability_deployment_id=selected.capability_deployment_id,
        compute_node_id=selected.deployment.compute_node_id,
        gateway_request_id=gateway.id,
        lab_execution_run_id=run.id,
        operation="invoke",
        status="queued",
        priority="lab",
        idempotency_key="queued-cancel-job",
        payload_reference=run.payload_reference,
        attempt_count=0,
        result_summary={},
    )
    session.add(job)
    run.status = "queued"
    gateway.status = "queued"
    session.commit()

    cancelled = run_service.cancel(run.id, actor_subject="operator-queued-cancel")

    session.refresh(job)
    assert cancelled.status == "cancelled"
    assert job.status == "cancelled"
    assert job.lease_token_hash is None
    assert "input_payload_expires_at" in cancelled.model_dump()
    assert payload_reference is not None
    assert serving.payload_store.get(payload_reference) is None