from __future__ import annotations

import json
import traceback
from dataclasses import dataclass
from typing import Any, Protocol

from argus.service.ray_resources import ClusterAvailable, get_cluster_available

from .config import V2Config
from .db import Db
from .serve_llm_config import build_llm_config_dict
from .serving_spec import ResolvedServingSpec


class ServeClient(Protocol):
    def ensure_started(self) -> None: ...

    def apply_app(self, *, app: Any, app_name: str, route_prefix: str = "/") -> Any: ...

    def delete_app(self, *, app_name: str) -> Any: ...

    def get_status(self) -> Any: ...


def _parse_engine_kwargs(row: dict[str, Any]) -> dict[str, Any] | None:
    raw = row.get("engine_kwargs_json")
    if raw in (None, ""):
        return None
    try:
        obj = json.loads(str(raw))
        return obj if isinstance(obj, dict) else None
    except Exception:
        return None


def _row_to_resolved_spec(row: dict[str, Any]) -> ResolvedServingSpec:
    return ResolvedServingSpec(
        user_id=str(row["user_id"]),
        model_id_suffix=str(row["model_id_suffix"]),
        model_id_prefix=str(row["model_id_prefix"]),
        model_id=str(row["model_id"]),
        model_source=str(row["model_source"]),
        num_replicas=int(row["num_replicas"]),
        gpus_per_replica=int(row["gpus_per_replica"]),
        engine_kwargs=_parse_engine_kwargs(row),
    )


def _needed_model_gpus(row: dict[str, Any]) -> int:
    return int(row.get("num_replicas") or 0) * int(row.get("gpus_per_replica") or 0)


@dataclass
class ServingReconciler:
    """
    v3.8 (per-model apps): reconcile serve_models (SQLite) into per-model Ray Serve apps.

    This reconciler is intentionally conservative:
    - Only acts on models in states QUEUED/DELETING.
    - Performs a minimal GPU precheck using ray available GPU totals (per-model).
    - Writes events and state transitions for explainability.
    """

    db: Db
    v2_cfg: V2Config
    ray_runtime_env_env_vars: dict[str, str]
    serve_client: ServeClient
    cpu_per_gpu: float = 1.0
    get_available_fn: Any = get_cluster_available

    def tick(self) -> None:
        # Pick the next desired change.
        change = self.db.pick_next_runnable_serve_change()
        if not change:
            return

        model_key = str(change["model_key"])
        state = str(change.get("state") or "")

        # Ensure Ray (and Serve) can be started before doing anything else.
        try:
            self.serve_client.ensure_started()
        except Exception as e:
            self.db.append_serve_event(model_key=model_key, event_type="SERVE_START_ERROR", payload_json=repr(e))
            return

        # Deleting is per-model: remove the model's app by name. This avoids touching other models.
        if state == "DELETING":
            try:
                self.db.append_serve_event(model_key=model_key, event_type="SERVE_DELETE_REQUESTED", payload_json=None)
                self.serve_client.delete_app(app_name=model_key)
            except Exception as e:
                err = f"{type(e).__name__}: {e}"
                tb = traceback.format_exc(limit=10)
                self.db.set_serve_model_state(
                    model_key=model_key, state="FAILED", error_summary=err, event_type="SERVE_DELETE_FAILED", payload_json=tb
                )
                return
            self.db.set_serve_model_state(model_key=model_key, state="DELETED", event_type="SERVE_DELETE_APPLIED")
            return

        # Precheck resources: per-model apply needs enough GPUs for this model only.
        needed = _needed_model_gpus(change)
        avail: ClusterAvailable = self.get_available_fn()
        if float(avail.total_available_gpus) < float(needed):
            msg = f"Insufficient GPUs: need {needed}, available {avail.total_available_gpus}"
            self.db.append_serve_event(model_key=model_key, event_type="SERVE_PENDING_RESOURCES", payload_json=msg)
            return

        # Build the single-model LLM config (dict form for unit tests).
        accelerator_type = str(self.v2_cfg.serving.llm.accelerator_type or "")
        resolved = _row_to_resolved_spec(change)
        llm_cfg_dicts = [
            build_llm_config_dict(
                resolved,
                accelerator_type=accelerator_type,
                runtime_env_env_vars=self.ray_runtime_env_env_vars,
                cpu_per_gpu=self.cpu_per_gpu,
            )
        ]

        # Build a Ray Serve OpenAI-compatible app if Ray Serve LLM is available.
        # Fall back to a plain dict so unit tests can run without real Ray Serve.
        app_obj: Any
        try:
            from ray.serve.llm import LLMConfig, build_openai_app  # type: ignore

            llm_cfgs = [LLMConfig(**d) for d in llm_cfg_dicts]
            app_obj = build_openai_app({"llm_configs": llm_cfgs})
        except Exception as e:
            self.db.append_serve_event(model_key=model_key, event_type="SERVE_LLM_IMPORT_ERROR", payload_json=repr(e))
            app_obj = {"llm_configs": llm_cfg_dicts}

        try:
            self.db.append_serve_event(model_key=model_key, event_type="SERVE_APPLY_REQUESTED", payload_json="1")
            self.serve_client.apply_app(app=app_obj, app_name=model_key, route_prefix=f"/serve/{model_key}")
        except Exception as e:
            err = f"{type(e).__name__}: {e}"
            tb = traceback.format_exc(limit=10)
            self.db.set_serve_model_state(model_key=model_key, state="FAILED", error_summary=err, event_type="SERVE_APPLY_FAILED", payload_json=tb)
            return

        # Mark as deploying; best-effort status probe can promote to RUNNING.
        self.db.set_serve_model_state(model_key=model_key, state="DEPLOYING", event_type="SERVE_DEPLOYING")
        try:
            _ = self.serve_client.get_status()
            self.db.set_serve_model_state(model_key=model_key, state="RUNNING", event_type="SERVE_RUNNING")
        except Exception as e:
            self.db.append_serve_event(model_key=model_key, event_type="SERVE_STATUS_ERROR", payload_json=repr(e))