from __future__ import annotations

import json
import os
from datetime import datetime
from pathlib import Path

from one_run.backends.base import RunContext
from one_run.backends.local import LocalBackend
from one_run.backends.slurm import SlurmBackend
from one_run.constants import DEFAULT_RUNS_DIR
from one_run.errors import OneRunError
from one_run.models import Manifest, RunMetadata, RunRuntime, RunStatus
from one_run.run_dir import RunPaths
from one_run.runtime import slurm_runtime_snapshot
from one_run.status_store import read_metadata, read_runtime, read_status, write_runtime, write_status
from one_run.timefmt import iso_short


def _load_run_paths(run_id: str, project_root: Path) -> RunPaths:
    run_dir = project_root / DEFAULT_RUNS_DIR / run_id
    if not run_dir.exists():
        raise OneRunError(f"run not found: {run_id}")
    logs = run_dir / "logs"
    return RunPaths(
        run_dir=run_dir,
        config_snapshot=run_dir / "config.snapshot.yaml",
        metadata=run_dir / "metadata.json",
        status=run_dir / "status.json",
        runtime=run_dir / "runtime.json",
        local_pid=run_dir / "local.pid",
        local_exit_code=run_dir / "local.exit_code",
        logs_dir=logs,
        stdout_log=logs / "stdout.log",
        stderr_log=logs / "stderr.log",
        artifacts_dir=run_dir / "artifacts",
        metrics_dir=run_dir / "metrics",
    )


def _refresh_slurm_if_possible(run_id: str, project_root: Path, status: RunStatus, metadata: RunMetadata) -> RunStatus:
    if metadata.backend_type != "slurm":
        return status

    dummy_manifest = Manifest.model_validate(
        {
            "experiment": {"name": run_id, "seed": 0},
            "entrypoint": {"command": "echo noop"},
            "environment": {"kind": "local"},
            "backend": {"kind": "slurm"},
        }
    )

    paths = _load_run_paths(run_id, project_root)
    run_ctx = RunContext(
        run_id=run_id,
        manifest_path=paths.config_snapshot,
        manifest=dummy_manifest,
        paths=paths,
        project_root=project_root,
        metadata=metadata,
    )

    query = SlurmBackend().query(run_ctx, status)
    if query is None:
        return status

    ended_at = status.ended_at
    if query.state in {"success", "failed"} and ended_at is None:
        ended_at = datetime.now()

    updated = status.model_copy(update={"state": query.state, "message": query.message, "ended_at": ended_at})
    write_status(paths.status, updated)
    if query.runtime:
        runtime = slurm_runtime_snapshot(
            elapsed=query.runtime.get("slurm_elapsed"),
            max_rss=query.runtime.get("slurm_max_rss"),
            alloc_cpus=query.runtime.get("slurm_alloc_cpus"),
        )
        write_runtime(paths.runtime, RunRuntime.model_validate(runtime.model_dump(mode="json")))
    return updated


def _refresh_local_if_possible(paths: RunPaths, status: RunStatus) -> RunStatus:
    if status.state not in {"pending", "running"}:
        return status

    if paths.local_exit_code.exists():
        try:
            exit_code = int(paths.local_exit_code.read_text(encoding="utf-8").strip())
        except ValueError:
            exit_code = 1
        ended_at = status.ended_at or datetime.fromtimestamp(paths.local_exit_code.stat().st_mtime)
        updated = status.model_copy(
            update={
                "state": "success" if exit_code == 0 else "failed",
                "exit_code": exit_code,
                "ended_at": ended_at,
                "message": None if exit_code == 0 else f"command exited with code {exit_code}",
            }
        )
        write_status(paths.status, updated)
        return updated

    if paths.local_pid.exists():
        try:
            pid = int(paths.local_pid.read_text(encoding="utf-8").strip())
            os.kill(pid, 0)
        except ProcessLookupError:
            updated = status.model_copy(
                update={
                    "state": "failed",
                    "ended_at": status.ended_at or datetime.now(),
                    "message": "process ended without exit code file",
                }
            )
            write_status(paths.status, updated)
            return updated
        except (ValueError, OSError):
            pass

    return status


def show_status(run_id: str, *, json_output: bool = False, long: bool = False) -> int:
    project_root = Path.cwd().resolve()
    paths = _load_run_paths(run_id, project_root)
    metadata = RunMetadata.model_validate(read_metadata(paths.metadata))
    status = read_status(paths.status)

    if metadata.backend_type == "slurm":
        status = _refresh_slurm_if_possible(run_id, project_root, status, metadata)
    else:
        status = _refresh_local_if_possible(paths, status)
        LocalBackend().query(
            RunContext(
                run_id=run_id,
                manifest_path=paths.config_snapshot,
                manifest=Manifest.model_validate(
                    {
                        "experiment": {"name": run_id, "seed": 0},
                        "entrypoint": {"command": "echo noop"},
                        "environment": {"kind": "local"},
                        "backend": {"kind": "local"},
                    }
                ),
                paths=paths,
                project_root=project_root,
                metadata=metadata,
            ),
            status,
        )

    runtime = read_runtime(paths.runtime) if paths.runtime.exists() else None
    if json_output:
        payload = {
            "run_id": run_id,
            "backend": metadata.backend_type,
            "metadata": metadata.model_dump(mode="json"),
            "status": status.model_dump(mode="json"),
            "runtime": runtime.model_dump(mode="json") if runtime else None,
        }
        print(json.dumps(payload))
        return 0

    if long:
        print(f"state={status.state}")
        print(f"backend={metadata.backend_type}")
        print(f"timestamp={iso_short(metadata.timestamp)}")
        print(f"started_at={iso_short(status.started_at)}")
        print(f"ended_at={iso_short(status.ended_at)}")
        print(f"exit_code={status.exit_code}")
        if status.message:
            print(f"message={status.message}")
        print(f"hostname={metadata.hostname}")
        print(f"cwd={metadata.cwd}")
        print(f"git_commit={metadata.git_commit}")
        print(f"git_dirty={metadata.git_dirty}")
        if runtime:
            if runtime.duration_sec is not None:
                print(f"duration_sec={runtime.duration_sec}")
            if runtime.peak_rss_mb is not None:
                print(f"peak_rss_mb={runtime.peak_rss_mb}")
        print(f"stdout_log={paths.stdout_log}")
        print(f"stderr_log={paths.stderr_log}")
        print(f"artifacts_dir={paths.artifacts_dir}")
        return 0

    print(f"run_id={run_id}")
    print(f"state={status.state}")
    print(f"backend={metadata.backend_type}")
    if metadata.slurm_job_id:
        print(f"slurm_job_id={metadata.slurm_job_id}")
    if status.exit_code is not None:
        print(f"exit_code={status.exit_code}")
    if status.message:
        print(f"message={status.message}")
    return 0
