from __future__ import annotations

import json
from pathlib import Path

from one_run.backends.base import RunContext
from one_run.backends.slurm import SlurmBackend
from one_run.commands.common import alloc_run_dir
from one_run.commands.deps_cmd import wait_for_dependencies
from one_run.errors import BackendError
from one_run.manifest import load_manifest
from one_run.metadata import build_metadata
from one_run.models import RunRuntime, RunStatus
from one_run.payloads import submit_result_payload
from one_run.runtime import local_runtime_snapshot
from one_run.status_store import write_metadata, write_runtime, write_status
from one_run.tracking import append_event, write_summary


def submit_slurm(
    config_path: str,
    *,
    json_output: bool = False,
    wait_for: list[str] | None = None,
    wait_timeout: str | None = None,
    wait_interval: float | None = None,
) -> int:
    manifest_path = Path(config_path).resolve()
    manifest = load_manifest(manifest_path)

    if manifest.backend.kind != "slurm":
        raise BackendError("this run requires backend.kind=slurm")

    project_root = Path.cwd().resolve()
    wait_for_dependencies(
        manifest,
        project_root=project_root,
        wait_for_run_ids=wait_for,
        wait_timeout=wait_timeout,
        wait_interval=wait_interval,
        emit_console=not json_output,
    )
    run_id, paths = alloc_run_dir(project_root, manifest.experiment.name)
    if not json_output:
        print(f"run_id={run_id}", flush=True)
    paths.config_snapshot.write_text(manifest_path.read_text(encoding="utf-8"), encoding="utf-8")

    metadata = build_metadata(run_id, manifest, project_root, Path.cwd())
    write_metadata(paths.metadata, metadata.model_dump(mode="json"))
    write_runtime(paths.runtime, local_runtime_snapshot(duration_sec=None, peak_rss_raw=None))
    write_status(paths.status, RunStatus(state="pending"))
    append_event(paths, run_id=run_id, event="created", state="pending")
    write_summary(paths, run_id)

    backend = SlurmBackend()
    run_ctx = RunContext(
        run_id=run_id,
        manifest_path=manifest_path,
        manifest=manifest,
        paths=paths,
        project_root=project_root,
        metadata=metadata,
    )

    backend.prepare(run_ctx)
    result = backend.start(run_ctx)

    if result.slurm_job_id:
        metadata = metadata.model_copy(update={"slurm_job_id": result.slurm_job_id})
        write_metadata(paths.metadata, metadata.model_dump(mode="json"))

    write_status(
        paths.status,
        RunStatus(
            state=result.state,
            started_at=result.started_at,
            ended_at=result.ended_at,
            exit_code=result.exit_code,
            message=result.message,
        ),
    )
    if result.runtime:
        write_runtime(paths.runtime, RunRuntime.model_validate(result.runtime))
    append_event(paths, run_id=run_id, event="submitted", state=result.state, message=result.message)
    write_summary(paths, run_id)

    if json_output:
        print(json.dumps(submit_result_payload(
            run_id=run_id,
            state=result.state,
            slurm_job_id=result.slurm_job_id,
            exit_code=result.exit_code,
            message=result.message,
            run_dir=str(paths.run_dir),
            summary=str(paths.summary),
            events=str(paths.events),
        )))
    else:
        if result.slurm_job_id:
            print(f"slurm_job_id={result.slurm_job_id}")
        print(f"state={result.state}")
    return 0 if result.state in {"pending", "running", "success"} else 1
