import hashlib
import logging
import random
from collections import Counter
from collections.abc import Iterable
from types import SimpleNamespace
from typing import Any, Optional, TypedDict

import attrs
import ray
from kubernetes import client

from geneva.config import ConfigBase
from geneva.runners.kuberay.client import KuberayClients

_LOG = logging.getLogger(__name__)

KUBERAY_API_GROUP = "ray.io"  # API group for KubeRay RayCluster CRD
KUBERAY_API_VERSION = "v1"  # CRD version
KUBERAY_API_GROUP_VERSION = f"{KUBERAY_API_GROUP}/{KUBERAY_API_VERSION}"
KUBERAY_JOB_API_KIND = "RayJob"
KUBERAY_JOB_API_NAME = "rayjobs"
KUBERAY_CLUSTER_PLURAL = "rayclusters"  # lowercase, plural form of the kind
GENEVA_NAMESPACE = "geneva"  # namespace where your clusters live


@attrs.define
class KuberayConfig(ConfigBase):
    checkpoint_store: str = attrs.field()
    ray_version: str = attrs.field()
    namespace: str = attrs.field(default="lancedb")
    worker_min_replicas: int = attrs.field(default=0)
    worker_max_replicas: int = attrs.field(default=10)

    @classmethod
    def name(cls) -> str:
        return "kuberay"


def generate_job_name(
    db_uri: str,
    table_name: str,
    column: str,
) -> str:
    db_name = db_uri.split("/")[-1]
    seed = random.randint(0, 1000000)
    job_name = f"ray-geneva-{db_name[:6]}-{table_name[6]}-{column[:6]}-{hashlib.md5(str(seed).encode()).hexdigest()[:6]}"  # noqa E501
    return job_name


Filter = tuple[str, str, str]  # e.g. ("class_name", "=", "ApplierActor")


def list_actors_all(
    *,
    filters: Iterable[Filter] | None = None,
    detail: bool = False,
    address: str = "auto",
    page_size: int = 1000,
) -> list[SimpleNamespace]:
    """
    Paginate Ray State API 'actors' until exhaustion.
    Returns a list of objects with attribute access (actor_id, state, class_name, ...).
    """

    # this requires connecting to the ray dashboard
    from ray.util.state import StateApiClient

    client = StateApiClient(address=address)
    out: list[SimpleNamespace] = []
    cursor = None

    while True:
        page = client.list(
            resource="actors",
            filters=list(filters or []),
            limit=page_size,  # server page size (not a hard cap)
            detail=detail,  # keep payload light unless you need details
            start=cursor,
            raise_on_missing_output=False,  # don’t raise on server truncation
        )

        for item in page.result or []:
            ns = SimpleNamespace()
            for k, v in item.items():
                setattr(ns, k, v)
            out.append(ns)

        if not page.next_page_token:
            break
        cursor = page.next_page_token

    return out


def _ray_status() -> dict[str, any]:
    """Check the status of the compute cluster
    Checks ray actors, ray and k8s to show progress
    """
    status: dict[str, any] = {
        "cnt_ray_nodes": 0,
        "cnt_geneva_workers_active": 0,
        "cnt_geneva_workers_pending": 0,
    }

    try:
        # of machines in ray cluster
        nodes = ray.nodes()
        live_nodes = [n for n in nodes if n["Alive"]]
        status["cnt_ray_nodes"] = len(live_nodes)
    except Exception as e:
        _LOG.exception("Problem listing ray nodes")
        status["ray_nodes_error"] = e

    try:
        # we want the ALIVE and PENDING_CREATION actors, but filter syntax does not
        # support OR.  instaed we and the non existnace of the other states.

        actors = list_actors_all(
            filters=[("state", "!=", "DEAD"), ("class_name", "=", "ApplierActor")],
            detail=True,
        )
        live_applier_actors = [
            a.actor_id
            for a in actors
            if a.class_name == "ApplierActor" and a.state == "ALIVE"
        ]
        status["cnt_geneva_workers_active"] = len(live_applier_actors)

        pending_applier_actors = [
            a.actor_id
            for a in actors
            if a.class_name == "ApplierActor"
            and a.state in ("PENDING_CREATION", "DEPENDENCIES_UNREADY", "RESTARTING")
        ]
        status["cnt_geneva_workers_pending"] = len(pending_applier_actors)
    except Exception as e:
        _LOG.exception("Problem listing actors")
        status["geneva_workers_error"] = e

    _LOG.debug(f"cluster status: {status}")
    return status


class PodStatus(TypedDict):
    name: str
    phase: str  # Pending | Running | Succeeded | Failed | Unknown
    ready: bool
    node_type: str | None  # head | worker | None (if unlabeled)
    node_name: str | None
    waiting_reasons: Counter[str]  # main containers
    init_waiting_reasons: Counter[str]  # init containers
    pulling_count: int  # total containers currently "Pulling..."
    gpu_requested: bool
    node_is_gpu: bool | None


# Which extended resources count as "GPU"
_GPU_KEYS = ("nvidia.com/gpu", "amd.com/gpu", "intel.com/gpu")


def _qty_gt_zero(q: Any) -> bool:
    try:
        # K8s quantities for extended resources are ints (e.g. "1")
        return int(q) > 0
    except Exception:
        return False


def _container_wants_gpu(c: client.V1Container) -> bool:
    r = c.resources or client.V1ResourceRequirements()
    for d in (r.limits or {}), (r.requests or {}):
        for k in _GPU_KEYS:
            if k in d and _qty_gt_zero(d[k]):
                return True
    return False


def _pod_requests_gpu(pod: client.V1Pod) -> bool:
    spec = pod.spec
    if not spec:
        return False
    for c in spec.init_containers or []:
        if _container_wants_gpu(c):
            return True
    return any(_container_wants_gpu(c) for c in spec.containers or [])


def _node_is_gpu(
    clients: KuberayClients, node_name: str, cache: dict[str, bool]
) -> bool:
    if node_name in cache:
        return cache[node_name]
    try:
        n = clients.core_api.read_node(node_name)
        alloc = n.status.allocatable or {}
        has_gpu = any(_qty_gt_zero(alloc.get(k)) for k in _GPU_KEYS)
    except Exception:
        has_gpu = False
    cache[node_name] = has_gpu
    return has_gpu


def k8s_status(
    clients: KuberayClients,
    namespace: str = "geneva",
) -> list[PodStatus]:
    # List all pods in the namespace
    pods = clients.core_api.list_namespaced_pod(namespace)

    # Filter Ray pods by label (adjust depending on your RayCluster setup)
    ray_pods = [
        pod
        for pod in pods.items
        if pod.metadata
        and pod.metadata.labels
        and "ray.io/node-type" in pod.metadata.labels
    ]

    # Summarize pod phases and conditions
    node_gpu_cache: dict[str, bool] = {}
    out: list[PodStatus] = []
    for pod in ray_pods:
        phase: str = pod.status.phase or "Unknown"
        conds: list[client.V1PodCondition] = pod.status.conditions or []
        cond_map = {c.type: c.status for c in conds}

        waiting: Counter[str] = Counter()
        init_waiting: Counter[str] = Counter()
        pulling_count = 0

        # Init containers
        for cs in pod.status.init_container_statuses or []:
            st = cs.state
            if st and st.waiting:
                reason = st.waiting.reason or "Waiting"
                init_waiting[reason] += 1
                if reason.lower().startswith("pull"):
                    pulling_count += 1

        # Main containers
        for cs in pod.status.container_statuses or []:
            st = cs.state
            if st and st.waiting:
                reason = st.waiting.reason or "Waiting"
                waiting[reason] += 1
                if reason.lower().startswith("pull"):
                    pulling_count += 1

        node_name = pod.spec.node_name if pod.spec else None
        node_is_gpu = (
            _node_is_gpu(clients, node_name, node_gpu_cache) if node_name else None
        )

        out.append(
            PodStatus(
                name=pod.metadata.name,
                phase=phase,
                ready=(cond_map.get("Ready") == "True"),
                node_type=pod.metadata.labels.get("ray.io/node-type"),
                node_name=pod.spec.node_name if pod.spec else None,
                waiting_reasons=waiting,
                init_waiting_reasons=init_waiting,
                pulling_count=pulling_count,
                gpu_requested=_pod_requests_gpu(pod),
                node_is_gpu=node_is_gpu,
            )
        )
    return out


def _node_ready_split_for(
    clients: KuberayClients, pod_statuses: list[PodStatus]
) -> tuple[int, int, int, int]:
    """
    Returns (gpu_ready, gpu_notready, cpu_ready, cpu_notready)
    for nodes that host any of the provided pods.
    """
    nodes_seen: dict[str, bool] = {}  # node_name -> is_gpu
    ready_map: dict[str, bool] = {}  # node_name -> Ready True/False

    # collect unique node names from these pods
    node_names = {s["node_name"] for s in pod_statuses if s["node_name"]}

    # read node objects once
    for name in node_names:
        try:
            n = clients.core_api.read_node(name)
            conds = {c.type: c.status for c in (n.status.conditions or [])}
            is_ready = conds.get("Ready") == "True"
            alloc = n.status.allocatable or {}
            is_gpu = any(_qty_gt_zero(alloc.get(k)) for k in _GPU_KEYS)
            nodes_seen[name] = is_gpu
            ready_map[name] = is_ready
        except Exception:  # noqa PERF203
            # if we can’t read, treat as not ready & cpu
            nodes_seen[name] = False
            ready_map[name] = False

    gpu_ready = gpu_notready = cpu_ready = cpu_notready = 0
    for name, is_gpu in nodes_seen.items():
        if ready_map.get(name, False):
            if is_gpu:
                gpu_ready += 1
            else:
                cpu_ready += 1
        else:
            if is_gpu:
                gpu_notready += 1
            else:
                cpu_notready += 1
    return gpu_ready, gpu_notready, cpu_ready, cpu_notready


def kuberay_cluster_status(
    clients: KuberayClients, namespace: str, name: str
) -> Optional[dict[str, Any]]:
    try:
        obj = clients.custom_api.get_namespaced_custom_object(
            group="ray.io",
            version="v1",
            namespace=namespace,
            plural="rayclusters",
            name=name,
        )
        return obj.get("status")
    except Exception:
        return None


class WorkerGroupBrief(TypedDict):
    name: str
    desired: int
    ready: int


class KuberaySummary(TypedDict):
    phase: (
        str  # "nodes-provisioning" | "pods-creating" | "cluster-cold" | "cluster-warm"
    )
    phase_idx: int  # monotonically increasing index for your progress bar
    pending: int
    running: int
    workers_ready: int
    waits_top3: list[tuple[str, int]]
    pulling: int
    total_pods: int

    # kuberay
    kr_state: str | None
    kr_desired_workers: int | None
    kr_available_workers: int | None
    kr_last_update: str | None
    kr_scaling: str | None  # "up" | "down" | "steady" | None
    kr_groups: list[WorkerGroupBrief]  # per workger-group ready/desired
    kr_last_condition: tuple[str, str] | None  # reason, type if present

    # pod splits
    pods_gpu_running: int
    pods_gpu_pending: int
    pods_cpu_running: int
    pods_cpu_pending: int
    workers_ready_gpu: int
    workers_ready_cpu: int

    # node splits (only nodes hosting this cluster’s pods)
    nodes_gpu_ready: int
    nodes_gpu_notready: int
    nodes_cpu_ready: int
    nodes_cpu_notready: int


def _get_int(d: dict[str, Any], *keys: str) -> Optional[int]:
    for k in keys:
        v = d.get(k)
        if isinstance(v, int):
            return v
    return None


def _parse_worker_groups(status: dict[str, Any]) -> list[WorkerGroupBrief]:
    out: list[WorkerGroupBrief] = []
    # Newer KubeRay surfaces workerGroupStatuses; older embeds in spec/status
    # differently
    wgs = status.get("workerGroupStatuses") or []
    for item in wgs:
        name = item.get("groupName") or item.get("name") or "worker"
        desired = _get_int(item, "desiredReplicas", "replicas", "desired") or 0
        ready = _get_int(item, "readyReplicas", "availableReplicas", "ready") or 0
        out.append(WorkerGroupBrief(name=name, desired=desired, ready=ready))
    return out


def _parse_scaling(status: dict[str, Any]) -> Optional[str]:
    # Heuristic: desired vs available (or ready) tells the story
    desired = _get_int(status, "desiredWorkerReplicas", "desiredWorkerCount")
    avail = _get_int(status, "availableWorkerReplicas", "availableWorkerCount")
    if desired is None or avail is None:
        return None
    if desired > avail:
        return "up"
    if desired < avail:
        return "down"
    return "steady"


def _last_condition(status: dict[str, Any]) -> Optional[tuple[str, str]]:
    # Surfaces *some* operator condition, if present
    conds = status.get("conditions") or []
    if not isinstance(conds, list) or not conds:
        return None
    last = conds[-1]
    reason = str(last.get("reason") or "?")
    typ = str(last.get("type") or "?")
    return (reason, typ)


def summarize_k8s(
    clients: KuberayClients,
    statuses: list[PodStatus],
    cluster_name: str,
    namespace: str = "geneva",
) -> KuberaySummary:
    phases = Counter(s["phase"] for s in statuses)
    waits = Counter()
    pulling = 0
    workers_ready = 0
    workers_ready_gpu = 0
    workers_ready_cpu = 0
    pods_gpu_running = pods_gpu_pending = 0
    pods_cpu_running = pods_cpu_pending = 0

    for s in statuses:
        waits.update(s["waiting_reasons"])
        waits.update(s["init_waiting_reasons"])
        pulling += s["pulling_count"]
        if s["node_type"] == "worker" and s["phase"] == "Running" and s["ready"]:
            workers_ready += 1
            if s.get("gpu_requested"):
                workers_ready_gpu += 1
            else:
                workers_ready_cpu += 1

        # pod splits by GPU/CPU and phase
        is_gpu = bool(s.get("gpu_requested"))
        if s["phase"] == "Running":
            if is_gpu:
                pods_gpu_running += 1
            else:
                pods_cpu_running += 1
        elif s["phase"] == "Pending":
            if is_gpu:
                pods_gpu_pending += 1
            else:
                pods_cpu_pending += 1

    # Phase selection focused on “cold start” (= no workers ready yet)
    if phases["Pending"] > 0 and any(
        k.lower().startswith(("unschedulable", "waiting")) for k in waits
    ):
        phase, phase_idx = "nodes-provisioning", 1
    elif pulling > 0 or any(k.lower().startswith("containercreating") for k in waits):
        phase, phase_idx = "pods-creating", 2
    elif workers_ready == 0:
        phase, phase_idx = "cluster-cold", 3
    else:
        phase, phase_idx = "cluster-warm", 4

    kr_status = kuberay_cluster_status(clients, namespace, cluster_name)
    if kr_status is None:
        _LOG.warning(
            f"Could not fetch KubeRay status for {namespace}/{cluster_name}..."
        )
        kr_status = {}

    desired = _get_int(kr_status, "desiredWorkerReplicas", "desiredWorkerCount")
    avail = _get_int(kr_status, "availableWorkerReplicas", "availableWorkerCount")
    groups = _parse_worker_groups(kr_status)
    scaling = _parse_scaling(kr_status)
    last_cond = _last_condition(kr_status)

    # Node readiness split (only nodes hosting our Ray pods)
    nodes_gpu_ready, nodes_gpu_notready, nodes_cpu_ready, nodes_cpu_notready = (
        _node_ready_split_for(clients, statuses)
    )

    def _int_or_none(v: Any) -> int | None:
        return int(v) if isinstance(v, int) else None

    return KuberaySummary(
        phase=phase,
        phase_idx=phase_idx,
        pending=phases.get("Pending", 0),
        running=phases.get("Running", 0),
        workers_ready=workers_ready,
        waits_top3=waits.most_common(3),
        pulling=pulling,
        total_pods=len(statuses),
        kr_state=kr_status.get("state"),
        kr_desired_workers=_int_or_none(desired),
        kr_available_workers=_int_or_none(avail),
        kr_last_update=kr_status.get("lastUpdateTime"),
        kr_scaling=scaling,
        kr_groups=groups,
        kr_last_condition=last_cond,
        # NEW pod splits
        pods_gpu_running=pods_gpu_running,
        pods_gpu_pending=pods_gpu_pending,
        pods_cpu_running=pods_cpu_running,
        pods_cpu_pending=pods_cpu_pending,
        workers_ready_gpu=workers_ready_gpu,
        workers_ready_cpu=workers_ready_cpu,
        # NEW node splits
        nodes_gpu_ready=nodes_gpu_ready,
        nodes_gpu_notready=nodes_gpu_notready,
        nodes_cpu_ready=nodes_cpu_ready,
        nodes_cpu_notready=nodes_cpu_notready,
    )


def summarize_kuberay_status(
    clients: KuberayClients, namespace: str | None, cluster_name: str | None
) -> KuberaySummary:
    k8s_statuses = k8s_status(clients, namespace=namespace)
    k8s_summary = summarize_k8s(
        clients, k8s_statuses, cluster_name=cluster_name, namespace=namespace
    )
    return k8s_summary
