from __future__ import annotations

import json
import os
import time
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.commands.common import load_run_paths
from one_run.utils.duration import parse_duration
from one_run.errors import OneRunError
from one_run.manifest import load_manifest
from one_run.models import Manifest, RunMetadata, RunRuntime, RunStatus
from one_run.payloads import status_result_payload
from one_run.run_dir import RunPaths
from one_run.runtime import slurm_runtime_snapshot
from one_run.utils.shell import run_capture
from one_run.status_store import read_metadata, read_runtime, read_status, write_metadata, write_runtime, write_status
from one_run.utils.timefmt import iso_short
from one_run.tracking import append_event, write_summary


_TERMINAL_STATES = {"success", "failed"}


def _retry_category(raw_state: str | None) -> str:
    state = (raw_state or "").upper()
    if "TIMEOUT" in state:
        return "timeout"
    if "CANCELLED" in state or "PREEMPTED" in state:
        return "cancelled"
    return "failed"


def _resubmit_slurm_job(paths: RunPaths, project_root: Path) -> str:
    script_path = paths.run_dir / "job.sbatch"
    if not script_path.exists():
        raise OneRunError(f"cannot retry slurm run: missing {script_path}")
    proc = run_capture(["sbatch", str(script_path)], cwd=project_root)
    if proc.returncode != 0:
        msg = (proc.stderr or proc.stdout).strip() or "sbatch failed"
        raise OneRunError(f"failed to retry slurm job: {msg}")
    job_id = SlurmBackend.parse_sbatch_job_id(proc.stdout or "")
    if not job_id:
        raise OneRunError(f"failed to parse slurm job id from retry output: {(proc.stdout or '').strip()}")
    return job_id


def _maybe_retry_slurm(
    run_id: str,
    project_root: Path,
    paths: RunPaths,
    metadata: RunMetadata,
    status: RunStatus,
    query_runtime: dict | None,
) -> RunStatus:
    try:
        manifest = load_manifest(paths.config_snapshot)
    except Exception:
        return status

    slurm_cfg = manifest.backend.slurm
    if not slurm_cfg:
        return status

    max_attempts = slurm_cfg.retry_max_attempts or 0
    if max_attempts <= 0:
        return status

    on_states = set(slurm_cfg.retry_on_states or ["failed", "timeout", "cancelled"])
    raw_state = (query_runtime or {}).get("slurm_state")
    category = _retry_category(raw_state)
    if category not in on_states:
        return status

    current_attempt = metadata.slurm_attempt or 1
    if current_attempt > max_attempts:
        return status

    backoff_sec = slurm_cfg.retry_backoff_sec or 0
    if backoff_sec > 0 and status.ended_at is not None:
        elapsed = (datetime.now() - status.ended_at).total_seconds()
        if elapsed < backoff_sec:
            remaining = int(backoff_sec - elapsed)
            pending = status.model_copy(
                update={"message": f"retry pending in {remaining}s (attempt {current_attempt + 1}/{max_attempts + 1})"}
            )
            if pending != status:
                write_status(paths.status, pending)
                append_event(paths, run_id=run_id, event="retry_pending", state=pending.state, message=pending.message)
                write_summary(paths, run_id)
            return pending

    new_job_id = _resubmit_slurm_job(paths, project_root)
    next_attempt = current_attempt + 1
    updated_metadata = metadata.model_copy(update={"slurm_job_id": new_job_id, "slurm_attempt": next_attempt})
    write_metadata(paths.metadata, updated_metadata.model_dump(mode="json"))

    retried = status.model_copy(
        update={
            "state": "pending",
            "started_at": datetime.now(),
            "ended_at": None,
            "exit_code": None,
            "message": f"retrying attempt {next_attempt}/{max_attempts + 1}",
        }
    )
    write_status(paths.status, retried)
    append_event(paths, run_id=run_id, event="retry_submitted", state="pending", message=retried.message)
    write_summary(paths, run_id)
    return retried


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(project_root, run_id)
    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})
    if updated != status:
        write_status(paths.status, updated)
        append_event(paths, run_id=run_id, event="state_changed", state=updated.state, message=updated.message)
        write_summary(paths, run_id)
    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")))
        write_summary(paths, run_id)
    if updated.state == "failed":
        updated = _maybe_retry_slurm(
            run_id=run_id,
            project_root=project_root,
            paths=paths,
            metadata=metadata,
            status=updated,
            query_runtime=query.runtime,
        )
    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)
        append_event(paths, run_id=paths.run_dir.name, event="state_changed", state=updated.state, message=updated.message)
        write_summary(paths, paths.run_dir.name)
        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)
            append_event(paths, run_id=paths.run_dir.name, event="state_changed", state=updated.state, message=updated.message)
            write_summary(paths, paths.run_dir.name)
            return updated
        except (ValueError, OSError):
            pass

    return status


def show_status(
    run_id: str,
    *,
    json_output: bool = False,
    long: bool = False,
    wait: bool = False,
    timeout: str | None = None,
    interval_sec: float = 2.0,
) -> int:
    project_root = Path.cwd().resolve()
    paths = load_run_paths(project_root, run_id)
    timeout_sec = None
    if timeout:
        try:
            timeout_sec = parse_duration(timeout).total_seconds()
        except ValueError as exc:
            raise OneRunError(str(exc)) from exc

    started_wait_at = datetime.now()
    timed_out = False

    while True:
        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,
            )

        if not wait or status.state in _TERMINAL_STATES:
            break

        if timeout_sec is not None:
            elapsed = (datetime.now() - started_wait_at).total_seconds()
            if elapsed >= timeout_sec:
                timed_out = True
                break

        time.sleep(max(interval_sec, 0.1))

    runtime = read_runtime(paths.runtime) if paths.runtime.exists() else None
    write_summary(paths, run_id)
    if json_output:
        payload = status_result_payload(
            run_id=run_id,
            wait=wait,
            timed_out=timed_out,
            backend=metadata.backend_type,
            metadata=metadata,
            status=status,
            runtime=runtime,
        )
        print(json.dumps(payload))
        return 1 if timed_out else 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 metadata.docker_image_digest:
            print(f"docker_image_digest={metadata.docker_image_digest}")
        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}")
    if wait:
        print(f"timed_out={str(timed_out).lower()}")
    return 1 if timed_out else 0
