from __future__ import annotations

import re
from datetime import datetime
from pathlib import Path

from one_run.backends.base import BackendQueryResult, BackendStartResult, RunContext
from one_run.constants import JOB_SCRIPT_NAME
from one_run.errors import BackendError
from one_run.utils.shell import command_exists, run_capture
from one_run.utils.templating import render_command


SLURM_STATE_MAP = {
    "PENDING": "pending",
    "CONFIGURING": "pending",
    "RUNNING": "running",
    "COMPLETING": "running",
    "COMPLETED": "success",
    "FAILED": "failed",
    "CANCELLED": "failed",
    "TIMEOUT": "failed",
    "NODE_FAIL": "failed",
    "PREEMPTED": "failed",
    "OUT_OF_MEMORY": "failed",
}


class SlurmBackend:
    def prepare(self, run_ctx: RunContext) -> None:
        if not command_exists("sbatch"):
            raise BackendError("sbatch not found. slurm backend is unavailable")

    def _build_sbatch_script(self, run_ctx: RunContext) -> Path:
        resources = run_ctx.manifest.resources
        slurm_cfg = run_ctx.manifest.backend.slurm
        lines = ["#!/usr/bin/env bash", "set -euo pipefail"]

        if slurm_cfg and slurm_cfg.partition:
            lines.append(f"#SBATCH --partition={slurm_cfg.partition}")
        if resources and resources.gpus:
            lines.append(f"#SBATCH --gpus={resources.gpus}")
        if resources and resources.cpus:
            lines.append(f"#SBATCH --cpus-per-task={resources.cpus}")
        if resources and resources.mem_gb:
            lines.append(f"#SBATCH --mem={resources.mem_gb}G")
        if resources and resources.time_min:
            lines.append(f"#SBATCH --time={resources.time_min}")

        lines.append(f"#SBATCH --output={run_ctx.paths.stdout_log}")
        lines.append(f"#SBATCH --error={run_ctx.paths.stderr_log}")
        lines.append("")

        command = render_command(
            run_ctx.manifest.entrypoint.command,
            {
                "config": str(run_ctx.paths.config_snapshot),
                "run_dir": str(run_ctx.paths.run_dir),
                "seed": run_ctx.manifest.experiment.seed,
                "project_root": str(run_ctx.project_root),
            },
        )

        lines.append(f"cd {run_ctx.project_root}")
        lines.append(command)

        script_path = run_ctx.paths.run_dir / JOB_SCRIPT_NAME
        script_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
        return script_path

    def start(self, run_ctx: RunContext) -> BackendStartResult:
        script_path = self._build_sbatch_script(run_ctx)
        proc = run_capture(["sbatch", str(script_path)], cwd=run_ctx.project_root)
        if proc.returncode != 0:
            msg = (proc.stderr or proc.stdout).strip() or "sbatch failed"
            return BackendStartResult(
                state="failed",
                started_at=datetime.now(),
                ended_at=datetime.now(),
                exit_code=proc.returncode,
                message=msg,
            )

        output = proc.stdout.strip()
        job_id = self.parse_sbatch_job_id(output)
        if not job_id:
            raise BackendError(f"could not parse slurm job id from sbatch output: {output}")

        return BackendStartResult(
            state="pending",
            started_at=datetime.now(),
            ended_at=None,
            exit_code=None,
            message=None,
            slurm_job_id=job_id,
        )

    @staticmethod
    def parse_sbatch_job_id(output: str) -> str | None:
        match = re.search(r"Submitted batch job\s+(\d+)", output.strip())
        if not match:
            return None
        return match.group(1)

    def _query_squeue(self, job_id: str) -> str | None:
        if not command_exists("squeue"):
            return None
        proc = run_capture(["squeue", "-h", "-j", job_id, "-o", "%T"])
        if proc.returncode != 0:
            return None
        state = proc.stdout.strip().splitlines()
        if not state:
            return None
        return state[0].strip().upper()

    def _query_sacct_row(self, job_id: str) -> dict[str, str] | None:
        if not command_exists("sacct"):
            return None
        proc = run_capture(
            [
                "sacct",
                "-n",
                "-j",
                job_id,
                "--format=State,Elapsed,MaxRSS,AllocCPUS,ExitCode",
                "--parsable2",
            ]
        )
        if proc.returncode != 0:
            return None
        lines = [ln.strip() for ln in proc.stdout.splitlines() if ln.strip()]
        if not lines:
            return None
        parts = lines[0].split("|")
        row = {
            "state": parts[0].strip().upper() if len(parts) > 0 else "",
            "elapsed": parts[1].strip() if len(parts) > 1 else "",
            "max_rss": parts[2].strip() if len(parts) > 2 else "",
            "alloc_cpus": parts[3].strip() if len(parts) > 3 else "",
            "exit_code": parts[4].strip() if len(parts) > 4 else "",
        }
        if not row["state"]:
            return None
        row["state"] = row["state"].split("+")[0]
        return row

    def query(self, run_ctx: RunContext, status) -> BackendQueryResult | None:
        job_id = run_ctx.metadata.slurm_job_id
        if not job_id:
            return None

        sacct_row = self._query_sacct_row(job_id)
        raw_state = self._query_squeue(job_id) or (sacct_row["state"] if sacct_row else None)
        if raw_state is None:
            return None

        mapped = SLURM_STATE_MAP.get(raw_state, "running")
        message = None
        if mapped == "failed":
            message = f"slurm state: {raw_state}"
        runtime = None
        if sacct_row:
            alloc_cpus = None
            if sacct_row["alloc_cpus"].isdigit():
                alloc_cpus = int(sacct_row["alloc_cpus"])
            runtime = {
                "slurm_state": raw_state,
                "slurm_elapsed": sacct_row["elapsed"] or None,
                "slurm_max_rss": sacct_row["max_rss"] or None,
                "slurm_alloc_cpus": alloc_cpus,
            }
        elif raw_state:
            runtime = {"slurm_state": raw_state}
        return BackendQueryResult(state=mapped, exit_code=None, message=message, runtime=runtime)

    def cancel(self, run_ctx: RunContext) -> None:
        return
