# DOCS_REF:
# - docs/user/reference/env-vars.md
# - docs/user/02-quickstart-distributed.md
# DOCS_SYNC:
# - If cluster membership env parsing or locality metadata changes, update the
#   env-vars reference and distributed quickstart notes.
# DESIGN_GUARDRAIL (2026-02-16):
# - ClusterView membership is node-centric (machine/VM level), not container-centric.
# - host_id/uds_path/shm_path/transport_caps are node endpoint metadata, not
#   schedulable membership identity.
# - Containers/workers may affect placement only via node-level summary signals
#   (for example resource_summary worker counts), never as membership rows.
# - io-edge hosts publish resources via URI/registry paths; they do not become
#   compute membership rows in ClusterView.
# - Freeze reaffirmed on 2026-02-18 after coordinator-drift regression:
#   issues/topic-system-end-to-end-container-event-coordinator-drift.md
# - Source of truth for this boundary:
#   issues/cluster-view-crdt-gossip.md
#   issues/flownet-io-edge-host-auto-resource-publication.md

from __future__ import annotations

import hashlib
import os
import threading
import time
from dataclasses import dataclass, replace
from typing import Any, Iterable, Mapping


ENV_CLUSTER_NODES = "FLOWNET_CLUSTER_NODES"
ENV_NODE_ZONE = "FLOWNET_NODE_ZONE"
ENV_NODE_REGION = "FLOWNET_NODE_REGION"

CLUSTER_VIEW_SCHEMA_V0 = "v0"
CLUSTER_VIEW_SCHEMA_V1 = "v1"
DEGRADED_HEALTH_STATES = frozenset({"down", "dead", "unhealthy", "offline", "draining"})


@dataclass(frozen=True)
class ClusterNode:
    node_id: str
    address: str
    zone: str | None = None
    region: str | None = None
    health: str = "unknown"
    last_seen: float | None = None
    resource_summary: dict[str, object] | None = None
    host_id: str | None = None
    uds_path: str | None = None
    shm_path: str | None = None
    transport_caps: tuple[str, ...] | None = None

    def as_dict(self) -> dict[str, object]:
        # Keep node metadata compact and scheduler-oriented. Do not emit
        # per-container membership internals here.
        data: dict[str, object] = {
            "node_id": self.node_id,
            "address": self.address,
            "zone": self.zone,
            "region": self.region,
            "health": self.health,
            "last_seen": self.last_seen,
            "resource_summary": dict(self.resource_summary or {}),
        }
        if self.host_id:
            data["host_id"] = self.host_id
        if self.uds_path:
            data["uds_path"] = self.uds_path
        if self.shm_path:
            data["shm_path"] = self.shm_path
        if self.transport_caps:
            data["transport_caps"] = list(self.transport_caps)
        return data


class ClusterView:
    def __init__(self, nodes: Iterable[ClusterNode], local_node_id: str):
        self._lock = threading.RLock()
        self._nodes: dict[str, ClusterNode] = {node.node_id: node for node in nodes}
        if local_node_id not in self._nodes:
            raise ValueError("local_node_id must exist in cluster view")
        self._local_node_id = local_node_id
        self._tombstones: dict[str, int] = {}
        self._node_versions: dict[str, int] = {
            node.node_id: _clock_ms_from_seconds(node.last_seen) for node in nodes
        }
        self._node_source_clocks: dict[str, dict[str, int]] = {}
        self._tombstone_source_clocks: dict[str, dict[str, int]] = {}
        self._view_epoch = 0

    def local_node_id(self) -> str:
        return self._local_node_id

    def list_nodes(self) -> list[ClusterNode]:
        with self._lock:
            return list(self._nodes.values())

    def iter_nodes(self) -> Iterable[ClusterNode]:
        return iter(self.list_nodes())

    def get_node(self, node_id: str) -> ClusterNode:
        with self._lock:
            if node_id not in self._nodes:
                raise KeyError(f"Unknown node: {node_id}")
            return self._nodes[node_id]

    def get_node_address(self, node_id: str) -> str:
        return self.get_node(node_id).address

    def get_node_meta(self, node_id: str) -> dict[str, object]:
        return self.get_node(node_id).as_dict()

    def update_local_resource_summary(
        self,
        patch: Mapping[str, object],
        *,
        merge: bool = True,
    ) -> dict[str, object]:
        """
        Update local node resource summary without changing cluster membership.

        This is the canonical path for container/worker lifecycle to influence
        placement signals while preserving node-centric membership freeze.
        """
        if not isinstance(patch, Mapping):
            raise ValueError("patch must be a mapping.")

        with self._lock:
            local = self._nodes[self._local_node_id]
            previous = dict(local.resource_summary or {})
            next_summary = dict(previous) if merge else {}
            for key, value in patch.items():
                normalized_key = str(key or "").strip()
                if not normalized_key:
                    continue
                if value is None:
                    next_summary.pop(normalized_key, None)
                    continue
                next_summary[normalized_key] = value

            if next_summary == previous:
                return {
                    "updated": False,
                    "view_epoch": self._view_epoch,
                    "resource_summary": dict(previous),
                }

            self._nodes[self._local_node_id] = replace(
                local,
                resource_summary=next_summary,
            )
            self._view_epoch += 1
            return {
                "updated": True,
                "view_epoch": self._view_epoch,
                "resource_summary": dict(next_summary),
            }

    def find_node_by_address(self, address: str) -> ClusterNode | None:
        with self._lock:
            for node in self._nodes.values():
                if node.address == address:
                    return node
        return None

    def resolve_node_for_key(self, key: str, affinity_hint: str | None = None) -> ClusterNode:
        with self._lock:
            nodes = sorted(self._nodes.values(), key=lambda node: node.node_id)
        if not nodes:
            raise ValueError("ClusterView has no nodes")
        material = key if affinity_hint is None else f"{key}:{affinity_hint}"
        digest = hashlib.sha256(material.encode("utf-8")).hexdigest()
        idx = int(digest, 16) % len(nodes)
        return nodes[idx]

    def view_epoch(self) -> int:
        with self._lock:
            return self._view_epoch

    def locality_tier(self, node_id: str) -> str:
        node = self.get_node(node_id)
        local = self.get_node(self._local_node_id)
        if node.node_id == local.node_id:
            return "local"
        if local.zone and node.zone and local.zone == node.zone:
            return "zone"
        if local.region and node.region and local.region == node.region:
            return "region"
        return "remote"

    def is_node_healthy(self, node_id: str) -> bool:
        node = self.get_node(node_id)
        return _is_healthy_state(node.health)

    def is_node_schedulable(self, node_id: str) -> bool:
        node = self.get_node(node_id)
        summary = node.resource_summary if isinstance(node.resource_summary, dict) else None
        marker = summary.get("schedulable") if summary else None
        return _is_schedulable_state(marker)

    def estimate_staleness_ms(self, now: float | None = None) -> float:
        """Best-effort staleness based on freshest observed node heartbeat."""
        with self._lock:
            freshest = max(
                (
                    node.last_seen
                    for node in self._nodes.values()
                    if node.last_seen is not None
                ),
                default=None,
            )
        if freshest is None:
            return 0.0
        now_ts = time.time() if now is None else now
        delta = now_ts - freshest
        if delta < 0:
            return 0.0
        return delta * 1000.0

    def snapshot(self, version: str = CLUSTER_VIEW_SCHEMA_V0) -> dict[str, object]:
        with self._lock:
            nodes = sorted(self._nodes.values(), key=lambda node: node.node_id)
            epoch = self._view_epoch
        if version == CLUSTER_VIEW_SCHEMA_V0:
            return {
                "schema_version": CLUSTER_VIEW_SCHEMA_V0,
                "local_node_id": self._local_node_id,
                "view_epoch": epoch,
                "nodes": [node.as_dict() for node in nodes],
            }
        if version != CLUSTER_VIEW_SCHEMA_V1:
            raise ValueError(f"Unsupported cluster view schema version: {version}")

        serialized_nodes: list[dict[str, object]] = []
        healthy_count = 0
        for node in nodes:
            healthy = self.is_node_healthy(node.node_id)
            if healthy:
                healthy_count += 1
            tier = self.locality_tier(node.node_id)
            serialized = node.as_dict()
            serialized["locality"] = {
                "tier": tier,
                "same_zone": tier == "zone",
                "same_region": tier in {"zone", "region"},
            }
            serialized["healthy"] = healthy
            serialized_nodes.append(serialized)

        return {
            "schema_version": CLUSTER_VIEW_SCHEMA_V1,
            "local_node_id": self._local_node_id,
            "view_epoch": epoch,
            "generated_at": time.time(),
            "health_summary": {
                "healthy": healthy_count,
                "degraded": len(nodes) - healthy_count,
                "total": len(nodes),
            },
            "nodes": serialized_nodes,
        }

    def apply_gossip_delta(
        self,
        *,
        membership: dict[str, dict[str, Any]],
        tombstones: dict[str, int] | None = None,
        delta_ts_ms: int | None = None,
        source_node: str | None = None,
        source_delta_seq: int | None = None,
    ) -> dict[str, int]:
        added = 0
        updated = 0
        removed = 0
        ignored = 0
        with self._lock:
            tombstone_ids = set(tombstones or {})
            source_seq = int(source_delta_seq) if source_delta_seq is not None else None
            if tombstones:
                for node_id, ts in tombstones.items():
                    ts_ms = int(ts)
                    if source_node is not None and source_seq is not None:
                        node_seen = _source_seq_for_node(
                            self._node_source_clocks, node_id=node_id, source=source_node
                        )
                        tomb_seen = _source_seq_for_node(
                            self._tombstone_source_clocks, node_id=node_id, source=source_node
                        )
                        if source_seq < node_seen or source_seq <= tomb_seen:
                            ignored += 1
                            continue
                    prev = self._tombstones.get(node_id)
                    if prev is not None and prev > ts_ms:
                        ignored += 1
                        continue
                    if node_id == self._local_node_id:
                        ignored += 1
                        continue
                    self._tombstones[node_id] = ts_ms
                    if source_node is not None and source_seq is not None:
                        _set_source_seq_for_node(
                            self._tombstone_source_clocks,
                            node_id=node_id,
                            source=source_node,
                            seq=source_seq,
                        )
                    current_version = self._node_versions.get(node_id, 0)
                    if ts_ms > current_version:
                        self._node_versions[node_id] = ts_ms
                    if node_id in self._nodes:
                        del self._nodes[node_id]
                        removed += 1

            seen_delta_ts = int(delta_ts_ms) if delta_ts_ms is not None else None
            for node_id, meta in membership.items():
                if not isinstance(meta, dict):
                    ignored += 1
                    continue
                if node_id == self._local_node_id:
                    ignored += 1
                    continue
                if source_node is not None and source_seq is not None:
                    node_seen = _source_seq_for_node(
                        self._node_source_clocks, node_id=node_id, source=source_node
                    )
                    tomb_seen = _source_seq_for_node(
                        self._tombstone_source_clocks, node_id=node_id, source=source_node
                    )
                    if source_seq <= node_seen:
                        ignored += 1
                        continue
                    if source_seq <= tomb_seen and node_id not in tombstone_ids:
                        ignored += 1
                        continue
                tomb = self._tombstones.get(node_id)
                if tomb is not None and seen_delta_ts is None:
                    if source_node is None or source_seq is None:
                        ignored += 1
                        continue
                    tomb_seen = _source_seq_for_node(
                        self._tombstone_source_clocks, node_id=node_id, source=source_node
                    )
                    if source_seq <= tomb_seen:
                        ignored += 1
                        continue
                if tomb is not None and seen_delta_ts is not None and tomb >= seen_delta_ts:
                    ignored += 1
                    continue
                current_version = self._node_versions.get(node_id, 0)
                if seen_delta_ts is not None and seen_delta_ts < current_version:
                    ignored += 1
                    continue
                try:
                    node = _node_from_gossip_entry(node_id=node_id, meta=meta)
                except ValueError:
                    ignored += 1
                    continue
                if node_id in self._nodes:
                    self._nodes[node_id] = node
                    updated += 1
                else:
                    self._nodes[node_id] = node
                    added += 1
                if tomb is not None and seen_delta_ts is not None and seen_delta_ts > tomb:
                    self._tombstones.pop(node_id, None)
                if source_node is not None and source_seq is not None:
                    _set_source_seq_for_node(
                        self._node_source_clocks,
                        node_id=node_id,
                        source=source_node,
                        seq=source_seq,
                    )
                next_version = current_version
                if seen_delta_ts is not None:
                    next_version = max(next_version, seen_delta_ts)
                else:
                    next_version = max(next_version, _clock_ms_from_seconds(node.last_seen))
                self._node_versions[node_id] = next_version

            if any((added, updated, removed)):
                self._view_epoch += 1

            return {
                "added": added,
                "updated": updated,
                "removed": removed,
                "ignored": ignored,
                "epoch": self._view_epoch,
            }

    def age_nodes(
        self,
        *,
        now: float | None = None,
        stale_after_s: float = 15.0,
        remove_after_s: float = 60.0,
        stale_health: str = "stale",
    ) -> dict[str, int]:
        if stale_after_s < 0:
            raise ValueError("stale_after_s must be >= 0")
        if remove_after_s < 0:
            raise ValueError("remove_after_s must be >= 0")
        if remove_after_s < stale_after_s:
            raise ValueError("remove_after_s must be >= stale_after_s")

        now_ts = time.time() if now is None else float(now)
        now_ms = max(0, int(now_ts * 1000.0))
        degraded = 0
        removed = 0
        ignored = 0
        with self._lock:
            for node_id, node in list(self._nodes.items()):
                if node_id == self._local_node_id:
                    continue
                if node.last_seen is None:
                    ignored += 1
                    continue
                age_s = now_ts - float(node.last_seen)
                if age_s < 0:
                    ignored += 1
                    continue
                if age_s >= remove_after_s:
                    prev_tomb = int(self._tombstones.get(node_id, 0))
                    self._tombstones[node_id] = max(prev_tomb, now_ms)
                    current_version = int(self._node_versions.get(node_id, 0))
                    self._node_versions[node_id] = max(current_version, now_ms)
                    del self._nodes[node_id]
                    removed += 1
                    continue
                if age_s >= stale_after_s and node.health != stale_health:
                    self._nodes[node_id] = ClusterNode(
                        node_id=node.node_id,
                        address=node.address,
                        zone=node.zone,
                        region=node.region,
                        health=stale_health,
                        last_seen=node.last_seen,
                        resource_summary=node.resource_summary,
                        host_id=node.host_id,
                        uds_path=node.uds_path,
                        shm_path=node.shm_path,
                        transport_caps=node.transport_caps,
                    )
                    degraded += 1
            if degraded or removed:
                self._view_epoch += 1
            return {
                "degraded": degraded,
                "removed": removed,
                "ignored": ignored,
                "epoch": self._view_epoch,
            }


def build_naive_cluster_view(
    *,
    local_node_id: str,
    local_address: str,
    nodes_env: str | None = None,
    local_schedulable: bool = True,
    local_host_id: str | None = None,
    local_uds_path: str | None = None,
    local_shm_path: str | None = None,
    local_transport_caps: Iterable[str] | None = None,
) -> ClusterView:
    # v0 bootstrap: build node-level membership from env/local runtime metadata.
    # Container lifecycle churn is intentionally out of this membership builder.
    raw = nodes_env if nodes_env is not None else os.environ.get(ENV_CLUSTER_NODES, "")
    nodes: list[ClusterNode] = []
    if raw.strip():
        entries = _split_node_entries(raw)
        for idx, entry in enumerate(entries):
            node = _parse_node_entry(entry, idx)
            nodes.append(node)
    local_idx = next((idx for idx, node in enumerate(nodes) if node.node_id == local_node_id), None)
    if local_idx is None:
        nodes.append(
            ClusterNode(
                node_id=local_node_id,
                address=local_address,
                zone=os.environ.get(ENV_NODE_ZONE),
                region=os.environ.get(ENV_NODE_REGION),
                health="local",
                last_seen=time.time(),
                resource_summary=_merge_resource_summary_schedulable(
                    None,
                    local_schedulable,
                ),
                host_id=local_host_id,
                uds_path=local_uds_path,
                shm_path=local_shm_path,
                transport_caps=_parse_transport_caps(local_transport_caps),
            )
        )
    else:
        local_node = nodes[local_idx]
        nodes[local_idx] = replace(
            local_node,
            resource_summary=_merge_resource_summary_schedulable(
                local_node.resource_summary,
                local_schedulable,
            ),
        )
    return ClusterView(nodes, local_node_id=local_node_id)


def _parse_node_entry(entry: str, idx: int) -> ClusterNode:
    parts = [part.strip() for part in entry.split("|") if part.strip()]
    head = parts[0]
    meta = _parse_meta(parts[1:])
    if "@" in head:
        node_id, address = head.split("@", 1)
    else:
        node_id = meta.get("node_id") or f"node-{idx}"
        address = head
    return ClusterNode(
        node_id=node_id,
        address=address,
        zone=meta.get("zone"),
        region=meta.get("region"),
        health=meta.get("health", "unknown"),
        last_seen=_parse_last_seen(meta.get("last_seen")),
        resource_summary=_parse_resource_summary(meta),
        host_id=meta.get("host_id"),
        uds_path=meta.get("uds_path"),
        shm_path=meta.get("shm_path"),
        transport_caps=_parse_transport_caps(meta.get("transport_caps")),
    )


def _parse_meta(items: Iterable[str]) -> dict[str, str]:
    result: dict[str, str] = {}
    for item in items:
        if "=" not in item:
            continue
        key, value = item.split("=", 1)
        result[key.strip()] = value.strip()
    return result


def _split_node_entries(raw: str) -> list[str]:
    # transport_caps may contain commas; only split on commas that look like a
    # new node entry head.
    parts = [item.strip() for item in raw.split(",") if item.strip()]
    if not parts:
        return []
    entries: list[str] = []
    current = parts[0]
    for part in parts[1:]:
        if _looks_like_node_head(part):
            entries.append(current)
            current = part
            continue
        current = f"{current},{part}"
    entries.append(current)
    return entries


def _looks_like_node_head(part: str) -> bool:
    head = part.split("|", 1)[0].strip()
    if "@" in head:
        return True
    if ":" in head and "=" not in head:
        return True
    return False


def _parse_transport_caps(value: Any) -> tuple[str, ...] | None:
    if value is None:
        return None
    raw_items: list[str]
    if isinstance(value, str):
        raw_items = [item.strip() for item in value.split(",")]
    elif isinstance(value, (list, tuple, set)):
        raw_items = [str(item).strip() for item in value]
    else:
        return None
    normalized: list[str] = []
    for item in raw_items:
        cap = item.lower()
        if not cap or cap in normalized:
            continue
        normalized.append(cap)
    if not normalized:
        return None
    return tuple(normalized)


def _parse_last_seen(value: str | None) -> float | None:
    if not value:
        return None
    try:
        return float(value)
    except ValueError:
        return None


def _clock_ms_from_seconds(value: float | None) -> int:
    if value is None:
        return 0
    try:
        return max(0, int(float(value) * 1000.0))
    except (TypeError, ValueError):
        return 0


def _source_seq_for_node(
    table: dict[str, dict[str, int]],
    *,
    node_id: str,
    source: str,
) -> int:
    per_node = table.get(node_id)
    if per_node is None:
        return -1
    return int(per_node.get(source, -1))


def _set_source_seq_for_node(
    table: dict[str, dict[str, int]],
    *,
    node_id: str,
    source: str,
    seq: int,
) -> None:
    per_node = table.get(node_id)
    if per_node is None:
        table[node_id] = {source: int(seq)}
        return
    prev = int(per_node.get(source, -1))
    if int(seq) > prev:
        per_node[source] = int(seq)


def _parse_resource_summary(meta: dict[str, str]) -> dict[str, object] | None:
    summary: dict[str, object] = {}
    for key in ("cpu", "memory_mb", "gpu", "gpu_mem_mb"):
        if key in meta:
            summary[key] = _parse_number(meta[key])
    if "schedulable" in meta:
        parsed = _coerce_bool(meta["schedulable"])
        summary["schedulable"] = parsed if parsed is not None else meta["schedulable"]
    for key, value in meta.items():
        if key.startswith("res_"):
            summary[key[4:]] = _parse_number(value)
    return summary or None


def _parse_number(value: str) -> Any:
    try:
        if "." in value:
            return float(value)
        return int(value)
    except ValueError:
        return value


def _coerce_bool(value: Any) -> bool | None:
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return bool(value)
    if not isinstance(value, str):
        return None
    normalized = value.strip().lower()
    if normalized in {"1", "true", "yes", "on"}:
        return True
    if normalized in {"0", "false", "no", "off"}:
        return False
    return None


def _merge_resource_summary_schedulable(
    summary: dict[str, object] | None,
    schedulable: Any,
) -> dict[str, object] | None:
    parsed = _coerce_bool(schedulable)
    if parsed is None:
        return dict(summary) if isinstance(summary, dict) else summary
    merged = dict(summary or {})
    merged["schedulable"] = parsed
    return merged


def _is_healthy_state(value: str | None) -> bool:
    if value is None:
        return True
    return value.strip().lower() not in DEGRADED_HEALTH_STATES


def _is_schedulable_state(value: Any) -> bool:
    parsed = _coerce_bool(value)
    if parsed is None:
        return True
    return parsed


def _node_from_gossip_entry(*, node_id: str, meta: dict[str, Any]) -> ClusterNode:
    address = meta.get("address")
    if not isinstance(address, str) or not address:
        raise ValueError("gossip membership entry requires non-empty address.")
    raw_last_seen = meta.get("last_seen")
    last_seen: float | None = None
    if isinstance(raw_last_seen, (int, float)):
        last_seen = float(raw_last_seen)
    zone = meta.get("zone") if isinstance(meta.get("zone"), str) else None
    region = meta.get("region") if isinstance(meta.get("region"), str) else None
    health = meta.get("health")
    if not isinstance(health, str):
        health = "unknown"
    resource_summary = meta.get("resource_summary")
    if not isinstance(resource_summary, dict):
        resource_summary = None
    resource_summary = _merge_resource_summary_schedulable(
        resource_summary,
        meta.get("schedulable"),
    )
    host_id = meta.get("host_id") if isinstance(meta.get("host_id"), str) else None
    uds_path = meta.get("uds_path") if isinstance(meta.get("uds_path"), str) else None
    shm_path = meta.get("shm_path") if isinstance(meta.get("shm_path"), str) else None
    transport_caps = _parse_transport_caps(meta.get("transport_caps"))
    return ClusterNode(
        node_id=node_id,
        address=address,
        zone=zone,
        region=region,
        health=health,
        last_seen=last_seen,
        resource_summary=resource_summary,
        host_id=host_id,
        uds_path=uds_path,
        shm_path=shm_path,
        transport_caps=transport_caps,
    )


__all__ = [
    "ClusterNode",
    "ClusterView",
    "CLUSTER_VIEW_SCHEMA_V0",
    "CLUSTER_VIEW_SCHEMA_V1",
    "build_naive_cluster_view",
]
