import json
import asyncio
import subprocess
from config import settings


class LoadBalancer:
    def __init__(self):
        self._rr_index = 0
        self._conn_counts: dict[str, int] = {n["id"]: 0 for n in settings.edge_nodes}

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

    def select_nodes(self, strategy: str | None = None):
        nodes = self.get_enabled_nodes()
        strategy = strategy or settings.load_balance_strategy

        if strategy == "round-robin":
            if not nodes:
                return []
            chosen = nodes[self._rr_index % len(nodes)]
            self._rr_index += 1
            return [chosen]
        elif strategy == "weighted":
            total = sum(n["weight"] for n in nodes)
            if total == 0:
                return nodes
            import random
            r = random.uniform(0, total)
            cumulative = 0
            for n in nodes:
                cumulative += n["weight"]
                if r <= cumulative:
                    return [n]
            return [nodes[-1]]
        elif strategy == "least-conn":
            sorted_nodes = sorted(nodes, key=lambda n: self._conn_counts.get(n["id"], 0))
            return [sorted_nodes[0]]
        return nodes

    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()