from __future__ import annotations

from typing import Any

from one_run.constants import JSON_SCHEMA_VERSION
from one_run.models import RunMetadata, RunRuntime, RunStatus


def run_result_payload(
    *,
    run_id: str,
    state: str,
    watch: bool,
    exit_code: int | None,
    message: str | None,
    run_dir: str,
    summary: str,
    events: str,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "run_result",
        "run_id": run_id,
        "state": state,
        "watch": watch,
        "exit_code": exit_code,
        "message": message,
        "paths": {
            "run_dir": run_dir,
            "summary": summary,
            "events": events,
        },
    }


def submit_result_payload(
    *,
    run_id: str,
    state: str,
    slurm_job_id: str | None,
    exit_code: int | None,
    message: str | None,
    run_dir: str,
    summary: str,
    events: str,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "submit_result",
        "run_id": run_id,
        "state": state,
        "slurm_job_id": slurm_job_id,
        "exit_code": exit_code,
        "message": message,
        "paths": {
            "run_dir": run_dir,
            "summary": summary,
            "events": events,
        },
    }


def status_result_payload(
    *,
    run_id: str,
    wait: bool,
    timed_out: bool,
    backend: str,
    metadata: RunMetadata,
    status: RunStatus,
    runtime: RunRuntime | None,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "status_result",
        "run_id": run_id,
        "wait": wait,
        "timed_out": timed_out,
        "backend": backend,
        "metadata": metadata.model_dump(mode="json"),
        "status": status.model_dump(mode="json"),
        "runtime": runtime.model_dump(mode="json") if runtime else None,
    }


def validate_result_payload(
    *,
    manifest: str,
    valid: bool,
    backend: str,
    environment: str,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "validate_result",
        "manifest": manifest,
        "valid": valid,
        "backend": backend,
        "environment": environment,
    }


def cancel_result_payload(
    *,
    run_id: str,
    state: str,
    force: bool,
    cancelled: bool,
    exit_code: int | None = None,
    message: str | None = None,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "cancel_result",
        "run_id": run_id,
        "state": state,
        "force": force,
        "cancelled": cancelled,
        "exit_code": exit_code,
        "message": message,
    }


def rm_cleanup_result_payload(*, items: list[dict[str, str]], count: int, dry_run: bool) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "rm_cleanup_result",
        "items": items,
        "count": count,
        "dry_run": dry_run,
    }


def ls_result_payload(*, items: list[dict[str, Any]]) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "ls_result",
        "items": items,
    }


def ps_result_payload(*, items: list[dict[str, str]]) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "ps_result",
        "items": items,
    }


def rm_result_payload(
    *,
    run_id: str,
    removed: bool,
    cancelled: bool,
    force: bool,
    state_before: str,
    state_after: str | None = None,
) -> dict[str, Any]:
    return {
        "schema_version": JSON_SCHEMA_VERSION,
        "kind": "rm_result",
        "run_id": run_id,
        "removed": removed,
        "cancelled": cancelled,
        "force": force,
        "state_before": state_before,
        "state_after": state_after,
    }
