from __future__ import annotations

from datetime import UTC, datetime
from uuid import uuid4

from wf_api.saved_subgraphs import SavedSubgraphTree
from wf_artifacts import (
    AvailableSource,
    CheckpointReason,
    DependencyDiagnostic,
    DiagnosticSeverity,
    PinnedRunEnvironment,
    ResumeReadiness,
    RunAdmission,
    RunCheckpoint,
    RunStore,
    StoredRunStatus,
    WorkflowArtifact,
    WorkflowDeployment,
    WorkflowRunRecord,
    validate_deployment_dependencies,
)
from wf_core import (
    PersistedRunState,
    RunState,
    RunStatus,
    dump_run_state,
    load_run_state,
    load_run_state_with_upgrade,
)


def create_pinned_environment(
    *,
    deployment: WorkflowDeployment,
    artifact: WorkflowArtifact,
    tree: SavedSubgraphTree,
) -> PinnedRunEnvironment:
    """Capture exact root, deployment, and child definitions for one run."""
    return PinnedRunEnvironment(
        deployment=deployment,
        root_artifact=artifact,
        child_artifacts=list(tree.artifacts_by_ref.values()),
    )


def persist_admission(
    *,
    store: RunStore,
    run_id: str,
    environment: PinnedRunEnvironment,
    resolved_input: dict[str, object],
    max_steps: int | None,
    scheduled_at: datetime | None = None,
    schedule_id: str | None = None,
    schedule_revision: int | None = None,
    deployment_revision: int | None = None,
) -> RunAdmission:
    """Persist the authoritative admission record before dispatch.

    The occurrence is decided here: the caller must treat this persist as the
    single authority for the run identity and frozen invocation. A failed
    durable admission must never dispatch.
    """
    from wf_core.models.json_values import validate_strict_json_value

    frozen = validate_strict_json_value(dict(resolved_input))
    if not isinstance(frozen, dict):
        raise ValueError("resolved workflow input must be a JSON object")
    admission = RunAdmission(
        id=run_id,
        environment=environment,
        resolved_input=frozen,
        max_steps=max_steps,
        scheduled_at=scheduled_at,
        schedule_id=schedule_id,
        schedule_revision=schedule_revision,
        deployment_revision=deployment_revision,
        created_at=datetime.now(UTC),
    )
    store.save_admission(admission)
    return admission


def materialize_admitted_view(
    *,
    store: RunStore,
    admission: RunAdmission,
) -> WorkflowRunRecord:
    """Materialize the run admission view using the admission identity.

    The view carries no checkpoint, trace, output, or step counts: the outcome
    is unknown until dispatch completes and persists a stopped checkpoint.
    """
    now = datetime.now(UTC)
    try:
        existing = store.get_run(admission.id)
    except KeyError:
        existing = None
    if existing is not None:
        return existing
    record = WorkflowRunRecord(
        id=admission.id,
        status=StoredRunStatus.ADMITTED,
        resume_readiness=ResumeReadiness.NOT_APPLICABLE,
        environment=admission.environment,
        latest_checkpoint_id=None,
        created_at=admission.created_at,
        updated_at=now,
    )
    store.save_run(record)
    return record


def recover_admission_view(*, store: RunStore, run_id: str) -> WorkflowRunRecord:
    """Reconcile a missing run view from its admission record.

    Recovery never executes work: it only completes the missing view so a
    later poll can dispatch the captured invocation exactly once.
    """
    try:
        return store.get_run(run_id)
    except KeyError:
        pass
    admission = store.get_admission(run_id)
    return materialize_admitted_view(store=store, admission=admission)


def persist_stopped_run(
    *,
    store: RunStore,
    environment: PinnedRunEnvironment,
    run: RunState,
    run_id: str | None = None,
) -> WorkflowRunRecord:
    """Persist one externally visible stopped state and its typed checkpoint."""
    if run.status not in {
        RunStatus.INTERRUPTED,
        RunStatus.COMPLETED,
        RunStatus.FAILED,
    }:
        raise ValueError(
            f"cannot persist active workflow run with status {run.status!s}"
        )

    key = run_id or f"run_{uuid4().hex}"
    now = datetime.now(UTC)
    sequence = 1
    created_at = now
    if run_id is not None:
        try:
            existing = store.get_run(run_id)
        except KeyError:
            existing = None
        if existing is not None:
            created_at = existing.created_at
            try:
                sequence = store.get_latest_checkpoint(run_id).sequence + 1
            except KeyError:
                sequence = 1

    status = StoredRunStatus(run.status.value)
    readiness = (
        ResumeReadiness.READY
        if status is StoredRunStatus.INTERRUPTED
        else ResumeReadiness.NOT_APPLICABLE
    )
    checkpoint_id = f"{key}.{sequence:06d}"
    checkpoint = RunCheckpoint(
        id=checkpoint_id,
        run_id=key,
        sequence=sequence,
        reason=CheckpointReason(status.value),
        state=PersistedRunState.model_validate(dump_run_state(run)),
        created_at=now,
    )
    record = WorkflowRunRecord(
        id=key,
        status=status,
        resume_readiness=readiness,
        environment=environment,
        latest_checkpoint_id=checkpoint_id,
        created_at=created_at,
        updated_at=now,
    )
    store.save_checkpoint(checkpoint)
    store.save_run(record)
    return record


def restore_interrupted_run(
    store: RunStore, run_id: str
) -> tuple[WorkflowRunRecord, RunState]:
    """Load a persisted interrupted run, persisting a v1 upgrade first.

    A pre-budget (v1) checkpoint receives its one-time defaults and is
    rewritten as a new v2 interrupted checkpoint under the same run id and
    pinned environment *before* the run is returned, so resume dispatch
    never runs on unmigrated state and a failed upgrade fails resume before
    any handler runs. Ordinary inspection uses :func:`load_stored_run`,
    which decodes v1 prospectively without mutating the store.
    """
    record = store.get_run(run_id)
    if record.status is not StoredRunStatus.INTERRUPTED:
        raise ValueError(f"workflow run {run_id!r} is not interrupted")
    checkpoint = store.get_latest_checkpoint(run_id)
    run, upgraded = load_run_state_with_upgrade(
        checkpoint.state.model_dump(mode="json")
    )
    if upgraded:
        record = persist_stopped_run(
            store=store,
            environment=record.environment,
            run=run,
            run_id=run_id,
        )
    return record, run


def load_stored_run(store: RunStore, run_id: str) -> tuple[WorkflowRunRecord, RunState]:
    """Load any stopped run record together with its latest typed checkpoint."""
    record = store.get_run(run_id)
    if record.status is StoredRunStatus.ADMITTED:
        raise ValueError(
            f"workflow run {run_id!r} is admitted but has no stopped checkpoint yet"
        )
    checkpoint = store.get_latest_checkpoint(run_id)
    return record, load_run_state(checkpoint.state.model_dump(mode="json"))


def validate_pinned_resume_environment(
    *,
    record: WorkflowRunRecord,
    sources: list[AvailableSource],
) -> list[DependencyDiagnostic]:
    """Revalidate exact stored graph definitions before a resume mutates state."""
    environment = record.environment
    diagnostics = validate_deployment_dependencies(
        artifact=environment.root_artifact,
        deployment=environment.deployment,
        sources=sources,
    )
    for child in environment.child_artifacts:
        diagnostics.extend(
            validate_deployment_dependencies(
                artifact=child,
                deployment=environment.deployment,
                sources=sources,
            )
        )
    return diagnostics


def has_blocking_diagnostics(diagnostics: list[DependencyDiagnostic]) -> bool:
    """Return whether dependency diagnostics prohibit executing a resume."""
    return any(item.severity is DiagnosticSeverity.ERROR for item in diagnostics)


def mark_resume_blocked(
    *,
    store: RunStore,
    record: WorkflowRunRecord,
    diagnostics: list[DependencyDiagnostic],
) -> WorkflowRunRecord:
    """Record blocked readiness without writing a new execution checkpoint."""
    blocked = record.model_copy(
        update={
            "resume_readiness": ResumeReadiness.BLOCKED,
            "diagnostics": diagnostics,
            "updated_at": datetime.now(UTC),
        }
    )
    store.save_run(blocked)
    return blocked