#!/usr/bin/env python3

import argparse
import ast
import gzip
import os
import pprint
import shutil
import subprocess
import sys
import tempfile
from datetime import datetime
from pathlib import Path


ATTRIBUTE_ID = "snmp_location"
EXPLICIT_ATTRIBUTE_ID = "_snmp_location"
DEFAULT_HOSTNAME_FILTER = ".bl.ch"

INVENTORY_SUBDIR = Path("var/check_mk/inventory")
WATO_SUBDIR = Path("etc/check_mk/conf.d/wato")
BACKUP_SUBDIR = Path("var/check_mk/snmp_location_backups")


def info(message):
    print(f"[INFO] {message}")


def warn(message):
    print(f"[WARN] {message}")


def error(message):
    print(f"[ERROR] {message}")


def get_omd_root():
    value = os.environ.get("OMD_ROOT")
    if not value:
        error("OMD_ROOT ist nicht gesetzt.")
        error("Das Script muss innerhalb einer Checkmk-OMD-Site ausgeführt werden.")
        sys.exit(1)

    root = Path(value)
    if not root.is_dir():
        error(f"OMD_ROOT existiert nicht: {root}")
        sys.exit(1)

    return root


def is_inventory_candidate(path):
    if not path.is_file():
        return False

    name = path.name

    if name in ("archive", "inv_files"):
        return False

    # Temporäre Checkmk-Dateien ignorieren, z.B.
    # .hostname.bl.ch.gz.newsbakec76
    if name.startswith("."):
        return False

    if ".new" in name or ".tmp" in name:
        return False

    return True


def inventory_hostname(path):
    name = path.name

    if name.endswith(".gz"):
        name = name[:-3]

    return name


def read_inventory_file(path):
    try:
        if path.name.endswith(".gz"):
            with gzip.open(path, "rt", encoding="utf-8") as handle:
                content = handle.read()
        else:
            with path.open("r", encoding="utf-8") as handle:
                content = handle.read()

        return ast.literal_eval(content)

    except Exception as exc:
        warn(f"Inventory konnte nicht gelesen werden: {path}: {exc}")
        return None


def is_valid_location(value):
    if not isinstance(value, str):
        return False

    value = value.strip()

    if not value:
        return False

    lower = value.lower()

    if "chassis" in lower or "module" in lower:
        return False

    return True


def extract_location_from_snmp_info(snmp_info):
    if not isinstance(snmp_info, dict):
        return None

    attributes = snmp_info.get("Attributes", {})
    if not isinstance(attributes, dict):
        return None

    pairs = attributes.get("Pairs", {})
    if not isinstance(pairs, dict):
        return None

    location = pairs.get("location")

    if is_valid_location(location):
        return location.strip()

    return None


def get_snmp_location(data):
    if not isinstance(data, dict):
        return None

    root = data.get("raw_tree", data)

    if not isinstance(root, dict):
        return None

    nodes = root.get("Nodes", {})

    if not isinstance(nodes, dict):
        return None

    software = nodes.get("software", {})

    if isinstance(software, dict):
        software_nodes = software.get("Nodes", {})

        if isinstance(software_nodes, dict):
            configuration = software_nodes.get("configuration", {})

            if isinstance(configuration, dict):
                configuration_nodes = configuration.get("Nodes", {})

                if isinstance(configuration_nodes, dict):
                    location = extract_location_from_snmp_info(
                        configuration_nodes.get("snmp_info", {})
                    )

                    if location:
                        return location

    networking = nodes.get("networking", {})

    if isinstance(networking, dict):
        networking_nodes = networking.get("Nodes", {})

        if isinstance(networking_nodes, dict):
            location = extract_location_from_snmp_info(
                networking_nodes.get("snmp_info", {})
            )

            if location:
                return location

    return None


def build_inventory_mapping(inventory_dir, hostname_filter):
    if not inventory_dir.is_dir():
        error(f"Inventory-Verzeichnis nicht gefunden: {inventory_dir}")
        sys.exit(1)

    files = []

    for path in inventory_dir.iterdir():
        if not is_inventory_candidate(path):
            continue

        hostname = inventory_hostname(path)

        if hostname_filter and hostname_filter not in hostname:
            continue

        files.append((hostname, path))

    files.sort(key=lambda item: item[0].lower())

    info(f"{len(files)} passende Inventory-Dateien gefunden.")

    fqdn_locations = {}
    short_locations = {}
    valid_count = 0

    for hostname, path in files:
        data = read_inventory_file(path)

        if data is None:
            continue

        location = get_snmp_location(data)

        if not location:
            continue

        valid_count += 1
        fqdn_locations[hostname] = location

        shortname = hostname.split(".", 1)[0]

        if shortname not in short_locations:
            short_locations[shortname] = location

        elif short_locations[shortname] != location:
            warn(
                f"Shortname-Konflikt für {shortname}: "
                f"'{short_locations[shortname]}' / '{location}'. "
                f"FQDN wird bevorzugt."
            )

    info(f"{valid_count} gültige Locations aus dem Inventory gelesen.")

    mapping = dict(short_locations)
    mapping.update(fqdn_locations)

    info(f"{len(mapping)} Host-/Shortname-Einträge im Speicher.")

    return mapping


def is_name(node, name):
    return isinstance(node, ast.Name) and node.id == name


def is_host_attributes_update(node):
    return (
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Attribute)
        and node.func.attr == "update"
        and is_name(node.func.value, "host_attributes")
    )


def is_explicit_location_update(node):
    if not (
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Attribute)
        and node.func.attr == "update"
    ):
        return False

    subscript = node.func.value

    if not isinstance(subscript, ast.Subscript):
        return False

    if not is_name(subscript.value, "explicit_host_conf"):
        return False

    key = subscript.slice

    if isinstance(key, ast.Constant):
        return key.value == EXPLICIT_ATTRIBUTE_ID

    if isinstance(key, ast.Str):
        return key.s == EXPLICIT_ATTRIBUTE_ID

    return False


def find_update_blocks(source):
    tree = ast.parse(source)

    result = {
        "host_attributes": None,
        "explicit_location": None,
    }

    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue

        if not node.args:
            continue

        dictionary = node.args[0]

        if not isinstance(dictionary, ast.Dict):
            continue

        if is_host_attributes_update(node):
            result["host_attributes"] = dictionary

        elif is_explicit_location_update(node):
            result["explicit_location"] = dictionary

    return result


def ast_dict_to_python_dict(node):
    try:
        value = ast.literal_eval(node)

        if isinstance(value, dict):
            return value

    except Exception:
        pass

    return None


def source_offsets(source):
    offsets = [0]

    for line in source.splitlines(keepends=True):
        offsets.append(offsets[-1] + len(line))

    return offsets


def node_start(node, offsets):
    return offsets[node.lineno - 1] + node.col_offset


def node_end(node, offsets):
    return offsets[node.end_lineno - 1] + node.end_col_offset


def replace_ast_node(source, node, replacement, offsets):
    start = node_start(node, offsets)
    end = node_end(node, offsets)

    return source[:start] + replacement + source[end:]


def serialize_dictionary(value):
    return pprint.pformat(
        value,
        width=160,
        sort_dicts=False,
        compact=False,
    )


def get_inventory_location(hostname, inventory_mapping):
    # FQDN zuerst, Shortname nur als Fallback.
    if hostname in inventory_mapping:
        return inventory_mapping[hostname], "FQDN"

    shortname = hostname.split(".", 1)[0]

    if shortname in inventory_mapping:
        return inventory_mapping[shortname], "Shortname"

    return None, None


def prepare_wato_file(path, inventory_mapping, hostname_filter, verbose):
    source = path.read_text(encoding="utf-8")

    try:
        blocks = find_update_blocks(source)

    except SyntaxError as exc:
        return {
            "changed": False,
            "invalid": True,
            "reason": f"Python-Syntaxfehler: {exc}",
            "content": source,
            "updated": [],
            "missing": [],
            "checked": [],
        }

    host_node = blocks["host_attributes"]

    if host_node is None:
        return {
            "changed": False,
            "invalid": False,
            "reason": "Kein host_attributes.update()-Block gefunden.",
            "content": source,
            "updated": [],
            "missing": [],
            "checked": [],
        }

    host_attributes = ast_dict_to_python_dict(host_node)

    if host_attributes is None:
        return {
            "changed": False,
            "invalid": True,
            "reason": "host_attributes-Dictionary ist nicht literal auswertbar.",
            "content": source,
            "updated": [],
            "missing": [],
            "checked": [],
        }

    explicit_node = blocks["explicit_location"]

    if explicit_node is not None:
        explicit_locations = ast_dict_to_python_dict(explicit_node)

        if explicit_locations is None:
            return {
                "changed": False,
                "invalid": True,
                "reason": (
                    "explicit_host_conf['_snmp_location'] "
                    "ist nicht literal auswertbar."
                ),
                "content": source,
                "updated": [],
                "missing": [],
                "checked": [],
            }

    else:
        explicit_locations = {}

    updated = []
    missing = []
    checked = []

    for hostname, attributes in host_attributes.items():
        if not isinstance(hostname, str):
            continue

        if hostname_filter and hostname_filter not in hostname:
            continue

        if not isinstance(attributes, dict):
            continue

        inventory_location, source_type = get_inventory_location(
            hostname,
            inventory_mapping,
        )

        explicit_location = explicit_locations.get(hostname)
        attribute_location = attributes.get(ATTRIBUTE_ID)

        if inventory_location is None:
            missing.append(hostname)

            if verbose:
                checked.append({
                    "hostname": hostname,
                    "inventory": None,
                    "source_type": None,
                    "explicit": explicit_location,
                    "attribute": attribute_location,
                    "status": "NO INVENTORY",
                })

            continue

        host_changed = attribute_location != inventory_location
        explicit_changed = explicit_location != inventory_location

        if host_changed or explicit_changed:
            attributes[ATTRIBUTE_ID] = inventory_location
            explicit_locations[hostname] = inventory_location

            updated.append({
                "hostname": hostname,
                "inventory": inventory_location,
                "source_type": source_type,
                "explicit_before": explicit_location,
                "attribute_before": attribute_location,
            })

            status = "CHANGE"

        else:
            status = "OK"

        if verbose:
            checked.append({
                "hostname": hostname,
                "inventory": inventory_location,
                "source_type": source_type,
                "explicit": explicit_location,
                "attribute": attribute_location,
                "status": status,
            })

    if not updated:
        return {
            "changed": False,
            "invalid": False,
            "reason": "keine Änderungen",
            "content": source,
            "updated": [],
            "missing": missing,
            "checked": checked,
        }

    offsets = source_offsets(source)

    replacements = [
        (
            host_node,
            serialize_dictionary(host_attributes),
        )
    ]

    if explicit_node is not None:
        replacements.append(
            (
                explicit_node,
                serialize_dictionary(explicit_locations),
            )
        )

    replacements.sort(
        key=lambda item: node_start(item[0], offsets),
        reverse=True,
    )

    new_source = source

    for node, replacement in replacements:
        new_source = replace_ast_node(
            new_source,
            node,
            replacement,
            offsets,
        )

    try:
        compile(new_source, str(path), "exec")

    except SyntaxError as exc:
        return {
            "changed": False,
            "invalid": True,
            "reason": f"Erzeugte Datei ist syntaktisch ungültig: {exc}",
            "content": source,
            "updated": [],
            "missing": missing,
            "checked": checked,
        }

    return {
        "changed": True,
        "invalid": False,
        "reason": "Änderungen vorbereitet",
        "content": new_source,
        "updated": updated,
        "missing": missing,
        "checked": checked,
    }


def print_verbose_details(result):
    for item in result["checked"]:
        hostname = item["hostname"]
        inventory = item["inventory"]
        explicit = item["explicit"]
        attribute = item["attribute"]
        status = item["status"]

        print(f"  HOST: {hostname}")

        if inventory is None:
            print("    Inventory:       -- kein Wert --")
        else:
            print(f"    Inventory:       {inventory}")

        if item["source_type"]:
            print(f"    Quelle:          {item['source_type']}")

        if explicit is None:
            print("    WATO explicit:   -- kein Wert --")
        else:
            print(f"    WATO explicit:   {explicit}")

        if attribute is None:
            print("    WATO attribute:  -- kein Wert --")
        else:
            print(f"    WATO attribute:  {attribute}")

        print(f"    Status:          {status}")
        print()


def create_backup(path, wato_dir, backup_root):
    relative = path.relative_to(wato_dir)
    destination = backup_root / relative

    destination.parent.mkdir(parents=True, exist_ok=True)

    shutil.copy2(path, destination)

    return destination


def write_atomic(path, content):
    fd, temp_name = tempfile.mkstemp(
        prefix=f".{path.name}.",
        suffix=".tmp",
        dir=str(path.parent),
        text=True,
    )

    temp_path = Path(temp_name)

    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(content)
            handle.flush()
            os.fsync(handle.fileno())

        shutil.copymode(path, temp_path)

        os.replace(temp_path, path)

    except Exception:
        try:
            temp_path.unlink(missing_ok=True)
        except Exception:
            pass

        raise


def restore_files(changed_files):
    warn("Stelle ursprüngliche Dateien wieder her...")

    for path, backup in changed_files:
        try:
            shutil.copy2(backup, path)
            info(f"Rollback: {path}")

        except Exception as exc:
            error(f"Rollback fehlgeschlagen: {path}: {exc}")


def run_cmk_reload():
    info("Führe 'cmk -R' aus...")
    print()

    result = subprocess.run(
        ["cmk", "-R"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    if result.stdout:
        print(result.stdout)

    return result.returncode


def main():
    parser = argparse.ArgumentParser(
        description="Synchronisiert Checkmk SNMP-Locations aus dem Inventory."
    )

    parser.add_argument(
        "--apply",
        action="store_true",
        help="Änderungen tatsächlich schreiben.",
    )

    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Für jeden passenden Host Inventory- und WATO-Werte anzeigen.",
    )

    parser.add_argument(
        "--no-reload",
        action="store_true",
        help="Nach Änderungen kein cmk -R ausführen.",
    )

    parser.add_argument(
        "--hostname-filter",
        default=DEFAULT_HOSTNAME_FILTER,
        help="Nur Hostnamen mit diesem Text bearbeiten.",
    )

    args = parser.parse_args()

    print()
    print("=" * 70)
    print(" Checkmk SNMP Location Synchronisation")
    print("=" * 70)
    print()

    omd_root = get_omd_root()

    inventory_dir = omd_root / INVENTORY_SUBDIR
    wato_dir = omd_root / WATO_SUBDIR

    info(f"OMD_ROOT:        {omd_root}")
    info(f"Inventory:       {inventory_dir}")
    info(f"WATO:            {wato_dir}")
    info(f"Attribut:        {ATTRIBUTE_ID}")
    info(f"Hostname-Filter: '{args.hostname_filter}'")
    print()

    if not args.apply:
        warn("DRY-RUN: Es werden KEINE Dateien verändert.")
        warn("Für tatsächliche Änderungen: --apply")
        print()

    info("Lese Inventory...")

    inventory_mapping = build_inventory_mapping(
        inventory_dir,
        args.hostname_filter,
    )

    print()
    info("Suche WATO hosts.mk-Dateien...")

    hosts_files = sorted(
        wato_dir.rglob("hosts.mk"),
        key=lambda path: str(path).lower(),
    )

    info(f"{len(hosts_files)} hosts.mk-Dateien gefunden.")
    print()

    total_checked = 0
    total_updated = 0
    total_missing = 0
    total_invalid = 0

    changed_files = []

    for path in hosts_files:
        result = prepare_wato_file(
            path,
            inventory_mapping,
            args.hostname_filter,
            args.verbose,
        )

        total_checked += len(result["checked"])
        total_updated += len(result["updated"])
        total_missing += len(result["missing"])

        if result["invalid"]:
            total_invalid += 1
            error(f"{path}: {result['reason']}")
            continue

        if result["changed"]:
            print(
                f"[CHANGE] {path}: "
                f"{len(result['updated'])} Host(s)"
            )

            for item in result["updated"]:
                print(
                    f"         {item['hostname']} -> "
                    f"{item['inventory']}"
                )

        elif args.verbose:
            info(f"{path}: {result['reason']}")

        if args.verbose and result["checked"]:
            print_verbose_details(result)

        if result["changed"]:
            changed_files.append((path, result["content"]))

    print()
    print("-" * 70)
    print(" Ergebnis")
    print("-" * 70)
    print(f"Hosts geprüft:                      {total_checked}")
    print(f"Hosts mit neuer/anderer Location:  {total_updated}")
    print(f"Hosts ohne Inventory-Wert:          {total_missing}")
    print(f"Ungültige WATO-Dateien:             {total_invalid}")
    print(f"Dateien mit Änderungen:             {len(changed_files)}")

    if total_invalid:
        print()
        error("Mindestens eine WATO-Datei ist ungültig.")
        error("Es wurden KEINE Dateien geschrieben.")
        return 1

    if not changed_files:
        print()
        print("[ OK ] Keine Änderungen erforderlich.")
        return 0

    if not args.apply:
        print()
        print(
            f"[ OK ] DRY-RUN: {total_updated} Host(s) "
            "würden geändert."
        )
        print()
        print("Für die tatsächliche Synchronisation:")
        print("    ./update_location_direct.py --apply")
        print()
        return 0

    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")

    backup_root = omd_root / BACKUP_SUBDIR / timestamp
    backup_root.mkdir(parents=True, exist_ok=False)

    info(f"Backup-Verzeichnis: {backup_root}")

    written_files = []

    try:
        for path, new_content in changed_files:
            backup = create_backup(
                path,
                wato_dir,
                backup_root,
            )

            info(f"Backup erstellt: {path}")

            write_atomic(
                path,
                new_content,
            )

            written_files.append(
                (path, backup)
            )

            info(f"Datei geschrieben: {path}")

    except Exception as exc:
        error(f"Fehler beim Schreiben: {exc}")

        if written_files:
            restore_files(written_files)

        return 1

    if args.no_reload:
        print()
        print("[ OK ] Dateien wurden geändert.")
        info(f"Backup: {backup_root}")
        info("Kein cmk -R wegen --no-reload.")
        return 0

    print()

    returncode = run_cmk_reload()

    if returncode == 0:
        print()
        print("=" * 70)
        print("[ OK ] Synchronisation erfolgreich.")
        print("=" * 70)
        print()
        info(f"Geänderte Hosts: {total_updated}")
        info(f"Backup: {backup_root}")

        return 0

    print()

    error(
        f"'cmk -R' ist fehlgeschlagen "
        f"(Returncode {returncode})."
    )

    restore_files(written_files)

    print()

    info(
        "Prüfe die wiederhergestellte Konfiguration "
        "mit cmk -R..."
    )

    rollback_returncode = run_cmk_reload()

    if rollback_returncode == 0:
        print()
        print("[ OK ] Rollback erfolgreich.")
        info(f"Backup: {backup_root}")

    else:
        print()
        error(
            "ACHTUNG: Auch der Reload nach dem "
            "Rollback ist fehlgeschlagen!"
        )
        error(f"Returncode: {rollback_returncode}")
        error(f"Backup: {backup_root}")

    return 1


if __name__ == "__main__":
    sys.exit(main())