from __future__ import annotations

from pathlib import Path

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": {}},
        },
        "data": {
            "user_root": str(tmp_path / "users"),
        },
        "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},
        },
        "serving": {
            "serve": {"http_port": 8000, "proxy_location": "HeadOnly"},
            "llm": {"accelerator_type": "H20"},
        },
    }
    p = tmp_path / "cfg.yaml"
    p.write_text(yaml.safe_dump(cfg), encoding="utf-8")
    return p


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

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")

    app = app_mod.create_app(str(cfg_path))

    admin_headers = {"authorization": "Bearer admin-token"}
    with TestClient(app) as c:
        r = c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        assert r.status_code == 200
        r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
        assert r2.status_code == 200
        user_token = r2.json()["token"]

        headers = {"authorization": f"Bearer {user_token}"}

        spec_yaml = (
            "model_id: qwen-0.5b\n"
            "model_source: $HOME/common/hf/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots/sha\n"
            "num_replicas: 1\n"
            "gpus_per_replica: 1\n"
        )
        r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
        assert r3.status_code == 200
        assert r3.json()["model_key"] == "mk-alice"
        assert r3.json()["state"] == "QUEUED"

        r4 = c.get("/api/v2/serve/models?limit=10&offset=0", headers=headers)
        assert r4.status_code == 200
        obj = r4.json()
        assert obj["openai_base_url"] is None
        assert len(obj["items"]) == 1
        assert obj["items"][0]["model_key"] == "mk-alice"
        assert obj["items"][0]["endpoint"]["openai_base_url"] == "http://testserver:8000/serve/mk-alice/v1"

        r5 = c.get("/api/v2/serve/models/mk-alice", headers=headers)
        assert r5.status_code == 200
        detail = r5.json()
        assert detail["model"]["model_key"] == "mk-alice"
        assert "model_id_prefix" in detail["model"]
        assert "resolved_spec_yaml" in detail
        assert isinstance(detail.get("events"), list)

        r6 = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 2})
        assert r6.status_code == 200
        assert r6.json()["state"] == "QUEUED"

        r7 = c.delete("/api/v2/serve/models/mk-alice", headers=headers)
        assert r7.status_code == 200
        assert r7.json()["state"] == "DELETING"

        # Admin status endpoint
        r8 = c.get("/api/v2/serve/status", headers=admin_headers)
        assert r8.status_code == 200
        assert r8.json()["http_port"] == 8000
        assert r8.json()["openai_base_url"] is None


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

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")

    app = app_mod.create_app(str(cfg_path))

    admin_headers = {"authorization": "Bearer admin-token"}
    with TestClient(app) as c:
        c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        r2 = c.post("/api/v2/users/alice/tokens", headers=admin_headers)
        user_token = r2.json()["token"]
        headers = {"authorization": f"Bearer {user_token}"}

        spec_yaml = (
            "model_id: x\n"
            "model_source: /private/users/bob/models/evil\n"
            "num_replicas: 1\n"
            "gpus_per_replica: 1\n"
        )
        r3 = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
        assert r3.status_code == 403


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

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
    app = app_mod.create_app(str(cfg_path))

    with TestClient(app) as c:
        # Create a user token
        admin_headers = {"authorization": "Bearer admin-token"}
        c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
        headers = {"authorization": f"Bearer {token}"}

        r = c.post("/api/v2/serve/models", headers=headers, data=": bad\n")
        assert r.status_code == 400

        r2 = c.post("/api/v2/serve/models", headers=headers, data="- 1\n- 2\n")
        assert r2.status_code == 400


def test_serving_api_engine_kwargs_binary_rejected(tmp_path: Path, monkeypatch):
    """
    yaml !!binary is parsed as bytes, which is not JSON-serializable.
    """
    from argus.service import app as app_mod

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: f"mk-{user_id}")
    app = app_mod.create_app(str(cfg_path))

    admin_headers = {"authorization": "Bearer admin-token"}
    with TestClient(app) as c:
        c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
        headers = {"authorization": f"Bearer {token}"}

        spec_yaml = (
            "model_id: x\n"
            "model_source: $HOME/common/hf/x\n"
            "engine_kwargs:\n"
            "  blob: !!binary \"AQID\"\n"
        )
        r = c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
        assert r.status_code == 400


def test_serving_api_list_include_deleted_and_forwarded_base_url(tmp_path: Path, monkeypatch):
    from argus.service import app as app_mod
    from argus.service.config import V2Config
    from argus.service.db import Db

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    keys = iter(["mk-alice-1", "mk-alice-2"])
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: next(keys))

    app = app_mod.create_app(str(cfg_path))

    admin_headers = {"authorization": "Bearer admin-token"}
    with TestClient(app) as c:
        c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
        headers = {"authorization": f"Bearer {token}"}

        spec_yaml = "model_id: x\nmodel_source: $HOME/common/hf/x\n"
        c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)
        c.post("/api/v2/serve/models", headers=headers, data=spec_yaml)

        # Mark one model as DELETED directly in DB (sets deleted_at).
        root = yaml.safe_load(cfg_path.read_text(encoding="utf-8"))
        v2_cfg = V2Config.from_root_dict(root)
        db = Db(v2_cfg.sqlite.db_path)
        db.set_serve_model_state(model_key="mk-alice-2", state="DELETED")

        r1 = c.get(
            "/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
            headers={**headers, "x-forwarded-host": "example.com:8080", "x-forwarded-proto": "https"},
        )
        assert r1.status_code == 200
        assert r1.json()["openai_base_url"] is None
        assert {m["model_key"] for m in r1.json()["items"]} == {"mk-alice-1"}
        assert r1.json()["items"][0]["endpoint"]["openai_base_url"] == "https://example.com:8000/serve/mk-alice-1/v1"

        # IPv6 bracketed host should be preserved.
        r1b = c.get(
            "/api/v2/serve/models?limit=10&offset=0&include_deleted=0",
            headers={**headers, "x-forwarded-host": "[::1]:8080", "x-forwarded-proto": "https"},
        )
        assert r1b.status_code == 200
        assert r1b.json()["items"][0]["endpoint"]["openai_base_url"] == "https://[::1]:8000/serve/mk-alice-1/v1"

        r2 = c.get("/api/v2/serve/models?include_deleted=1", headers=headers)
        assert r2.status_code == 200
        assert {m["model_key"] for m in r2.json()["items"]} == {"mk-alice-1", "mk-alice-2"}


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

    cfg_path = _write_config(tmp_path)
    monkeypatch.setenv("MVP_INTERNAL_TOKEN", "admin-token")

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

        def run_forever(self, stop_flag):
            return None

    monkeypatch.setattr(app_mod, "Scheduler", _Scheduler)
    monkeypatch.setattr(app_mod, "new_model_key", lambda user_id: "mk-alice")

    app = app_mod.create_app(str(cfg_path))

    admin_headers = {"authorization": "Bearer admin-token"}
    with TestClient(app) as c:
        c.post("/api/v2/users", headers=admin_headers, json={"user_id": "alice"})
        token = c.post("/api/v2/users/alice/tokens", headers=admin_headers).json()["token"]
        headers = {"authorization": f"Bearer {token}"}

        c.post("/api/v2/serve/models", headers=headers, data="model_id: x\nmodel_source: $HOME/common/hf/x\n")
        r = c.patch("/api/v2/serve/models/mk-alice", headers=headers, json={"num_replicas": 0})
        assert r.status_code == 422