from __future__ import annotations

from dataclasses import asdict, dataclass, fields, field
from ipaddress import IPv4Address, IPv4Interface, IPv4Network
from typing import Literal, TypeVar, ClassVar, AsyncIterator
from enum import Enum, auto
import json
import aiofiles
from pathlib import Path
from flux_networking_shared.helpers import exec_binary, ExecBinaryError
from flux_networking_shared.log import log
from contextlib import asynccontextmanager

T = TypeVar("T")


@dataclass
class SourceAddress:
    ip: IPv4Address
    validated: bool = False


@dataclass
class Route:
    dst: IPv4Network
    gateway: str | None
    scope: Literal["universe", "link"]
    proto: Literal["static", "kernel"]
    link: str
    prefsrc: IPv4Address | None

    @classmethod
    def from_dict(cls, data: dict) -> Route:
        fieldSet = {f.name for f in fields(cls) if f.init}
        filtered = {k: v for k, v in data.items() if k in fieldSet}

        filtered["dst"] = IPv4Network(filtered["dst"])

        if filtered["prefsrc"]:
            filtered["prefsrc"] = IPv4Address(filtered["prefsrc"])

        return cls(**filtered)

    @property
    def is_default(self) -> bool:
        return self.dst.with_prefixlen == "0.0.0.0/0" and bool(self.gateway)

    @property
    def data_column(self) -> list[str]:
        return [field.name for field in fields(self)]

    @property
    def data_row(self) -> list[str]:
        return [getattr(self, field.name) for field in fields(self)]

    def __hash__(self) -> int:
        return hash(self.dst) + hash(self.link) + hash(self.gateway)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Route):
            return False

        return (
            self.dst == other.dst
            and self.link == other.link
            and self.gateway == other.gateway
        )


class InterfaceState(Enum):
    UP = auto()
    DOWN = auto()


@dataclass
class InterfaceShapingPolicy:
    interface: str
    rate: Literal[35, 75, 135, 250] | None = None

    @classmethod
    def from_dict(cls, data: dict) -> InterfaceShapingPolicy:
        name = next(iter(data))
        payload = data[name]

        try:
            in_limit_raw: str = payload["incoming"]["protocol=ip"]["rate"]
        except KeyError:
            return cls(name, None)

        # try:
        #     out_limit_raw: str = payload["outgoing"]["protocol=ip"]["rate"]
        # except KeyError:
        #     return None

        try:
            in_limit = int(in_limit_raw.removesuffix("Mbps"))
        except ValueError:
            return cls(name, None)

        # out_limit = out_limit_raw.split(":")[0]

        return cls(name, in_limit)

    @property
    def rate_with_unit(self) -> str:
        return f"{self.rate}Mbps"

    @property
    def is_shaped(self) -> bool:
        return bool(self.rate)

    def as_dict(self) -> dict:
        default: dict[str, dict] = {"outgoing": {}, "incoming": {}}

        if not self.rate:
            return default

        return {
            "outgoing": {
                "protocol=ip": {
                    "limit": 1000,
                    "rate": self.rate_with_unit,
                }
            },
            "incoming": {
                "protocol=ip": {
                    "limit": 1000,
                    "rate": self.rate_with_unit,
                }
            },
        }


@dataclass
class FluxShapingPolicy:
    policy_file: Path = Path("/mnt/root/config/flux-shaping.json")
    interfaces: dict[str, InterfaceShapingPolicy] = field(default_factory=dict)
    raw: str = ""

    @classmethod
    async def from_policy_file(cls) -> FluxShapingPolicy:
        try:
            async with aiofiles.open(FluxShapingPolicy.policy_file, "r") as f:
                data = await f.read()
        except FileNotFoundError:
            data = None

        if not data:
            return cls()

        return cls.from_policy_str(data)

    @classmethod
    def from_policy_str(cls, data: str) -> FluxShapingPolicy:
        try:
            raw_interfaces: dict = json.loads(data)
        except json.JSONDecodeError:
            return cls()

        interfaces = {
            k: InterfaceShapingPolicy.from_dict({k: v})
            for k, v in raw_interfaces.items()
        }

        return cls(interfaces=interfaces, raw=data)

    def __str__(self) -> str:
        as_dict = {x.interface: x.as_dict() for x in self.interfaces.values()}

        return json.dumps(as_dict)

    @asynccontextmanager
    async def override_ratelimit(self) -> AsyncIterator[None]:
        try:
            for interface, policy in self.interfaces.items():
                if policy.rate:
                    await self.clear_interface_policy(interface)

            yield

        finally:
            for interface, policy in self.interfaces.items():
                if policy.rate:
                    await self.set_interface_policy(
                        interface, policy.rate, force=True
                    )

    async def set_interface_policy(
        self,
        name: str,
        bandwidth_limit: Literal[35, 75, 135, 250] | None,
        force: bool = False,
    ) -> bool:
        """Sets the tc policy on an interface

        Args:
            name (str): the interface name
            bandwidth_limit (Literal[35, 75, 135, 250] | None): The limit to set

        Returns:
            bool: If the policy was changed
        """

        policy = self.get_interface_policy(name)

        if not force and policy.rate == bandwidth_limit:
            return False

        if not bandwidth_limit:
            await self.clear_interface_policy(name)

        policy.rate = bandwidth_limit
        self.interfaces[name] = policy

        await self.sync_policy()

        return True

    def get_interface_policy(self, name: str) -> InterfaceShapingPolicy:
        policy = self.interfaces.get(name)

        if not policy:
            policy = InterfaceShapingPolicy(name)

        return policy

    async def clear_interface_policy(self, name: str) -> None:
        # We shouldn't need to do this. It should clear when importing the config,
        # but it doesn't. I have a bug open.

        cmd = ["tcdel", name, "--all"]

        try:
            await exec_binary(cmd)
        except ExecBinaryError:
            log.exception("Unable to clear interface shaping policy")

    async def sync_policy(self, write_only: bool = False) -> bool:
        async with aiofiles.open(self.policy_file, "w") as f:
            await f.write(str(self))

        if write_only:
            return True

        cmd = ["tcset", str(self.policy_file), "--import-setting"]

        try:
            await exec_binary(cmd)
        except ExecBinaryError:
            log.exception("Unable to sync traffic control policy")

        return True

    async def rollback_policy(self) -> bool:
        prior_policy = FluxShapingPolicy.from_policy_str(self.raw)

        for name, policy in self.interfaces.items():
            if name not in prior_policy.interfaces:
                await self.clear_interface_policy(name)

        for name, policy in prior_policy.interfaces.items():
            if not policy.is_shaped:
                await self.clear_interface_policy(name)

        synced = await prior_policy.sync_policy()

        return synced


@dataclass
class NetworkInterface:
    type: str
    name: str
    state: InterfaceState
    ifindex: int
    mac_addr: str
    ip_allocation: str
    network_vendor: str
    network_type: str
    address: IPv4Interface | None = None
    vlan: str | None = None
    parent_ifindex: int | None = None
    bridge_id: str | None = None
    asdict = asdict

    @classmethod
    def from_probe_data(cls, data: dict, connected: bool) -> NetworkInterface:
        netlink: dict = data["netlink_data"]
        udev: dict = data["udev_data"]

        int_type = data["type"]
        int_name = netlink["name"]
        int_state = InterfaceState.UP if connected else InterfaceState.DOWN
        int_ifindex = netlink["ifindex"]
        int_mac_addr = udev["attrs"]["address"]

        int_network_vendor = udev.get("ID_VENDOR_FROM_DATABASE", None)
        int_network_type = udev.get("ID_MODEL_FROM_DATABASE", None)
        int_vlan = netlink.get("vlan_id", None)
        int_parent_ifindex = netlink.get("vlan_link", None)
        int_bridge_id = data["bridge"]["options"].get("bridge_id", None)
        int_designated_bridge = data["bridge"]["options"].get(
            "designated_bridge", None
        )

        bridge_id = int_bridge_id if int_bridge_id else int_designated_bridge

        int_details = next(
            filter(
                lambda x: x["family"] == 2,
                data["addresses"],
            ),
            None,
        )

        address = IPv4Interface(int_details["address"]) if int_details else None
        source = int_details.get("source", None) if int_details else None
        ip_allocation = source if source else "disabled"

        return cls(
            int_type,
            int_name,
            int_state,
            int_ifindex,
            int_mac_addr,
            ip_allocation,
            int_network_vendor,
            int_network_type,
            address,
            int_vlan,
            int_parent_ifindex,
            bridge_id,
        )

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, NetworkInterface):
            return False

        return self.asdict() == other.asdict()


# class RouteGroup(list[Route]):
#     @classmethod
#     def from_dict(cls, data: dict) -> RouteGroup:
#         parsed = [Route.from_dict(x) for x in data if x["family"] == 2]
#         return cls(parsed)

#     @property
#     def default_route(self) -> Route | None:
#         return next(filter(lambda x: x.is_default, self), None)


# class NetworkInterfaceGroup(list[NetworkInterface]):
#     @classmethod
#     def from_dict(cls, data: dict) -> NetworkInterfaceGroup:
#         parsed = [NetworkInterface.from_dict(x) for x in data]

#         # Just filter eth interfaces
#         eth_ints = [
#             x for x in parsed if x.type == "eth" and x.bridge_id is None
#         ]
#         # rest = [x for x in parsed if x.type != "eth"]
#         return cls(eth_ints)


class NetworkInterfaceGroup(dict[str, NetworkInterface]):
    @property
    def index_to_name_map(self) -> dict[int, str]:
        return {v.ifindex: k for k, v in self.items()}

    def changed(self, target: NetworkInterface) -> bool:
        interface = self.get(target.name)

        if not interface:
            return True

        return interface != target


# @dataclass
# class NetworkInfo:
#     routes: RouteGroup
#     interfaces: NetworkInterfaceGroup

#     @classmethod
#     def from_probe_dict(cls, data: dict) -> NetworkInfo:
#         network = data.get("network", {})
#         route_data = network.get("routes", [])
#         interface_data = network.get("links", [])

#         routes = RouteGroup.from_dict(route_data)
#         interfaces = NetworkInterfaceGroup.from_dict(interface_data)

#         return cls(routes, interfaces)

#     @property
#     def primary_interface(self) -> NetworkInterface:
#         default_ifindex = self.routes.default_route.ifindex

#         return next(
#             filter(lambda x: x.ifindex == default_ifindex, self.interfaces),
#             None,
#         )
