import ipaddress
import time
from config import settings


class LoadBalancer:
    def __init__(self):
        self._rr_index = 0
        self._conn_counts: dict[str, int] = {}
        self._node_health: dict[str, dict] = {}

    def get_enabled_nodes(self):
        return [n for n in settings.edge_nodes if n.get("enabled", True)]

    def get_node_by_id(self, node_id: str):
        for n in settings.edge_nodes:
            if n["id"] == node_id:
                return n
        return None

    def select_best_node(self, client_ip: str | None = None):
        nodes = self.get_enabled_nodes()
        if not nodes:
            return None

        if client_ip:
            best = self._find_subnet_match(client_ip, nodes)
            if best:
                return best

        strategy = settings.load_balance_strategy
        if strategy == "round-robin":
            online = [n for n in nodes if self.is_online(n["id"])]
            if not online:
                online = nodes
            chosen = online[self._rr_index % len(online)]
            self._rr_index += 1
            return chosen
        elif strategy == "weighted":
            return self._weighted_select(nodes)
        elif strategy == "least-conn":
            return self._least_conn_select(nodes)
        else:
            return self._weighted_select(nodes)

    def _find_subnet_match(self, client_ip: str, nodes: list[dict]):
        try:
            ip = ipaddress.ip_address(client_ip)
        except ValueError:
            return None

        best = None
        best_prefix = -1

        for node in nodes:
            if not self.is_online(node["id"]):
                continue

            subnet_str = node.get("subnet", "0.0.0.0/0")
            if not subnet_str:
                continue

            try:
                subnet = ipaddress.ip_network(subnet_str, strict=False)
            except ValueError:
                continue

            if ip in subnet:
                if subnet.prefixlen > best_prefix:
                    best_prefix = subnet.prefixlen
                    best = node

        return best

    def _weighted_select(self, nodes: list[dict]):
        online = [n for n in nodes if self.is_online(n["id"])]
        if not online:
            online = nodes
        total = sum(n["weight"] for n in online)
        if total == 0:
            return online[0] if online else None
        import random
        r = random.uniform(0, total)
        cumulative = 0
        for n in online:
            cumulative += n["weight"]
            if r <= cumulative:
                return n
        return online[-1]

    def _least_conn_select(self, nodes: list[dict]):
        online = [n for n in nodes if self.is_online(n["id"])]
        if not online:
            online = nodes
        sorted_nodes = sorted(online, key=lambda n: self._conn_counts.get(n["id"], 0))
        return sorted_nodes[0] if sorted_nodes else None

    def report_health(self, node_id: str, metrics: dict):
        self._node_health[node_id] = {
            "last_seen": time.time(),
            "online": True,
            **metrics,
        }

    def is_online(self, node_id: str) -> bool:
        h = self._node_health.get(node_id)
        if h is None:
            if any(n["id"] == node_id and n.get("is_origin") for n in settings.edge_nodes):
                return True
            return False
        if time.time() - h.get("last_seen", 0) > 30:
            return False
        return h.get("online", False)

    def get_health_summary(self):
        result = {}
        for node in settings.edge_nodes:
            nid = node["id"]
            h = self._node_health.get(nid)
            result[nid] = {
                "name": node["name"],
                "host": node["host"],
                "online": self.is_online(nid),
                "weight": node["weight"],
                "enabled": node.get("enabled", True),
                "subnet": node.get("subnet"),
                "is_origin": node.get("is_origin", False),
                "last_seen": h.get("last_seen") if h else None,
                "metrics": {
                    "cpu_percent": h.get("cpu_percent"),
                    "memory_percent": h.get("memory_percent"),
                    "bandwidth_mbps": h.get("bandwidth_mbps"),
                    "active_connections": h.get("active_connections"),
                    "hls_readers": h.get("hls_readers"),
                } if h else None,
            }
        return result

    def increment_connection(self, node_id: str):
        self._conn_counts[node_id] = self._conn_counts.get(node_id, 0) + 1

    def decrement_connection(self, node_id: str):
        if self._conn_counts.get(node_id, 0) > 0:
            self._conn_counts[node_id] -= 1


lb = LoadBalancer()