from __future__ import annotations

from pathlib import Path

import yaml
from types import SimpleNamespace

from argus.ray.models import RayConfig
from argus.service.config import V2Config
from argus.service.db import Db


def _mk_cfg(tmp_path: Path) -> tuple[RayConfig, V2Config]:
    root = {
        "ray": {
            "address": "http://127.0.0.1:8265",
            "shared_root": "/private",
            "entrypoint_resources": {"worker_node": 1},
            "runtime_env": {"env_vars": {}},
        },
        "service": {
            "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
            "scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
        },
    }
    return RayConfig.from_dict(root), V2Config.from_root_dict(root)


def test_tick_submits_one_task(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
    )

    class _Tool:
        def __init__(self, cfg):
            self.submitted = []
            self.job_dirs = []
            self.extra_env = []
            self.extra_overrides = []

        def submit(
            self,
            spec,
            no_wait: bool,
            job_dir: str | None = None,
            extra_env_vars: dict | None = None,
            extra_overrides: list[str] | None = None,
        ):
            self.submitted.append(spec.submission_id)
            self.job_dirs.append(job_dir)
            self.extra_env.append(extra_env_vars)
            self.extra_overrides.append(extra_overrides)
            return str(spec.submission_id)

        def status(self, submission_id: str):
            return "RUNNING"

        def logs(self, submission_id: str):
            return ""

    monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()

    row = db.get_task("t1")
    assert row and row["state"] == "SUBMITTED"
    attempts = db.list_attempts("t1")
    assert len(attempts) == 1
    assert attempts[0]["ray_submission_id"] == "t1--a01"
    assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
    env = s.tool.extra_env[-1]
    assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
    assert env["MVP_WANDB_RUN"] == "t1--a01"
    assert env["MVP_TRAINER_LOGGER"] == "console"
    assert s.tool.extra_overrides[-1] is None


def test_tick_submits_one_task_advanced(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="advanced",
        jobspec_yaml=yaml.safe_dump(
            {
                "kind": "advanced",
                "nnodes": 2,
                "n_gpus_per_node": 4,
                "command": "python3 -m verl.trainer.main_ppo +ray_kwargs.ray_init.address=auto",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
    )

    class _Tool:
        def __init__(self, cfg):
            self.submitted = []
            self.job_dirs = []
            self.extra_env = []

        def submit(self, spec, no_wait: bool, job_dir: str | None = None):
            raise AssertionError("should not call submit() for advanced")

        def submit_advanced(
            self,
            spec,
            no_wait: bool,
            job_dir: str | None = None,
            user_id: str | None = None,
            extra_env_vars: dict | None = None,
            extra_overrides: list[str] | None = None,
        ):
            self.submitted.append(spec.submission_id)
            self.job_dirs.append(job_dir)
            self.extra_env.append(extra_env_vars)
            return str(spec.submission_id)

        def status(self, submission_id: str):
            return "RUNNING"

        def logs(self, submission_id: str):
            return ""

    monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()

    row = db.get_task("t1")
    assert row and row["state"] == "SUBMITTED"
    attempts = db.list_attempts("t1")
    assert len(attempts) == 1
    assert attempts[0]["ray_submission_id"] == "t1--a01"
    assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
    env = s.tool.extra_env[-1]
    assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
    assert env["MVP_WANDB_RUN"] == "t1--a01"
    assert env["MVP_TRAINER_LOGGER"] == "console"


def test_tick_injects_wandb_env_when_enabled(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    root = {
        "ray": {
            "address": "http://127.0.0.1:8265",
            "shared_root": "/private",
            "entrypoint_resources": {"worker_node": 1},
            "runtime_env": {"env_vars": {}},
        },
        "service": {
            "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
            "scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
        },
        "tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
    }
    ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.setenv("WANDB_API_KEY", "k")
    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
    )

    class _Tool:
        def __init__(self, cfg):
            self.extra_env = []
            self.job_dirs = []

        def submit(
            self,
            spec,
            no_wait: bool,
            job_dir: str | None = None,
            extra_env_vars: dict | None = None,
            extra_overrides: list[str] | None = None,
        ):
            self.job_dirs.append(job_dir)
            self.extra_env.append(extra_env_vars)
            return str(spec.submission_id)

        def status(self, submission_id: str):
            return "RUNNING"

        def logs(self, submission_id: str):
            return ""

    monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()

    assert s.tool.job_dirs[-1] == "/private/users/alice/jobs/t1--a01"
    env = s.tool.extra_env[-1]
    assert env and env["WANDB_API_KEY"] == "k"
    assert env["WANDB_BASE_URL"] == "http://172.22.0.1:8090"
    assert env["WANDB_MODE"] == "online"
    assert env["WANDB_DIR"] == "/private/users/alice/jobs/t1--a01/wandb"
    assert env["MVP_WANDB_PROJECT"] == "alice_project"
    assert env["MVP_WANDB_RUN"] == "t1--a01"
    assert env["MVP_TRAINER_LOGGER"] == "[console,wandb]"


def test_tick_wandb_enabled_missing_key_degrades(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    root = {
        "ray": {
            "address": "http://127.0.0.1:8265",
            "shared_root": "/private",
            "entrypoint_resources": {"worker_node": 1},
            "runtime_env": {"env_vars": {}},
        },
        "service": {
            "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
            "scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
        },
        "tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
    }
    ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.delenv("WANDB_API_KEY", raising=False)
    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
    )

    class _Tool:
        def __init__(self, cfg):
            self.extra_env = []

        def submit(
            self,
            spec,
            no_wait: bool,
            job_dir: str | None = None,
            extra_env_vars: dict | None = None,
            extra_overrides: list[str] | None = None,
        ):
            self.extra_env.append(extra_env_vars)
            return str(spec.submission_id)

        def status(self, submission_id: str):
            return "RUNNING"

        def logs(self, submission_id: str):
            return ""

    monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()

    env = s.tool.extra_env[-1]
    assert env and env["MVP_WANDB_PROJECT"] == "alice_project"
    assert env["MVP_WANDB_RUN"] == "t1--a01"
    assert env["MVP_TRAINER_LOGGER"] == "console"
    attempts = db.list_attempts("t1")
    assert attempts and "WANDB_API_KEY" in str(attempts[0].get("message") or "")


def test_tick_wandb_enabled_updates_jobspec_logger(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    root = {
        "ray": {
            "address": "http://127.0.0.1:8265",
            "shared_root": "/private",
            "entrypoint_resources": {"worker_node": 1},
            "runtime_env": {"env_vars": {}},
        },
        "service": {
            "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
            "scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
        },
        "tracking": {"wandb": {"enabled": True, "base_url": "http://172.22.0.1:8090"}},
    }
    ray_cfg, v2_cfg = RayConfig.from_dict(root), V2Config.from_root_dict(root)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.setenv("WANDB_API_KEY", "k")
    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=999.0, total_available_npus=0.0),
    )

    class _Tool:
        def __init__(self, cfg):
            self.specs = []
            self.extra_overrides = []

        def submit(
            self,
            spec,
            no_wait: bool,
            job_dir: str | None = None,
            extra_env_vars: dict | None = None,
            extra_overrides: list[str] | None = None,
        ):
            self.specs.append(spec)
            self.extra_overrides.append(extra_overrides)
            return str(spec.submission_id)

        def status(self, submission_id: str):
            return "RUNNING"

        def logs(self, submission_id: str):
            return ""

    monkeypatch.setattr(sched_mod, "RayJobTool", _Tool)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()

    spec = s.tool.specs[-1]
    # v3.6: built-in training should be configured to use wandb when enabled and API key exists.
    assert spec.workload == "ppo"
    ov = s.tool.extra_overrides[-1]
    assert ov and "trainer.logger=[console,wandb]" in ov
    assert ov and "trainer.project_name=alice_project" in ov
    assert ov and "trainer.experiment_name=t1--a01" in ov


def test_tick_marks_pending_resources(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )

    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(
        sched_mod,
        "get_cluster_available",
        lambda: SimpleNamespace(total_available_gpus=0.0, total_available_npus=0.0),
    )
    monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)
    s.tick()
    row = db.get_task("t1")
    assert row and row["state"] == "PENDING_RESOURCES"
    assert row["next_run_at"]


def test_sync_failed_insufficient_resources(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump(
            {
                "workload": "ppo",
                "code_path": "/private/common/code/verl",
                "model_id": "m",
                "train_file": "/private/common/datasets/t",
                "val_file": "/private/common/datasets/v",
            }
        ),
        nnodes=2,
        n_gpus_per_node=4,
    )
    db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="t1--a01")
    db.set_task_state(task_id="t1", state="RUNNING", latest_attempt_no=1)

    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)

    class _Tool:
        def status(self, sid: str):
            return "FAILED"

        def logs(self, sid: str):
            # Match the service's regex exactly:
            # it expects literal backslashes and repeats of 's'/'d' (because of double-escaping).
            return "Total available GPUs\\ss\\dd\\ssis less than total desired GPUs\\ss\\dd"

    s.tool = _Tool()
    s.tick()

    row = db.get_task("t1")
    assert row and row["state"] == "PENDING_RESOURCES"
    attempts = db.list_attempts("t1")
    assert attempts[-1]["failure_kind"] == "INSUFFICIENT_RESOURCES"


def test_sync_status_error_keeps_state(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()
    db.create_task_v25(
        task_id="t1",
        user_id="alice",
        workload="ppo",
        jobspec_yaml=yaml.safe_dump({"workload": "ppo", "code_path": "/private/common/code/verl", "model_id": "m", "train_file": "/private/common/datasets/t"}),
        nnodes=2,
        n_gpus_per_node=4,
    )
    db.create_attempt(task_id="t1", attempt_no=1, ray_submission_id="t1--a01")
    db.set_task_state(task_id="t1", state="RUNNING", latest_attempt_no=1)

    monkeypatch.setattr(sched_mod, "ensure_ray_connected", lambda: None)
    monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)

    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)

    class _Tool:
        def status(self, sid: str):
            raise RuntimeError("boom")

    s.tool = _Tool()
    s.tick()
    row = db.get_task("t1")
    assert row and row["state"] == "RUNNING"


def test_run_forever_swallows_tick_exceptions(monkeypatch, tmp_path: Path):
    from argus.service import scheduler as sched_mod

    ray_cfg, v2_cfg = _mk_cfg(tmp_path)
    db = Db(v2_cfg.sqlite.db_path)
    db.init()

    monkeypatch.setattr(sched_mod, "RayJobTool", lambda cfg: None)
    s = sched_mod.Scheduler(db=db, ray_cfg=ray_cfg, v2_cfg=v2_cfg)

    calls = {"n": 0}

    def _tick():
        calls["n"] += 1
        raise RuntimeError("boom")

    monkeypatch.setattr(s, "tick", _tick)
    monkeypatch.setattr(sched_mod.time, "sleep", lambda _: None)

    class _Stop:
        def __init__(self):
            self._n = 0

        def is_set(self):
            self._n += 1
            return self._n > 1

    s.run_forever(_Stop())
    assert calls["n"] == 1
