from __future__ import annotations

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

from one_run.commands.common import load_run_paths
from one_run.errors import OneRunError
from one_run.models import RunMetadata, RunStatus
from one_run.payloads import cancel_result_payload
from one_run.run_dir import RunPaths
from one_run.utils.shell import command_exists, run_capture
from one_run.status_store import read_metadata, read_status, write_status
from one_run.tracking import append_event, write_summary

def _cancel_local(paths: RunPaths, status: RunStatus, *, force: bool = False) -> RunStatus:
    now = datetime.now()

    if paths.local_pid.exists():
        try:
            pid = int(paths.local_pid.read_text(encoding="utf-8").strip())
            os.kill(pid, signal.SIGKILL if force else signal.SIGTERM)
        except ProcessLookupError:
            pass
        except (ValueError, OSError) as exc:
            raise OneRunError(f"failed to cancel local run: {exc}") from exc

    exit_code = 137 if force else 130
    message = "force cancelled by user" if force else "cancelled by user"
    paths.local_exit_code.write_text(f"{exit_code}\n", encoding="utf-8")
    updated = status.model_copy(
        update={
            "state": "failed",
            "ended_at": now,
            "exit_code": exit_code,
            "message": message,
        }
    )
    write_status(paths.status, updated)
    append_event(paths, run_id=paths.run_dir.name, event="cancelled", state=updated.state, message=updated.message)
    write_summary(paths, paths.run_dir.name)
    return updated


def _cancel_slurm(metadata: RunMetadata, status: RunStatus, paths: RunPaths, *, force: bool = False) -> RunStatus:
    job_id = metadata.slurm_job_id
    if not job_id:
        raise OneRunError("slurm run has no job id")
    if not command_exists("scancel"):
        raise OneRunError("scancel not found. cannot cancel slurm job")

    cmd = ["scancel", "--signal=KILL", job_id] if force else ["scancel", job_id]
    proc = run_capture(cmd)
    if proc.returncode != 0:
        msg = (proc.stderr or proc.stdout).strip() or "scancel failed"
        raise OneRunError(f"failed to cancel slurm job {job_id}: {msg}")

    exit_code = 137 if force else 130
    message = "force cancelled by user" if force else "cancelled by user"
    updated = status.model_copy(
        update={
            "state": "failed",
            "ended_at": datetime.now(),
            "exit_code": exit_code,
            "message": message,
        }
    )
    write_status(paths.status, updated)
    append_event(paths, run_id=paths.run_dir.name, event="cancelled", state=updated.state, message=updated.message)
    write_summary(paths, paths.run_dir.name)
    return updated


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

    if status.state in {"success", "failed"}:
        if json_output:
            print(json.dumps(cancel_result_payload(
                run_id=run_id,
                state=status.state,
                force=force,
                cancelled=False,
            )))
        else:
            print(f"run_id={run_id}")
            print(f"state={status.state}")
            print("cancelled=false")
        return 0

    if metadata.backend_type == "slurm":
        updated = _cancel_slurm(metadata, status, paths, force=force)
    else:
        updated = _cancel_local(paths, status, force=force)

    if json_output:
        print(json.dumps(cancel_result_payload(
            run_id=run_id,
            state=updated.state,
            force=force,
            cancelled=True,
            exit_code=updated.exit_code,
            message=updated.message,
        )))
    else:
        print(f"run_id={run_id}")
        print(f"state={updated.state}")
        print(f"force={str(force).lower()}")
        print("cancelled=true")
    return 0
