from __future__ import annotations

import json
from pathlib import Path


class _FakeServeClient:
    def __init__(self):
        self.started = 0
        self.applied = []
        self.deleted = []
        self.status_calls = 0
        self.fail_apply = False
        self.fail_status = False
        self.fail_delete = False
        self.fail_start = False

    def ensure_started(self) -> None:
        if self.fail_start:
            raise RuntimeError("start boom")
        self.started += 1

    def apply_app(self, *, app, app_name: str, route_prefix: str = "/"):
        if self.fail_apply:
            raise RuntimeError("boom")
        self.applied.append({"app": app, "app_name": app_name, "route_prefix": route_prefix})
        return {"ok": True}

    def delete_app(self, *, app_name: str):
        if self.fail_delete:
            raise RuntimeError("delete boom")
        self.deleted.append({"app_name": app_name})
        return {"ok": True}

    def get_status(self):
        self.status_calls += 1
        if self.fail_status:
            raise RuntimeError("status boom")
        return {"ok": True}


def _seed_model(db, *, model_key: str, user_id: str, state: str, num_replicas: int = 1, gpus_per_replica: int = 1):
    spec_yaml = "model_id: x\nmodel_source: $HOME/common/hf/x\n"
    resolved_yaml = f"user_id: {user_id}\nmodel_id: {user_id}-202601061235-x\n"
    db.create_serve_model(
        model_key=model_key,
        user_id=user_id,
        model_id_suffix="x",
        model_id_prefix=f"{user_id}-202601061235",
        model_id=f"{user_id}-202601061235-x",
        model_source="/private/hf/x",
        num_replicas=num_replicas,
        gpus_per_replica=gpus_per_replica,
        engine_kwargs_json=json.dumps({"gpu_memory_utilization": 0.9}),
        spec_yaml=spec_yaml,
        resolved_spec_yaml=resolved_yaml,
    )
    db.set_serve_model_state(model_key=model_key, state=state, event_type="TEST_SEED")


def test_reconciler_skips_when_no_changes(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    root = {
        "ray": {"shared_root": "/private"},
        "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
        "data": {"sftpgo": {}, "retention": {}},
        "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
    }
    cfg = V2Config.from_root_dict(root)
    db = Db(cfg.sqlite.db_path)
    db.init()

    client = _FakeServeClient()
    rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client, get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})())
    rec.tick()
    assert client.started == 0
    assert client.applied == []


def test_reconciler_pending_resources_no_apply(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=2, gpus_per_replica=4)

    client = _FakeServeClient()
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 1, "total_available_npus": 0})(),
    )
    rec.tick()
    # Serve may be started even when resources are insufficient, but apply should not happen.
    assert client.started == 1
    assert client.applied == []
    # State remains QUEUED.
    row = db.get_serve_model("mk1")
    assert row and row["state"] == "QUEUED"
    ev = db.list_serve_events("mk1", limit=50)
    assert any(e["event_type"] == "SERVE_PENDING_RESOURCES" for e in ev)


def test_reconciler_apply_success_marks_running(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)

    client = _FakeServeClient()
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={"HF_ENDPOINT": "https://hf-mirror.com"},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
    )
    rec.tick()
    assert client.started == 1
    assert len(client.applied) == 1
    assert client.applied[0]["app_name"] == "mk1"
    assert client.applied[0]["route_prefix"] == "/serve/mk1"
    applied = client.applied[0]["app"]["llm_configs"]
    assert applied[0]["engine_kwargs"]["tensor_parallel_size"] == 1
    assert applied[0]["runtime_env"]["env_vars"]["HF_HUB_OFFLINE"] == "1"
    row = db.get_serve_model("mk1")
    assert row and row["state"] == "RUNNING"


def test_reconciler_delete_removes_and_marks_deleted(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="keep", user_id="alice", state="RUNNING", num_replicas=1, gpus_per_replica=1)
    _seed_model(db, model_key="del", user_id="alice", state="DELETING", num_replicas=1, gpus_per_replica=1)

    client = _FakeServeClient()
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
    )
    rec.tick()
    assert client.applied == []
    assert client.deleted == [{"app_name": "del"}]
    row = db.get_serve_model("del")
    assert row and row["state"] == "DELETED"
    assert row.get("deleted_at")


def test_reconciler_apply_failure_marks_failed(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="QUEUED")

    client = _FakeServeClient()
    client.fail_apply = True
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
    )
    rec.tick()
    row = db.get_serve_model("mk1")
    assert row and row["state"] == "FAILED"
    assert row.get("error_summary")


def test_reconciler_applies_one_model_per_tick(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)
    _seed_model(db, model_key="mk2", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)

    client = _FakeServeClient()
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
    )

    rec.tick()
    assert len(client.applied) == 1
    assert client.applied[0]["app_name"] == "mk1"
    assert db.get_serve_model("mk1")["state"] == "RUNNING"
    assert db.get_serve_model("mk2")["state"] == "QUEUED"

    rec.tick()
    assert len(client.applied) == 2
    assert client.applied[1]["app_name"] == "mk2"
    assert db.get_serve_model("mk2")["state"] == "RUNNING"


def test_reconciler_start_error_is_recorded(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="QUEUED", num_replicas=1, gpus_per_replica=1)

    client = _FakeServeClient()
    client.fail_start = True
    rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client)
    rec.tick()
    assert db.get_serve_model("mk1")["state"] == "QUEUED"
    ev = db.list_serve_events("mk1", limit=50)
    assert any(e["event_type"] == "SERVE_START_ERROR" for e in ev)


def test_reconciler_delete_failure_marks_failed(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()
    _seed_model(db, model_key="mk1", user_id="alice", state="DELETING", num_replicas=1, gpus_per_replica=1)

    client = _FakeServeClient()
    client.fail_delete = True
    rec = ServingReconciler(db=db, v2_cfg=cfg, ray_runtime_env_env_vars={}, serve_client=client)
    rec.tick()
    row = db.get_serve_model("mk1")
    assert row and row["state"] == "FAILED"
    assert row.get("error_summary")


def test_reconciler_engine_kwargs_bad_json_is_ignored(tmp_path: Path):
    from argus.service.config import V2Config
    from argus.service.db import Db
    from argus.service.serving_reconciler import ServingReconciler

    cfg = V2Config.from_root_dict(
        {
            "ray": {"shared_root": "/private"},
            "service": {"api": {}, "auth": {}, "sqlite": {"db_path": str(tmp_path / "mvp.sqlite3")}, "scheduler": {}},
            "data": {"sftpgo": {}, "retention": {}},
            "serving": {"serve": {"http_port": 8000}, "llm": {"accelerator_type": "H20"}},
        }
    )
    db = Db(cfg.sqlite.db_path)
    db.init()

    db.create_serve_model(
        model_key="mk1",
        user_id="alice",
        model_id_suffix="x",
        model_id_prefix="alice-202601061235",
        model_id="alice-202601061235-x",
        model_source="/private/hf/x",
        num_replicas=1,
        gpus_per_replica=1,
        engine_kwargs_json="not-json",
        spec_yaml="model_id: x\nmodel_source: $HOME/common/hf/x\n",
        resolved_spec_yaml="user_id: alice\nmodel_id: alice-202601061235-x\n",
    )
    db.set_serve_model_state(model_key="mk1", state="QUEUED", event_type="TEST_SEED")

    client = _FakeServeClient()
    rec = ServingReconciler(
        db=db,
        v2_cfg=cfg,
        ray_runtime_env_env_vars={},
        serve_client=client,
        get_available_fn=lambda: type("A", (), {"total_available_gpus": 8, "total_available_npus": 0})(),
    )
    rec.tick()
    cfgs = client.applied[0]["app"]["llm_configs"]
    assert cfgs[0]["engine_kwargs"]["tensor_parallel_size"] == 1
