from __future__ import annotations

import os
from pathlib import Path

import pytest
import yaml
from fastapi.testclient import TestClient


def _write_config(tmp_path: Path) -> Path:
    cfg = {
        "ray": {
            "address": "http://127.0.0.1:8265",
            "shared_root": "/private",
            "entrypoint_resources": {"worker_node": 1},
            "runtime_env": {"env_vars": {}},
        },
        "service": {
            "api": {"host": "127.0.0.1", "port": 0},
            "auth": {"token_env": "MVP_INTERNAL_TOKEN"},
            "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")},
            "scheduler": {"tick_s": 1, "retry_interval_s": 1, "max_running_tasks": 1},
        },
    }
    p = tmp_path / "cfg.yaml"
    p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
    return p


def test_auth_requires_token_env(tmp_path: Path, monkeypatch):
    from argus.service import app as app_mod

    cfg_path = _write_config(tmp_path)
    monkeypatch.delenv("MVP_INTERNAL_TOKEN", raising=False)

    class _Scheduler:
        def __init__(self, **kwargs):
            self.tool = object()

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    app = app_mod.create_app(str(cfg_path))

    with TestClient(app) as c:
        r = c.get("/api/v2/queue")
        assert r.status_code == 500


def test_task_submit_get_cancel_logs_queue(tmp_path: Path, monkeypatch):
    from argus.service import app as app_mod

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")
    monkeypatch.setattr(app_mod, "new_task_id", lambda workload, **kwargs: "tid1")

    class _Tool:
        def __init__(self):
            self.stopped = []

        def stop(self, sid: str):
            self.stopped.append(sid)
            return True

        def logs(self, sid: str):
            return "a\nb\nc\n"

    class _Scheduler:
        def __init__(self, **kwargs):
            self.tool = _Tool()

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    app = app_mod.create_app(str(cfg_path))

    headers = {"authorization": "Bearer token1"}
    with TestClient(app) as c:
        r = c.post(
            "/api/v2/tasks",
            headers=headers,
            data="workload: ppo\ncode_path: /private/common/code/verl\nmodel_id: m\ntrain_file: /private/common/datasets/t\n",
        )
        assert r.status_code == 200
        assert r.json()["task_id"] == "tid1"

        r2 = c.get("/api/v2/tasks/tid1", headers=headers)
        assert r2.status_code == 200
        assert r2.json()["desired_resources"]["total_gpus"] == 8

        r3 = c.get("/api/v2/queue", headers=headers)
        assert r3.status_code == 200
        assert "pending" in r3.json()

        r4 = c.post("/api/v2/tasks/tid1:cancel", headers=headers)
        assert r4.status_code == 200
        assert r4.json()["state"] == "CANCELED"

        # Seed an attempt then fetch logs
        from argus.service.db import Db
        from argus.service.config import V2Config
        from argus.ray.models import RayConfig

        root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
        ray_cfg = RayConfig.from_dict(root)
        v2_cfg = V2Config.from_root_dict(root)
        db = Db(v2_cfg.sqlite.db_path)
        db.create_task(
            task_id="tid2",
            workload="ppo",
            jobspec_yaml="workload: ppo\ncode_path: /private/common/code/verl\nmodel_id: m\ntrain_file: /private/common/datasets/t\n",
            nnodes=2,
            n_gpus_per_node=4,
        )
        db.create_attempt(task_id="tid2", attempt_no=1, ray_submission_id="sid2")
        db.set_task_state(task_id="tid2", state="RUNNING", latest_attempt_no=1)

        r5 = c.get("/api/v2/tasks/tid2/logs?tail=1", headers=headers)
        assert r5.status_code == 200
        assert r5.text.strip() == "c"


def test_submit_rejects_non_mapping_jobspec(tmp_path: Path, monkeypatch):
    from argus.service import app as app_mod

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")

    class _Scheduler:
        def __init__(self, **kwargs):
            self.tool = object()

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    app = app_mod.create_app(str(cfg_path))

    with TestClient(app) as c:
        r = c.post("/api/v2/tasks", headers={"authorization": "Bearer token1"}, data="- 1\n- 2\n")
        assert r.status_code == 400


def test_submit_rejects_invalid_jobspec(tmp_path: Path, monkeypatch):
    from argus.service import app as app_mod

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "token1")

    class _Scheduler:
        def __init__(self, **kwargs):
            self.tool = object()

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    app = app_mod.create_app(str(cfg_path))

    with TestClient(app) as c:
        r = c.post("/api/v2/tasks", headers={"authorization": "Bearer token1"}, data="workload: nope\n")
        assert r.status_code == 400
