from __future__ import annotations

import asyncio
from dataclasses import dataclass, field
from itertools import count
from typing import ClassVar, Literal

import miniupnpc
from yarl import URL
from ipaddress import IPv4Address

from flux_networking_shared.log import log


@dataclass
class UpnpLease:
    host: str
    remaining_s: int
    protocol: str
    description: str = ""


@dataclass
class AddMappingResult:
    port: int
    lease: UpnpLease | None = None

    def as_dict(self) -> dict:
        return {
            "port": self.port,
            "host": self.lease.host if self.lease else "",
            "remaining": self.lease.remaining_s if self.lease else "",
        }


FluxApiPorts = Literal[16127, 16137, 16147, 16157, 16167, 16177, 16187, 16197]

type FluxPortMap = dict[FluxApiPorts, UpnpLease]


@dataclass
class UpnpQuerier:
    api_ports: ClassVar[set[FluxApiPorts]] = set(
        [16127, 16137, 16147, 16157, 16167, 16177, 16187, 16197]
    )
    MAPPING_DESCRIPTION: ClassVar[str] = "FluxOS Reserved"
    used_api_ports: set[FluxApiPorts] = field(default_factory=set)
    used_client_ports: set[int] = field(default_factory=set)
    port_to_host_map: FluxPortMap = field(default_factory=dict)
    api_port: FluxApiPorts | None = None
    external_ip: str | None = None
    control_url: str | None = None
    igd_count: int = 0
    igd_available: bool = False
    initiated: bool = False

    def __post_init__(self) -> None:
        # Interesting properties:
        #   lanaddr, multicastif, minissdpdsocket, lanaddr
        #
        # Interesting methods:
        #   externalipaddress(), statusinfo(), connectiontype()
        #   addportmapping(64000, 'TCP',
        #       '192.168.1.166', 63000, 'port mapping test', '')
        #   deleteportmapping(64000, 'TCP')

        self.client = miniupnpc.UPnP()
        self.client.discoverdelay = 20  # ms

    @property
    def free_api_ports(self) -> set:
        return UpnpQuerier.api_ports - self.used_api_ports

    @property
    def igd_address(self) -> IPv4Address | None:
        if not self.control_url:
            return None

        parsed = URL(self.control_url)
        return IPv4Address(parsed.host)

    @property
    def client_address(self) -> IPv4Address | None:
        return IPv4Address(self.client.lanaddr) if self.client.lanaddr else None

    def _initiate(self) -> None:
        self.initiated = True

        try:
            self.igd_count = self.client.discover()
        except Exception:
            log.error("Timed out searching for IGD's")
            return

        try:
            self.control_url = self.client.selectigd()
        except Exception:
            log.error("IGD Unavailable")
            return

        self.igd_available = True

    def _run_mapping(self, port: int, lifetime: int = 3600):
        retry_permament = False
        success: bool = False

        try:
            success = self.client.addportmapping(
                port,
                "TCP",
                self.client.lanaddr,
                port,
                self.MAPPING_DESCRIPTION,
                "",
                lifetime,
            )
        except Exception as e:
            log.exception("Unable to set mapping for port: %s", port)

            if str(e) == "OnlyPermanentLeasesSupported":
                retry_permament = True

        return retry_permament, success

    def _add_mapping(self, port: int) -> AddMappingResult:
        if not self.initiated:
            self._initiate()

        result = AddMappingResult(port)

        if not self.igd_available:
            return result

        retry_permament, success = self._run_mapping(port)

        if retry_permament:
            _, success = self._run_mapping(port, 0)

        # not sure if this ever happens
        if not success:
            return result

        lease_time = 0 if retry_permament else 3600

        result.lease = UpnpLease(self.client.lanaddr, lease_time, "TCP", self.MAPPING_DESCRIPTION)

        return result

    def _remove_mapping(self, port: int) -> tuple[int, bool]:
        if not self.initiated:
            self._initiate()

        if not self.igd_available:
            return port, False

        try:
            success: bool = self.client.deleteportmapping(
                port,
                "TCP",
                # this is optional. Depending on if the upnp server is running
                # in strict mode, determines if you can remove mappings for other
                # addresses, most run in strict mode.
                self.client.lanaddr,
            )
        except Exception:
            log.exception("Unable to remove mapping for port: %s", port)
            return port, False

        return port, success

    def _build_mapping_tables(self) -> None:
        if not self.initiated:
            self._initiate()

        if not self.igd_available:
            return

        self.used_api_ports.clear()
        self.used_client_ports.clear()

        counter = count()

        for index in counter:
            mapping = self.client.getgenericportmapping(index)
            if mapping is None:
                break

            (port, proto, (ihost, iport), desc, _, _, lease_time) = mapping

            if proto != "TCP":
                continue

            log.info(
                "Mapping found for host: %s, port: %s, iport: %s, desc: %s",
                ihost,
                port,
                iport,
                desc,
            )

            if port in self.api_ports:
                self.used_api_ports.add(port)
                self.port_to_host_map[port] = UpnpLease(ihost, lease_time, proto, desc)

            if ihost == self.client.lanaddr:
                self.used_client_ports.add(port)

    async def get_used_port_map(self, force: bool = False) -> FluxPortMap:
        """Builds a map of ports and which host has used what port

        Args:
            force (bool, optional): Forces a lookup via the router. Defaults to False.

        Returns:
            dict[int, str]: The ports to host map
        """
        if not self.port_to_host_map or force:
            await asyncio.to_thread(self._build_mapping_tables)

        return self.port_to_host_map

    async def get_all_mappings(self) -> dict[int, UpnpLease]:
        """Get ALL UPnP port mappings from the router.

        Returns:
            dict[int, UpnpLease]: All port mappings (not filtered to FluxOS API ports)
        """
        if not self.initiated:
            await asyncio.to_thread(self._initiate)

        if not self.igd_available:
            return {}

        def _get_all():
            all_mappings: dict[int, UpnpLease] = {}
            counter = count()

            for index in counter:
                mapping = self.client.getgenericportmapping(index)
                if mapping is None:
                    break

                (port, proto, (ihost, iport), desc, _, _, lease_time) = mapping

                all_mappings[port] = UpnpLease(ihost, lease_time, proto, desc)

            return all_mappings

        return await asyncio.to_thread(_get_all)

    async def find_free_port(self) -> int | None:
        """Returns the highest non used FluxOS api port.

        Returns:
            int | None: The next free port
        """

        if not self.port_to_host_map:
            await asyncio.to_thread(self._build_mapping_tables)

        try:
            return sorted(self.free_api_ports).pop()
        except IndexError:
            return None

    async def assign_mapping(self, port: int) -> AddMappingResult:
        result = await asyncio.to_thread(self._add_mapping, port)
        return result

    async def unassign_mapping(self, port: int) -> tuple[int, bool]:
        result = await asyncio.to_thread(self._remove_mapping, port)
        return result

    async def remove_all_mappings(self) -> list[int]:
        # we refresh the mappings here as they can be stale
        await asyncio.to_thread(self._build_mapping_tables)

        log.info("Removing used client ports: %s", self.used_client_ports)

        removed: list[int] = []

        for port in self.used_client_ports:
            _, success = await self.unassign_mapping(port)

            if not success:
                return removed

            removed.append(port)

        return removed


def clean_upnp() -> list[int]:
    """Remove all UPnP port mappings created by this machine."""
    async def _run() -> list[int]:
        querier = UpnpQuerier()
        return await querier.remove_all_mappings()

    return asyncio.run(_run())


if __name__ == "__main__":

    async def main():
        builder = UpnpQuerier()
        print(await builder.find_free_port())
        # await builder.assign_mapping(23112)

    asyncio.run(main())
