#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
# Copyright (C) 2026 Bryan Joshua Pedini
"""
Ansible dynamic inventory script that sources hosts from NetBox.

Hosts are read from either the NetBox IP address list or the virtual machine
list, optionally narrowed down by a single filter option; everything is set in
the configuration file and can be overridden through environment variables.
"""

import argparse
import json
import os
import sys
from typing import Any

import requests
import yaml

DEFAULT_CONFIG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml")
GROUP_NAME = "netbox"

SOURCE_IP_ADDRESSES = "ip-addresses"
SOURCE_VIRTUAL_MACHINES = "virtual-machines"
SOURCES = (SOURCE_IP_ADDRESSES, SOURCE_VIRTUAL_MACHINES)

# filter option -> sources it applies to
FILTERS = {
    "tenant": SOURCES,
    "domain": SOURCES,
    "named": (SOURCE_IP_ADDRESSES,),
    "vrf": (SOURCE_IP_ADDRESSES,),
    "cluster": (SOURCE_VIRTUAL_MACHINES,),
    "site": (SOURCE_VIRTUAL_MACHINES,),
}


class InventoryError(Exception):
    """Raised when the inventory cannot be built (configuration or API problems)."""


class NetBoxAPIError(InventoryError):
    """
    Raised when the NetBox API cannot be reached or returns an error.

    Attributes:
        message: Human readable description of the failure.
        status_code: HTTP status code of the response, if one was received.
    """

    def __init__(self, message: str, status_code: int | None = None):
        super().__init__(message)
        self.message = message
        self.status_code = status_code


def _extract_error_message(response: requests.Response) -> str:
    """Pull a readable error message out of a NetBox error response."""
    try:
        data = response.json()
        if isinstance(data, dict) and "detail" in data:
            return f"NetBox API error: {data['detail']}"
    except ValueError:
        pass
    return f"NetBox API error: HTTP {response.status_code}"


class NetBoxClient:
    """Thin wrapper around the NetBox REST API."""

    def __init__(self, url: str, token: str, verify_ssl: bool = True, timeout: int = 30):
        self._url = url.rstrip("/")
        self._timeout = timeout
        self._session = requests.Session()
        self._session.verify = verify_ssl
        self._session.headers.update(
            {
                "Authorization": f"Token {token}",
                "Accept": "application/json",
            }
        )

    def _request(self, url: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
        """
        Perform a GET request and return the decoded JSON body.

        Raises:
            NetBoxAPIError: On connection problems, non-200 responses or
                unparsable bodies.
        """
        try:
            response = self._session.get(url, params=params, timeout=self._timeout)
        except requests.exceptions.Timeout as exc:
            raise NetBoxAPIError("NetBox API request timed out.") from exc
        except requests.exceptions.SSLError as exc:
            raise NetBoxAPIError("SSL error while connecting to the NetBox API.") from exc
        except requests.exceptions.ConnectionError as exc:
            raise NetBoxAPIError("Unable to connect to the NetBox API.") from exc
        except requests.exceptions.RequestException as exc:
            raise NetBoxAPIError(f"NetBox API request failed: {exc}") from exc
        if response.status_code != 200:
            raise NetBoxAPIError(
                _extract_error_message(response), status_code=response.status_code
            )
        try:
            return response.json()
        except ValueError as exc:
            raise NetBoxAPIError(
                "NetBox API returned invalid JSON.", status_code=response.status_code
            ) from exc

    def get_all(self, path: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
        """Return every result of a paginated list endpoint."""
        results: list[dict[str, Any]] = []
        url: str | None = f"{self._url}/api/{path.strip('/')}/"
        while url:
            data = self._request(url, params=params)
            results.extend(data.get("results", []))
            url = data.get("next")
            params = None  # the "next" URL already carries the query string
        return results

    def resolve_id(self, path: str, name: str, label: str) -> int:
        """
        Look up an object by name and return its numeric id.

        Raises:
            NetBoxAPIError: If no object with that name exists.
        """
        matches = self.get_all(path, params={"name": name})
        if not matches:
            raise NetBoxAPIError(f'No {label} named "{name}" found in NetBox.')
        return matches[0]["id"]


def load_config(path: str | None = None) -> dict[str, Any]:
    """
    Load the YAML configuration file and apply environment overrides.

    Args:
        path: Configuration file path; falls back to $NETBOX_INVENTORY_CONFIG,
            then to config.yaml next to this script.

    Returns:
        A flat configuration dict with url, token, verify_ssl, timeout,
        source, filter and filter_value keys.

    Raises:
        InventoryError: If the file is malformed or url/token are missing.
    """
    path = path or os.environ.get("NETBOX_INVENTORY_CONFIG") or DEFAULT_CONFIG_FILE
    config: dict[str, Any] = {}
    if os.path.exists(path):
        with open(path, encoding="utf-8") as handle:
            loaded = yaml.safe_load(handle) or {}
        if not isinstance(loaded, dict):
            raise InventoryError(f"Invalid configuration file: {path}")
        config = loaded
    netbox = config.get("netbox") or {}
    inventory = config.get("inventory") or {}
    merged = {
        "url": os.environ.get("NETBOX_URL", netbox.get("url") or ""),
        "token": os.environ.get("NETBOX_TOKEN", netbox.get("token") or ""),
        "verify_ssl": bool(netbox.get("verify_ssl", True)),
        "timeout": netbox.get("timeout", 30),
        "source": os.environ.get(
            "NETBOX_INVENTORY_SOURCE", inventory.get("source") or SOURCE_IP_ADDRESSES
        ),
        "filter": os.environ.get("NETBOX_INVENTORY_FILTER", inventory.get("filter") or ""),
        "filter_value": os.environ.get(
            "NETBOX_INVENTORY_FILTER_VALUE", inventory.get("filter_value") or ""
        ),
    }
    try:
        merged["timeout"] = int(merged["timeout"])
    except (TypeError, ValueError) as exc:
        raise InventoryError("Configuration value netbox.timeout must be a number.") from exc
    if not merged["url"] or not merged["token"]:
        raise InventoryError(
            "NetBox url and token are required, set them in the configuration file "
            "or through the NETBOX_URL / NETBOX_TOKEN environment variables."
        )
    return merged


def parse_bool(value: str) -> bool:
    """
    Parse a yes/no filter value, an empty string counts as yes.

    Raises:
        InventoryError: If the value is not recognisable as a boolean.
    """
    lowered = value.strip().lower()
    if lowered in ("", "1", "true", "yes", "on"):
        return True
    if lowered in ("0", "false", "no", "off"):
        return False
    raise InventoryError(f'Filter value "{value}" is not a valid boolean.')


def domain_matches(hostname: str, domain: str) -> bool:
    """
    Check whether a hostname belongs to a domain.

    The hostname must end with the domain on a label boundary, so
    "web.example.com" matches "example.com" but "notexample.com" does not.
    """
    if not hostname:
        return False
    domain = domain.lstrip(".").rstrip(".").lower()
    hostname = hostname.rstrip(".").lower()
    return hostname == domain or hostname.endswith("." + domain)


def validate_options(source: str, filter_name: str, filter_value: str) -> None:
    """
    Validate the source/filter combination before touching the API.

    Raises:
        InventoryError: On unknown sources or filters, a filter that does not
            apply to the source, or a missing filter value.
    """
    if source not in SOURCES:
        raise InventoryError(f'Unknown source "{source}", valid sources: {", ".join(SOURCES)}.')
    if not filter_name:
        return
    if filter_name not in FILTERS:
        raise InventoryError(
            f'Unknown filter "{filter_name}", valid filters: {", ".join(FILTERS)}.'
        )
    if source not in FILTERS[filter_name]:
        raise InventoryError(f'Filter "{filter_name}" does not apply to source "{source}".')
    if filter_name != "named" and not filter_value:
        raise InventoryError(f'Filter "{filter_name}" requires a value.')


def server_side_params(
    client: NetBoxClient, filter_name: str, filter_value: str
) -> dict[str, Any]:
    """Translate name-based filters into NetBox API query parameters."""
    if filter_name == "tenant":
        return {"tenant_id": client.resolve_id("tenancy/tenants", filter_value, "tenant")}
    if filter_name == "vrf":
        return {"vrf_id": client.resolve_id("ipam/vrfs", filter_value, "VRF")}
    if filter_name == "cluster":
        return {"cluster_id": client.resolve_id("virtualization/clusters", filter_value, "cluster")}
    if filter_name == "site":
        return {"site_id": client.resolve_id("dcim/sites", filter_value, "site")}
    return {}


def apply_client_side_filters(
    entries: list[tuple[str, dict[str, Any]]], filter_name: str, filter_value: str
) -> list[tuple[str, dict[str, Any]]]:
    """
    Apply the filters that cannot be expressed as API query parameters.

    "domain" keeps hostnames belonging to the given domain, "named" keeps only
    IP addresses with a DNS name and deduplicates them by hostname.
    """
    if filter_name == "domain":
        return [(name, host) for name, host in entries if domain_matches(name, filter_value)]
    if filter_name == "named" and parse_bool(filter_value):
        seen: set[str] = set()
        named = []
        for name, host in entries:
            if not host.get("netbox_dns_name") or name in seen:
                continue
            seen.add(name)
            named.append((name, host))
        return named
    return entries


def strip_prefix(address: str) -> str:
    """Return a bare IP address without the CIDR prefix length."""
    return address.split("/", 1)[0]


def ip_host_entry(ip: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """
    Turn a NetBox IP address record into an inventory (name, hostvars) pair.

    The DNS name is used as the inventory hostname when set, the bare address
    otherwise; the bare address always ends up in ansible_host.
    """
    address = strip_prefix(ip.get("address") or "")
    name = ip.get("dns_name") or address
    hostvars = {
        "ansible_host": address,
        "netbox_ip_address": ip.get("address") or "",
        "netbox_dns_name": ip.get("dns_name") or "",
        "netbox_tenant": (ip.get("tenant") or {}).get("name") or "",
        "netbox_vrf": (ip.get("vrf") or {}).get("name") or "",
        "netbox_description": ip.get("description") or "",
    }
    return name, hostvars


def vm_host_entry(vm: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """
    Turn a NetBox virtual machine record into an inventory (name, hostvars) pair.

    ansible_host is only set when the virtual machine has a primary IP, so that
    Ansible falls back to resolving the name otherwise.
    """
    name = vm.get("name") or ""
    hostvars = {
        "netbox_status": (vm.get("status") or {}).get("value") or "",
        "netbox_cluster": (vm.get("cluster") or {}).get("name") or "",
        "netbox_site": (vm.get("site") or {}).get("name") or "",
        "netbox_tenant": (vm.get("tenant") or {}).get("name") or "",
    }
    primary = vm.get("primary_ip") or {}
    if primary.get("address"):
        hostvars["ansible_host"] = strip_prefix(primary["address"])
    return name, hostvars


def fetch_entries(
    client: NetBoxClient, source: str, params: dict[str, Any] | None = None
) -> list[tuple[str, dict[str, Any]]]:
    """Fetch all hosts of the configured source from NetBox."""
    if source == SOURCE_IP_ADDRESSES:
        records = client.get_all("ipam/ip-addresses", params)
        return [ip_host_entry(record) for record in records if record.get("address")]
    params = {"exclude": "config_context", **(params or {})}
    records = client.get_all("virtualization/virtual-machines", params)
    return [vm_host_entry(record) for record in records if record.get("name")]


def build_inventory(config: dict[str, Any], client: NetBoxClient) -> dict[str, Any]:
    """
    Build the full --list inventory structure for Ansible.

    Raises:
        InventoryError: On invalid source/filter options or NetBox API failures.
    """
    source = config["source"]
    filter_name = config["filter"]
    filter_value = config["filter_value"]
    validate_options(source, filter_name, filter_value)
    params = server_side_params(client, filter_name, filter_value) if filter_name else {}
    entries = fetch_entries(client, source, params)
    entries = apply_client_side_filters(entries, filter_name, filter_value)
    hostvars: dict[str, dict[str, Any]] = {}
    for name, host in entries:
        if name in hostvars:
            # duplicate inventory names keep both hosts by falling back to the
            # bare address as the name; identical duplicates are dropped
            name = host.get("ansible_host") or name
            if name in hostvars:
                continue
        hostvars[name] = host
    return {
        GROUP_NAME: {"hosts": sorted(hostvars)},
        "_meta": {"hostvars": hostvars},
    }


def main() -> int:
    """Entry point for the Ansible dynamic inventory protocol."""
    parser = argparse.ArgumentParser(description="Ansible dynamic inventory backed by NetBox.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--list", action="store_true", help="output the full inventory as JSON")
    group.add_argument("--host", metavar="HOSTNAME", help="output variables for a single host")
    args = parser.parse_args()
    if args.host:
        # all host variables are returned through _meta in --list
        print(json.dumps({}))
        return 0
    try:
        config = load_config()
        client = NetBoxClient(
            config["url"], config["token"], config["verify_ssl"], config["timeout"]
        )
        print(json.dumps(build_inventory(config, client), indent=2, sort_keys=True))
    except InventoryError as exc:
        print(f"netbox-inventory: {exc}", file=sys.stderr)
        return 1
    return 0


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