from __future__ import annotations

import os
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 AdvancedTaskSpec, JobSpec, RayConfig, parse_taskspec
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
from .serve_client import RayServeClient
from .serving_reconciler import ServingReconciler


_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)
        self._serving: ServingReconciler | None = None
        if bool(self.v2_cfg.serving.enabled):
            self._serving = ServingReconciler(
                db=self.db,
                v2_cfg=self.v2_cfg,
                ray_runtime_env_env_vars=self.ray_cfg.runtime_env_env_vars,
                serve_client=RayServeClient(
                    http_port=int(self.v2_cfg.serving.serve.http_port),
                    proxy_location=str(self.v2_cfg.serving.serve.proxy_location),
                    ray_init_address="auto",
                ),
            )

    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 _wandb_env_vars(self, *, ray_submission_id: str, job_dir: str) -> tuple[dict[str, str], str | None]:
        """
        Returns (env_vars, warn_message).

        env_vars is always non-empty so that AdvancedTaskSpec examples can safely
        reference injected env vars without requiring the user to edit their command:
        - MVP_TRAINER_LOGGER: "console" or "[console,wandb]"
        - MVP_WANDB_PROJECT: "<user_id>_project"
        - MVP_WANDB_RUN: "<ray_submission_id>"

        When W&B is effectively enabled (config enabled + API key present),
        WANDB_* env vars are included as well.
        """
        env, warn, _enabled = self._tracking_env_vars(user_id=None, ray_submission_id=ray_submission_id, job_dir=job_dir)
        return env, warn

    def _tracking_env_vars(
        self, *, user_id: str | None, ray_submission_id: str, job_dir: str
    ) -> tuple[dict[str, str], str | None, bool]:
        """
        Returns (env_vars, warn_message, wandb_enabled_effective).

        - env_vars always includes MVP_* vars so advanced commands can use them.
        - wandb_enabled_effective is True only when tracking.wandb.enabled AND api key exists.
        """
        cfg = self.v2_cfg.tracking.wandb
        project_suffix = cfg.project_suffix or "_project"
        project_name = f"{user_id}{project_suffix}" if user_id else f"mvp{project_suffix}"

        api_key = os.environ.get(cfg.api_key_env, "").strip()
        wandb_enabled_effective = bool(cfg.enabled and api_key)
        warn: str | None = None
        if cfg.enabled and not api_key:
            warn = f"WARN_WANDB_DISABLED_MISSING_ENV:{cfg.api_key_env}"

        env: dict[str, str] = {
            "MVP_TRAINER_LOGGER": ("[console,wandb]" if wandb_enabled_effective else "console"),
            "MVP_WANDB_PROJECT": project_name,
            "MVP_WANDB_RUN": ray_submission_id,
        }

        if wandb_enabled_effective:
            env.update(
                {
                    "WANDB_BASE_URL": cfg.base_url,
                    "WANDB_API_KEY": api_key,
                    "WANDB_MODE": "online",
                    "WANDB_DIR": f"{job_dir.rstrip('/')}/wandb",
                }
            )

        return env, warn, wandb_enabled_effective

    def _wandb_training_overrides(self, *, user_id: str | None, ray_submission_id: str) -> list[str] | None:
        """
        Returns hydra overrides to enable W&B logging for built-in training workloads.

        Only call this when W&B is effectively enabled for the job (i.e. API key exists).
        """
        cfg = self.v2_cfg.tracking.wandb
        project_suffix = cfg.project_suffix or "_project"
        project_name = f"{user_id}{project_suffix}" if user_id else f"mvp{project_suffix}"
        return [
            "trainer.logger=[console,wandb]",
            f"trainer.project_name={project_name}",
            f"trainer.experiment_name={ray_submission_id}",
        ]

    def _parse_taskspec(self, jobspec_yaml: str) -> JobSpec | AdvancedTaskSpec:
        obj = yaml.safe_load(jobspec_yaml) or {}
        if not isinstance(obj, dict):
            raise ValueError("jobspec must be a YAML mapping")
        return parse_taskspec(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_taskspec(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 taskspec (v1.1 compatible)
        tracking_env, wandb_warn, wandb_enabled = self._tracking_env_vars(
            user_id=user_id_s, ray_submission_id=ray_sid, job_dir=job_dir
        )
        if isinstance(spec, JobSpec):
            d = spec.to_public_dict()
            d["submission_id"] = ray_sid
            spec2 = JobSpec.from_dict(d)
            wandb_overrides = self._wandb_training_overrides(user_id=user_id_s, ray_submission_id=ray_sid) if wandb_enabled else None
            submit = lambda: self.tool.submit(
                spec2,
                no_wait=True,
                job_dir=job_dir,
                extra_env_vars=tracking_env,
                extra_overrides=wandb_overrides,
            )
        else:
            d = spec.to_public_dict()
            d["submission_id"] = ray_sid
            spec2 = AdvancedTaskSpec.from_dict(d)
            submit = lambda: self.tool.submit_advanced(
                spec2, no_wait=True, job_dir=job_dir, user_id=user_id_s, extra_env_vars=tracking_env
            )

        try:
            submitted = submit()
            # 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",
                message=wandb_warn,
            )
            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()

        # v3.8: reconcile serve_models (best-effort).
        if self._serving is not None:
            try:
                self._serving.tick()
            except Exception:
                # Keep scheduler alive even if serving tick fails.
                pass

        # 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)))
