from __future__ import annotations

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


JOB_STATUS_QUEUED = "QUEUED"
JOB_STATUS_RUNNING = "RUNNING"
JOB_STATUS_COMPLETED = "COMPLETED"
JOB_STATUS_PARTIAL = "PARTIAL"
JOB_STATUS_FAILED = "FAILED"

STEP_STATUS_PENDING = "PENDING"
STEP_STATUS_RUNNING = "RUNNING"
STEP_STATUS_COMPLETED = "COMPLETED"
STEP_STATUS_PARTIAL = "PARTIAL"
STEP_STATUS_FAILED = "FAILED"


def utc_now() -> str:
    return (
        datetime.now(timezone.utc)
        .replace(microsecond=0)
        .isoformat()
        .replace("+00:00", "Z")
    )


class StateError(RuntimeError):
    pass


def _atomic_write_json(
    path: Path,
    data: dict[str, Any],
) -> None:
    path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    temp_path = path.with_name(
        f".{path.name}.tmp"
    )

    payload = json.dumps(
        data,
        ensure_ascii=False,
        indent=2,
    ) + "\n"

    temp_path.write_text(
        payload,
        encoding="utf-8",
    )

    temp_path.replace(path)


def _load_json(path: Path) -> dict[str, Any]:
    try:
        data = json.loads(
            path.read_text(
                encoding="utf-8",
            )
        )
    except FileNotFoundError as exc:
        raise StateError(
            f"Missing state file: {path}"
        ) from exc
    except json.JSONDecodeError as exc:
        raise StateError(
            f"Invalid JSON in state file {path}: {exc}"
        ) from exc

    if not isinstance(data, dict):
        raise StateError(
            f"State file must contain an object: {path}"
        )

    return data


def state_path(job_dir: Path) -> Path:
    return job_dir / "state" / "client_state.json"


def load_state(job_dir: Path) -> dict[str, Any]:
    return _load_json(
        state_path(job_dir)
    )


def save_state(
    job_dir: Path,
    state: dict[str, Any],
) -> None:
    _atomic_write_json(
        state_path(job_dir),
        state,
    )


def initialize_state(
    job_dir: Path,
    job: dict[str, Any],
) -> dict[str, Any]:
    job_id = job.get("job_id")

    if not isinstance(job_id, str) or not job_id:
        raise StateError(
            "Job is missing a valid job_id."
        )

    steps = job.get("steps")

    if not isinstance(steps, list) or not steps:
        raise StateError(
            f"Job {job_id} has no steps."
        )

    now = utc_now()

    state: dict[str, Any] = {
        "job_id": job_id,
        "status": JOB_STATUS_QUEUED,
        "created_at": now,
        "updated_at": now,
        "steps": {},
    }

    for step in steps:
        step_no = step.get("step_no")

        if not isinstance(step_no, int):
            raise StateError(
                f"Job {job_id} contains invalid step_no."
            )

        state["steps"][str(step_no)] = {
            "step_no": step_no,
            "step_id": step.get("step_id"),
            "session_name": step.get("session_name"),
            "status": STEP_STATUS_PENDING,
            "started_at": None,
            "completed_at": None,
            "exit_code": None,
            "restore_seen": False,
            "error": None,
        }

    save_state(
        job_dir,
        state,
    )

    return state


def update_job_state(
    job_dir: Path,
    status: str,
    *,
    error: str | None = None,
) -> dict[str, Any]:
    state = load_state(job_dir)

    state["status"] = status
    state["updated_at"] = utc_now()

    if error is not None:
        state["error"] = error
    elif status != JOB_STATUS_FAILED:
        state.pop("error", None)

    save_state(
        job_dir,
        state,
    )

    return state


def update_step_state(
    job_dir: Path,
    step_no: int,
    status: str,
    *,
    exit_code: int | None = None,
    restore_seen: bool | None = None,
    error: str | None = None,
    started: bool = False,
    completed: bool = False,
) -> dict[str, Any]:
    state = load_state(job_dir)

    key = str(step_no)
    steps = state.get("steps")

    if not isinstance(steps, dict) or key not in steps:
        raise StateError(
            f"Unknown step {step_no}."
        )

    step_state = steps[key]

    step_state["status"] = status

    if started:
        step_state["started_at"] = utc_now()

    if completed:
        step_state["completed_at"] = utc_now()

    if exit_code is not None:
        step_state["exit_code"] = exit_code

    if restore_seen is not None:
        step_state["restore_seen"] = restore_seen

    if error is not None:
        step_state["error"] = error
    elif status != STEP_STATUS_FAILED:
        step_state["error"] = None

    state["updated_at"] = utc_now()

    save_state(
        job_dir,
        state,
    )

    return state


def get_step_state(
    job_dir: Path,
    step_no: int,
) -> dict[str, Any]:
    state = load_state(job_dir)

    steps = state.get("steps")

    if not isinstance(steps, dict):
        raise StateError(
            f"Invalid steps state in {job_dir}."
        )

    step = steps.get(str(step_no))

    if not isinstance(step, dict):
        raise StateError(
            f"Unknown step {step_no}."
        )

    return step