from __future__ import annotations

import re
import time
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any

import yaml

from argus.core.ids import attempt_submission_id
from argus.ray.models import JobSpec, RayConfig
from argus.ray.ray_job_tool import RayJobTool

from .config import V2Config
from .db import Db
from .ray_resources import ensure_ray_connected, get_cluster_available


_INSUFFICIENT_RE = re.compile(r"Total available GPUs\\s+\\d+\\s+is less than total desired GPUs\\s+\\d+")


def _utc_now_iso() -> str:
    return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


def _utc_after_s(seconds: int) -> str:
    return (datetime.utcnow() + timedelta(seconds=seconds)).replace(microsecond=0).isoformat() + "Z"


@dataclass
class Scheduler:
    db: Db
    ray_cfg: RayConfig
    v2_cfg: V2Config

    def __post_init__(self) -> None:
        self.tool = RayJobTool(self.ray_cfg)

    def _job_dir_for_task(self, *, user_id: str | None, ray_submission_id: str) -> str:
        root = self.ray_cfg.shared_root.rstrip("/")
        if user_id:
            return f"{root}/users/{user_id}/jobs/{ray_submission_id}"
        return f"{root}/jobs/{ray_submission_id}"

    def _resources_sufficient(self, *, nnodes: int, n_gpus_per_node: int) -> bool:
        avail = get_cluster_available()
        required = float(nnodes * n_gpus_per_node)
        return avail.total_available_gpus >= required

    def _parse_jobspec(self, jobspec_yaml: str) -> JobSpec:
        obj = yaml.safe_load(jobspec_yaml) or {}
        if not isinstance(obj, dict):
            raise ValueError("jobspec must be a YAML mapping")
        return JobSpec.from_dict(obj)

    def _submit_one(self, task_row: dict[str, Any]) -> None:
        task_id = str(task_row["task_id"])
        jobspec_yaml = str(task_row["jobspec_yaml"])
        user_id = task_row.get("user_id")
        user_id_s = str(user_id) if user_id not in (None, "") else None

        spec = self._parse_jobspec(jobspec_yaml)
        attempt_no = int(task_row.get("latest_attempt_no", 0)) + 1
        ray_sid = attempt_submission_id(task_id, attempt_no)
        job_dir = self._job_dir_for_task(user_id=user_id_s, ray_submission_id=ray_sid)

        # Record attempt first so that we can surface it even if submit crashes.
        self.db.create_attempt(task_id=task_id, attempt_no=attempt_no, ray_submission_id=ray_sid)
        self.db.set_task_state(task_id=task_id, state="SUBMITTING", latest_attempt_no=attempt_no)

        # Override submission_id in jobspec (v1.1 compatible)
        d = spec.to_public_dict()
        d["submission_id"] = ray_sid
        spec2 = JobSpec.from_dict(d)

        try:
            submitted = self.tool.submit(spec2, no_wait=True, job_dir=job_dir)
            # submitted should equal ray_sid; keep as source of truth.
            self.db.update_attempt(task_id=task_id, attempt_no=attempt_no, ray_status="SUBMITTED")
            self.db.set_task_state(task_id=task_id, state="SUBMITTED")
            if submitted != ray_sid:
                self.db.set_task_state(task_id=task_id, state="SUBMITTED", event_type="WARN_SUBMISSION_ID_MISMATCH")
        except Exception as e:
            msg = repr(e)
            self.db.update_attempt(task_id=task_id, attempt_no=attempt_no, ray_status="FAILED", failure_kind="UNKNOWN", message=msg, end_time=_utc_now_iso())
            self.db.set_task_state(task_id=task_id, state="FAILED", error_summary=msg)

    def _sync_one_running(self, task_row: dict[str, Any]) -> None:
        task_id = str(task_row["task_id"])
        latest_attempt_no = int(task_row.get("latest_attempt_no", 0))
        if latest_attempt_no <= 0:
            return

        # Look up ray_submission_id
        attempts = self.db.list_attempts(task_id)
        if not attempts:
            return
        ray_sid = str(attempts[-1]["ray_submission_id"])

        try:
            st = self.tool.status(ray_sid)
        except Exception as e:
            # Keep current state; transient failures should not flap tasks.
            self.db.set_task_state(task_id=task_id, state=str(task_row["state"]), event_type="RAY_STATUS_ERROR", payload_json=repr(e))
            return

        st_s = str(st)
        if st_s in ("PENDING", "RUNNING"):
            self.db.update_attempt(task_id=task_id, attempt_no=latest_attempt_no, ray_status=st_s)
            self.db.set_task_state(task_id=task_id, state=("RUNNING" if st_s == "RUNNING" else "SUBMITTED"))
            return

        if st_s in ("SUCCEEDED",):
            self.db.update_attempt(task_id=task_id, attempt_no=latest_attempt_no, ray_status=st_s, end_time=_utc_now_iso())
            self.db.set_task_state(task_id=task_id, state="SUCCEEDED")
            return

        if st_s in ("FAILED", "STOPPED"):
            logs = ""
            try:
                logs = self.tool.logs(ray_sid)
            except Exception:
                logs = ""
            failure_kind = "UNKNOWN"
            msg = ""
            if _INSUFFICIENT_RE.search(logs):
                failure_kind = "INSUFFICIENT_RESOURCES"
                msg = "Insufficient resources (verl fail-fast): " + (_INSUFFICIENT_RE.search(logs).group(0))
                self.db.update_attempt(
                    task_id=task_id,
                    attempt_no=latest_attempt_no,
                    ray_status=st_s,
                    failure_kind=failure_kind,
                    message=msg,
                    end_time=_utc_now_iso(),
                )
                self.db.set_task_state(
                    task_id=task_id,
                    state="PENDING_RESOURCES",
                    error_summary=msg,
                    next_run_at=_utc_after_s(self.v2_cfg.scheduler.retry_interval_s),
                    event_type="RETRY_SCHEDULED",
                )
                return

            msg = f"Ray job {st_s}"
            self.db.update_attempt(
                task_id=task_id,
                attempt_no=latest_attempt_no,
                ray_status=st_s,
                failure_kind=failure_kind,
                message=msg,
                end_time=_utc_now_iso(),
            )
            self.db.set_task_state(task_id=task_id, state="FAILED", error_summary=msg)

    def tick(self) -> None:
        ensure_ray_connected()

        # Sync active tasks
        for row in self.db.list_active_tasks(limit=50):
            self._sync_one_running(row)

        # Submit new tasks if capacity allows
        if self.db.count_running() >= self.v2_cfg.scheduler.max_running_tasks:
            return

        row = self.db.pick_next_runnable_task()
        if not row:
            return

        nnodes = int(row["nnodes"])
        n_gpus_per_node = int(row["n_gpus_per_node"])

        if not self._resources_sufficient(nnodes=nnodes, n_gpus_per_node=n_gpus_per_node):
            self.db.set_task_state(
                task_id=str(row["task_id"]),
                state="PENDING_RESOURCES",
                next_run_at=_utc_after_s(self.v2_cfg.scheduler.retry_interval_s),
                event_type="PENDING_RESOURCES",
            )
            return

        self._submit_one(row)

    def run_forever(self, stop_flag: Any) -> None:
        while not stop_flag.is_set():
            try:
                self.tick()
            except Exception:
                # Best-effort: don't crash the scheduler loop
                pass
            time.sleep(max(1, int(self.v2_cfg.scheduler.tick_s)))
