#!/usr/bin/env python3
# -*- encoding: utf-8; py-indent-offset: 4 -*-
#   _____  __          __  _____
#  / ____| \ \        / / |  __ \
# | (___    \ \  /\  / /  | |__) |
#  \___ \    \ \/  \/ /   |  _  /
#  ____) |    \  /\  /    | | \ \
# |_____/      \/  \/     |_|  \_\
#
# (c) 2025 SWR
# @author Frank Baier <frank.baier@swr.de>
#
import argparse
import os
import pickle
import re
import time
from pathlib import Path
from typing import Any
from labelpicker_ng import (
    Config,
    LabelpickerConfig,
    ServiceLabelDefinition,
    cmk,
    init_logger,
    logger,
)

INTERFACE_PREFIXES = {
    "Twe": "TwentyFiveGigE",
    "Fo": "FortyGigabitEthernet",
    "Hu": "HundredGigE",
    "Te": "TenGigabitEthernet",
    "Gi": "GigabitEthernet",
    "Fa": "FastEthernet",
    "Eth": "Ethernet",
    "Po": "Port-channel",
}


def expand_interface_name(interface_name: str) -> str:
    interface_name = interface_name.strip()

    # Die langen Namen müssen vor ihren Kurzformen geprüft werden: Andernfalls
    # würde beispielsweise "Ethernet1/1" auf "Eth" passen und zu
    # "Etherneternet1/1" erweitert werden.
    if any(
        interface_name.startswith(long_name)
        for long_name in INTERFACE_PREFIXES.values()
    ):
        return interface_name

    for short_name, long_name in INTERFACE_PREFIXES.items():
        if interface_name.startswith(short_name):
            return long_name + interface_name[len(short_name):]

    return interface_name


ProtocolState = dict[str, int | bool]
InterfaceCache = dict[str, dict[str, ProtocolState]]
NeighborCache = dict[str, InterfaceCache]


class Checkmk:
    def __init__(
        self,
        wato: cmk.CMKInstance,
        *,
        label_definitions: dict[str, ServiceLabelDefinition],
        pickle_file: str | Path,
    ) -> None:
        self.wato = wato
        self.label_definitions = label_definitions
        self.search_labels = [
            definition.label
            for definition in label_definitions.values()
        ]
        self.retention_days = {
            definition.label: definition.retention_time
            for definition in label_definitions.values()
        }
        self.pickle_file = Path(pickle_file)

    @staticmethod
    def _get_inventory_node(
        parent: dict[str, Any],
        node_name: str,
    ) -> dict[str, Any] | None:
        nodes = parent.get("Nodes")

        if not isinstance(nodes, dict):
            return None

        node = nodes.get(node_name)

        if not isinstance(node, dict):
            return None

        return node

    def get_inventory_table(
        self,
        inventory: dict[str, Any],
        inventory_path: str,
    ) -> dict[str, Any] | None:
        node = inventory
        for node_name in inventory_path.split("."):
            node = self._get_inventory_node(node, node_name)
            if node is None:
                return None

        table = node.get("Table")
        return table if isinstance(table, dict) else None

    def get_label_interfaces(
        self,
        inventory: dict[str, Any],
        definition: ServiceLabelDefinition,
    ) -> set[str]:
        table = self.get_inventory_table(inventory, definition.inventory_path)
        if table is None:
            return set()

        rows = table.get("Rows")
        if not isinstance(rows, list):
            return set()

        interfaces: set[str] = set()
        for row in rows:
            if not isinstance(row, dict):
                continue
            if definition.match_field is not None:
                field_value = row.get(definition.match_field)
                if (
                    not isinstance(field_value, str)
                    or definition.match_value is None
                    or re.search(definition.match_value, field_value) is None
                ):
                    continue
            interface_name = row.get(definition.inventory_field)
            if isinstance(interface_name, str) and interface_name:
                interfaces.add(expand_interface_name(interface_name))
        return interfaces

    def collect_label_interfaces(
        self,
    ) -> dict[str, dict[str, set[str]]]:
        """Read interfaces matching every configured service-label rule."""
        result: dict[str, dict[str, set[str]]] = {}
        inventories = self.wato.get_inventory()

        for host_name, inventory in inventories.items():
            if not isinstance(inventory, dict):
                continue

            result[host_name] = {
                definition.label: self.get_label_interfaces(inventory, definition)
                for definition in self.label_definitions.values()
            }

        return result

    def _load_neighbor_cache(self) -> NeighborCache:
        if not self.pickle_file.exists():
            return {}

        with self.pickle_file.open("rb") as file_handle:
            data = pickle.load(file_handle)

        if not isinstance(data, dict):
            raise TypeError("Das Pickle-File enthält kein Dictionary.")

        cache: NeighborCache = {}
        for host_name, interfaces in data.items():
            if not isinstance(host_name, str) or not isinstance(interfaces, dict):
                continue

            host_cache: InterfaceCache = {}
            for interface_name, protocols in interfaces.items():
                if (
                    not isinstance(interface_name, str)
                    or not isinstance(protocols, dict)
                ):
                    continue

                interface_cache: dict[str, ProtocolState] = {}
                for label in self.search_labels:
                    state = protocols.get(label)
                    if (
                        isinstance(state, dict)
                        and isinstance(state.get("last_seen"), int)
                    ):
                        interface_cache[label] = {
                            "last_seen": state["last_seen"],
                            "current": bool(state.get("current", False)),
                        }

                if interface_cache:
                    host_cache[interface_name] = interface_cache

            if host_cache:
                cache[host_name] = host_cache

        return cache

    def update_neighbor_cache(
        self,
        cache: NeighborCache,
        discovered_interfaces: dict[str, dict[str, set[str]]],
        *,
        now: int | None = None,
    ) -> NeighborCache:
        """Merge a discovery run, retaining vanished neighbors temporarily."""
        timestamp = int(time.time()) if now is None else now

        # A removed configuration entry must disappear immediately and does
        # not use the normal neighbor-retention period.
        configured_labels = set(self.search_labels)
        for host_cache in cache.values():
            for interface_cache in host_cache.values():
                for label in set(interface_cache) - configured_labels:
                    del interface_cache[label]

        for host_name, interfaces in discovered_interfaces.items():
            host_cache = cache.setdefault(host_name, {})

            # Successful inventory retrieval makes all older entries stale
            # until they are seen again in this discovery run.
            for protocol_states in host_cache.values():
                for state in protocol_states.values():
                    state["current"] = False

            for label, interface_names in interfaces.items():
                for interface_name in interface_names:
                    interface_cache = host_cache.setdefault(interface_name, {})
                    interface_cache[label] = {
                        "last_seen": timestamp,
                        "current": True,
                    }

        for host_name in list(cache):
            host_cache = cache[host_name]
            for interface_name in list(host_cache):
                interface_cache = host_cache[interface_name]
                for protocol in list(interface_cache):
                    last_seen = interface_cache[protocol]["last_seen"]
                    retention_days = self.retention_days[protocol]
                    expiry = timestamp - (retention_days * 24 * 60 * 60)
                    if isinstance(last_seen, int) and last_seen < expiry:
                        del interface_cache[protocol]

                if not interface_cache:
                    del host_cache[interface_name]

            if not host_cache:
                del cache[host_name]

        return cache

    def save_pickle(
        self,
        data: NeighborCache,
    ) -> None:
        self.pickle_file.parent.mkdir(
            parents=True,
            exist_ok=True,
        )

        temporary_file = self.pickle_file.with_suffix(
            f"{self.pickle_file.suffix}.tmp"
        )

        with temporary_file.open("wb") as file_handle:
            pickle.dump(
                data,
                file_handle,
                protocol=pickle.HIGHEST_PROTOCOL,
            )

        temporary_file.replace(self.pickle_file)

    def update_service_label_pickle(self) -> NeighborCache:
        cache = self._load_neighbor_cache()
        discovered_interfaces = self.collect_label_interfaces()
        cache = self.update_neighbor_cache(cache, discovered_interfaces)
        self.save_pickle(cache)

        return cache

    def load_pickle(self) -> NeighborCache:
        return self._load_neighbor_cache()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Cache CDP- und LLDP-Nachbarn pro Host und Interface."
    )
    parser.add_argument(
        "-i",
        "--init",
        action="store_true",
        help="Labelpicker-Standardkonfiguration initialisieren.",
    )
    parser.add_argument("-u", "--username", help="Name des Automation-Users.")
    parser.add_argument("-p", "--password", help="Secret des Automation-Users.")
    parser.add_argument("-c", "--config", default="", help="Labelpicker-Konfiguration.")
    parser.add_argument(
        "-w",
        "--print_config",
        action="store_true",
        help="Labelpicker-Standardkonfiguration Base64-kodiert ausgeben.",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        default=False,
        help="Debug-Logging aktivieren.",
    )
    parser.add_argument(
        "-t",
        "--testmode",
        action="store_true",
        help="Discovery durchführen, aber die Pickle-Datei nicht verändern.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.print_config:
        Config.generate_init_cfg(args.config)

    config_instance = Config(args)
    if args.init:
        config_instance.init_cfg()

    config: LabelpickerConfig = config_instance.get_cfg()
    init_logger(config)
    logger.info(__file__.split("/")[-1], extra={"color": "h1"})

    service_labels = config.servicelabels
    if service_labels is None:
        logger.info(
            "Keine 'servicelabels'-Sektion konfiguriert; nichts zu tun."
        )
        return

    if not service_labels.definitions:
        logger.info("Keine Service-Label-Regeln konfiguriert; nichts zu tun.")
        return

    try:
        wato = cmk.CMKInstance(config.checkmk)
    except (ValueError, TypeError) as error:
        logger.error(f"Checkmk-Konfiguration ungültig: {error}")
        raise SystemExit(2) from error
    except Exception:
        logger.error("Verbindung zur Checkmk-REST-API nicht möglich.")
        raise SystemExit(1)

    omd_root = Path(config.checkmk.omd_root or os.environ["OMD_ROOT"])
    pickle_file = (
        omd_root / "var/labelpicker_ng" / service_labels.pickle_name
    )
    checkmk = Checkmk(
        wato,
        label_definitions=service_labels.definitions,
        pickle_file=pickle_file,
    )

    if args.testmode:
        cache = checkmk._load_neighbor_cache()
        discovered_interfaces = checkmk.collect_label_interfaces()
        cache = checkmk.update_neighbor_cache(cache, discovered_interfaces)
        logger.warning("Testmode: Pickle-Datei wird nicht verändert.")
    else:
        cache = checkmk.update_service_label_pickle()

    logger.info(
        f"{len(cache)} Hosts mit Service-Label-Interfaces in {pickle_file}."
    )


if __name__ == "__main__":
    main()
