"""Pods config management — load/save/active pod."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

from pi_agents.pods.types import Config, GPU, Pod, PodModel


def _get_config_path() -> str:
    """Get pods config file path."""
    return os.path.join(Path.home(), ".pi", "pods.json")


def _pod_from_dict(d: dict[str, Any]) -> Pod:
    """Deserialize a Pod from a dict."""
    gpus = [GPU(id=g.get("id", 0), name=g.get("name", ""), memory=g.get("memory", "")) for g in d.get("gpus", [])]
    models: dict[str, PodModel] = {}
    for k, v in d.get("models", {}).items():
        models[k] = PodModel(
            model=v.get("model", ""),
            port=v.get("port", 0),
            gpu=v.get("gpu", []),
            pid=v.get("pid", 0),
        )
    return Pod(
        ssh=d.get("ssh", ""),
        gpus=gpus,
        models=models,
        models_path=d.get("modelsPath"),
        vllm_version=d.get("vllmVersion"),
    )


def _pod_to_dict(pod: Pod) -> dict[str, Any]:
    """Serialize a Pod to a dict."""
    result: dict[str, Any] = {
        "ssh": pod.ssh,
        "gpus": [{"id": g.id, "name": g.name, "memory": g.memory} for g in pod.gpus],
        "models": {
            k: {"model": v.model, "port": v.port, "gpu": v.gpu, "pid": v.pid}
            for k, v in pod.models.items()
        },
    }
    if pod.models_path is not None:
        result["modelsPath"] = pod.models_path
    if pod.vllm_version is not None:
        result["vllmVersion"] = pod.vllm_version
    return result


def load_config() -> Config:
    """Load pods configuration from disk."""
    path = _get_config_path()
    if not os.path.exists(path):
        return Config()
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        pods = {k: _pod_from_dict(v) for k, v in data.get("pods", {}).items()}
        return Config(pods=pods, active=data.get("active"))
    except Exception:
        return Config()


def save_config(config: Config) -> None:
    """Save pods configuration to disk."""
    path = _get_config_path()
    os.makedirs(os.path.dirname(path), exist_ok=True)
    data: dict[str, Any] = {
        "pods": {k: _pod_to_dict(v) for k, v in config.pods.items()},
    }
    if config.active is not None:
        data["active"] = config.active
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)


def get_active_pod() -> tuple[str, Pod] | None:
    """Get the active pod name and config. Returns None if no active pod."""
    config = load_config()
    if not config.active or config.active not in config.pods:
        return None
    return (config.active, config.pods[config.active])


def set_active_pod(name: str) -> None:
    """Set the active pod."""
    config = load_config()
    if name not in config.pods:
        raise ValueError(f"Pod '{name}' not found")
    config.active = name
    save_config(config)


def add_pod(name: str, pod: Pod) -> None:
    """Add a pod to the config."""
    config = load_config()
    config.pods[name] = pod
    if config.active is None:
        config.active = name
    save_config(config)


def remove_pod(name: str) -> None:
    """Remove a pod from the config."""
    config = load_config()
    config.pods.pop(name, None)
    if config.active == name:
        config.active = None
    save_config(config)
