#!/usr/bin/env python3
"""
Prometheus exporter for Cisco Nexus NX-OS PTP data.

This exporter runs on a Linux monitoring host, NOT on the Nexus itself.
It connects to NX-OS over SSH and collects:

  show ptp clock
  show ptp parent
  show ptp brief
  show ptp counters all

SSH keys are preferred. Password authentication can be supplied via an
environment variable if required.
"""

from __future__ import annotations

import argparse
import math
import os
import re
import socket
import threading
import time
from typing import Dict, List, Optional, Tuple

import paramiko
from prometheus_client import REGISTRY, start_http_server
from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily, InfoMetricFamily


KV_RE = re.compile(r"^\s*([^:]+?)\s*:\s*(.*?)\s*$")
BRIEF_RE = re.compile(r"^\s*(Eth\S+)\s+([A-Za-z_-]+)\s*$")
COUNTER_HEADER_RE = re.compile(r"PTP Packet Counters of Interface\s+(\S+):", re.I)
COUNTER_ROW_RE = re.compile(
    r"^\s*(Announce|Sync|Follow\s*Up|FollowUp|Delay Request|Delay Response|"
    r"PDelay Request|PDelay Response|PDelay Follow\s*Up|PDelay FollowUp|"
    r"Management)\s+(\d+)\s+(\d+)\s*$",
    re.I,
)


def number(value: str, default: float = math.nan) -> float:
    value = value.strip()
    try:
        if value.lower().startswith("0x"):
            return float(int(value, 16))
        return float(value)
    except ValueError:
        return default


def parse_clock(text: str) -> Dict[str, str]:
    result: Dict[str, str] = {}
    in_quality = False

    for line in text.splitlines():
        stripped = line.strip()

        if stripped == "Clock Quality:":
            in_quality = True
            continue

        m = KV_RE.match(line)
        if not m:
            continue

        key = m.group(1).strip()
        value = m.group(2).strip()

        if in_quality and key in {"Class", "Accuracy", "Offset (log variance)"}:
            result[f"Clock Quality {key}"] = value
            continue

        # Any normal top-level line ends the Clock Quality sub-section.
        in_quality = False
        result[key] = value

    return result


def parse_parent(text: str) -> Dict[str, str]:
    result: Dict[str, str] = {}
    in_gm_quality = False

    for line in text.splitlines():
        stripped = line.strip()

        if stripped == "Grandmaster Clock Quality:":
            in_gm_quality = True
            continue

        m = KV_RE.match(line)
        if not m:
            continue

        key = m.group(1).strip()
        value = m.group(2).strip()

        if in_gm_quality and key in {
            "Class",
            "Accuracy",
            "Offset (log variance)",
            "Priority1",
            "Priority2",
        }:
            result[f"GM {key}"] = value
            continue

        if stripped in {"Parent Clock:", "Grandmaster Clock:"}:
            continue

        # Keep the quality context only while parsing its fields.
        if key not in {
            "Class",
            "Accuracy",
            "Offset (log variance)",
            "Priority1",
            "Priority2",
        }:
            in_gm_quality = False

        result[key] = value

    return result


def parse_brief(text: str) -> List[Tuple[str, str]]:
    rows: List[Tuple[str, str]] = []
    for line in text.splitlines():
        m = BRIEF_RE.match(line)
        if m:
            rows.append((m.group(1), m.group(2).lower()))
    return rows


def normalize_message_type(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_")


def parse_counters(text: str) -> List[Tuple[str, str, float, float]]:
    """
    Returns (interface, message_type, tx, rx) tuples.
    """
    rows: List[Tuple[str, str, float, float]] = []
    current_interface: Optional[str] = None

    for line in text.splitlines():
        m = COUNTER_HEADER_RE.search(line)
        if m:
            current_interface = m.group(1).rstrip(":")
            continue

        if not current_interface:
            continue

        m = COUNTER_ROW_RE.match(line)
        if m:
            rows.append(
                (
                    current_interface,
                    normalize_message_type(m.group(1)),
                    float(m.group(2)),
                    float(m.group(3)),
                )
            )
    return rows


class NXOSClient:
    def __init__(
        self,
        host: str,
        port: int,
        username: str,
        password: Optional[str],
        key_file: Optional[str],
        timeout: float,
        accept_new_host_key: bool,
    ):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_file = key_file
        self.timeout = timeout
        self.accept_new_host_key = accept_new_host_key

    def run_commands(self, commands: List[str]) -> Dict[str, str]:
        client = paramiko.SSHClient()
        client.load_system_host_keys()
        if self.accept_new_host_key:
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        else:
            client.set_missing_host_key_policy(paramiko.RejectPolicy())

        client.connect(
            hostname=self.host,
            port=self.port,
            username=self.username,
            password=self.password,
            key_filename=self.key_file,
            timeout=self.timeout,
            banner_timeout=self.timeout,
            auth_timeout=self.timeout,
            look_for_keys=self.key_file is None and self.password is None,
            allow_agent=True,
        )

        results: Dict[str, str] = {}
        try:
            for command in commands:
                stdin, stdout, stderr = client.exec_command(command, timeout=self.timeout)
                output = stdout.read().decode(errors="replace")
                error = stderr.read().decode(errors="replace").strip()
                status = stdout.channel.recv_exit_status()
                if status != 0:
                    raise RuntimeError(
                        f"{command!r} failed with status {status}: {error or output}"
                    )
                results[command] = output
        finally:
            client.close()

        return results


class NexusPTPCollector:
    COMMANDS = [
        "show ptp clock",
        "show ptp parent",
        "show ptp brief",
        "show ptp counters all",
    ]

    def __init__(self, nxos: NXOSClient, target_name: str, cache_seconds: float):
        self.nxos = nxos
        self.target_name = target_name
        self.cache_seconds = cache_seconds
        self.lock = threading.Lock()
        self.last_fetch = 0.0
        self.cached: Optional[Dict[str, str]] = None
        self.last_fetch_success = 0.0
        self.last_error = ""

    def _fetch(self) -> Dict[str, str]:
        now = time.monotonic()
        with self.lock:
            if self.cached is not None and (now - self.last_fetch) < self.cache_seconds:
                return self.cached
            try:
                self.cached = self.nxos.run_commands(self.COMMANDS)
                self.last_fetch_success = 1.0
                self.last_error = ""
            except Exception as exc:
                self.last_fetch_success = 0.0
                self.last_error = type(exc).__name__
                if self.cached is None:
                    raise
            finally:
                self.last_fetch = time.monotonic()
            return self.cached or {}

    def collect(self):
        start = time.monotonic()

        try:
            outputs = self._fetch()
            success = self.last_fetch_success
        except Exception as exc:
            outputs = {}
            success = 0.0
            self.last_error = type(exc).__name__

        clock = parse_clock(outputs.get("show ptp clock", ""))
        parent = parse_parent(outputs.get("show ptp parent", ""))
        ports = parse_brief(outputs.get("show ptp brief", ""))
        counters = parse_counters(outputs.get("show ptp counters all", ""))

        target_info = InfoMetricFamily(
            "ptp_nexus_target",
            "Nexus PTP exporter target",
            labels=["target", "host"],
        )
        target_info.add_metric([self.target_name, self.nxos.host], {})
        yield target_info

        if clock:
            info = InfoMetricFamily(
                "ptp_nexus_local_clock",
                "Nexus local PTP clock identity and source address",
                labels=["clock_identity", "source_ip", "device_type"],
            )
            info.add_metric(
                [
                    clock.get("Clock Identity", ""),
                    clock.get("PTP Source IP Address", ""),
                    clock.get("PTP Device Type", ""),
                ],
                {},
            )
            yield info

            locked = GaugeMetricFamily(
                "ptp_nexus_clock_locked",
                "1 when the Nexus PTP clock state is Locked",
            )
            locked.add_metric([], 1.0 if clock.get("PTP Clock state", "").lower() == "locked" else 0.0)
            yield locked

            fields = [
                ("ptp_nexus_offset_from_master_nanoseconds", "Offset From Master", "Nexus offset from immediate PTP master"),
                ("ptp_nexus_mean_path_delay_nanoseconds", "Mean Path Delay", "Nexus mean path delay"),
                ("ptp_nexus_steps_removed", "Steps removed", "Boundary-clock steps from grandmaster"),
                ("ptp_nexus_domain_number", "Clock Domain", "PTP domain"),
                ("ptp_nexus_priority1", "Priority1", "Nexus PTP priority1"),
                ("ptp_nexus_priority2", "Priority2", "Nexus PTP priority2"),
                ("ptp_nexus_clock_class", "Clock Quality Class", "Nexus local clockClass"),
                ("ptp_nexus_clock_accuracy_code", "Clock Quality Accuracy", "Nexus clockAccuracy code"),
                ("ptp_nexus_offset_scaled_log_variance", "Clock Quality Offset (log variance)", "Nexus offsetScaledLogVariance"),
                ("ptp_nexus_port_count", "Number of PTP ports", "Number of PTP-enabled ports"),
            ]
            for metric_name, field, help_text in fields:
                if field in clock:
                    g = GaugeMetricFamily(metric_name, help_text)
                    g.add_metric([], number(clock[field]))
                    yield g

        if parent:
            info = InfoMetricFamily(
                "ptp_nexus_parent",
                "Nexus immediate PTP parent and grandmaster",
                labels=[
                    "parent_clock_identity",
                    "parent_port_number",
                    "parent_ip",
                    "grandmaster_clock_identity",
                ],
            )
            info.add_metric(
                [
                    parent.get("Parent Clock Identity", ""),
                    parent.get("Parent Port Number", ""),
                    parent.get("Parent IP", ""),
                    parent.get("Grandmaster Clock Identity", ""),
                ],
                {},
            )
            yield info

            fields = [
                ("ptp_nexus_grandmaster_clock_class", "GM Class", "Selected grandmaster clockClass"),
                ("ptp_nexus_grandmaster_clock_accuracy_code", "GM Accuracy", "Selected grandmaster clockAccuracy code"),
                ("ptp_nexus_grandmaster_offset_scaled_log_variance", "GM Offset (log variance)", "Selected grandmaster offsetScaledLogVariance"),
                ("ptp_nexus_grandmaster_priority1", "GM Priority1", "Selected grandmaster priority1"),
                ("ptp_nexus_grandmaster_priority2", "GM Priority2", "Selected grandmaster priority2"),
            ]
            for metric_name, field, help_text in fields:
                if field in parent:
                    g = GaugeMetricFamily(metric_name, help_text)
                    g.add_metric([], number(parent[field]))
                    yield g

        if ports:
            metric = GaugeMetricFamily(
                "ptp_nexus_port_state",
                "Current Nexus PTP port state; one sample with value 1 per port",
                labels=["interface", "state"],
            )
            for interface, state in ports:
                metric.add_metric([interface, state], 1.0)
            yield metric

        if counters:
            metric = CounterMetricFamily(
                "ptp_nexus_port_messages",
                "PTP packet counters reported by NX-OS",
                labels=["interface", "direction", "message_type"],
            )
            for interface, message_type, tx, rx in counters:
                metric.add_metric([interface, "tx", message_type], tx)
                metric.add_metric([interface, "rx", message_type], rx)
            yield metric

        scrape = GaugeMetricFamily(
            "ptp_nexus_exporter_scrape_success",
            "1 if the most recent NX-OS SSH collection succeeded",
            labels=["target", "error"],
        )
        scrape.add_metric([self.target_name, self.last_error or "none"], success)
        yield scrape

        duration = GaugeMetricFamily(
            "ptp_nexus_exporter_scrape_duration_seconds",
            "Time spent serving this exporter collection",
            labels=["target"],
        )
        duration.add_metric([self.target_name], time.monotonic() - start)
        yield duration


def main() -> None:
    p = argparse.ArgumentParser(description="Prometheus exporter for Cisco Nexus PTP")
    p.add_argument("--host", required=True)
    p.add_argument("--target-name", default=None)
    p.add_argument("--ssh-port", type=int, default=22)
    p.add_argument("--username", required=True)
    p.add_argument("--key-file", default=None)
    p.add_argument(
        "--password-env",
        default=None,
        help="Read SSH password from this environment variable",
    )
    p.add_argument("--ssh-timeout", type=float, default=5.0)
    p.add_argument("--cache-seconds", type=float, default=5.0)
    p.add_argument("--accept-new-host-key", action="store_true")
    p.add_argument("--listen", default="0.0.0.0")
    p.add_argument("--port", type=int, default=9560)
    args = p.parse_args()

    password = os.environ.get(args.password_env) if args.password_env else None

    nxos = NXOSClient(
        host=args.host,
        port=args.ssh_port,
        username=args.username,
        password=password,
        key_file=args.key_file,
        timeout=args.ssh_timeout,
        accept_new_host_key=args.accept_new_host_key,
    )
    collector = NexusPTPCollector(
        nxos=nxos,
        target_name=args.target_name or args.host,
        cache_seconds=args.cache_seconds,
    )
    REGISTRY.register(collector)

    start_http_server(args.port, addr=args.listen)
    print(
        f"Nexus PTP exporter listening on {args.listen}:{args.port}, target={args.host}",
        flush=True,
    )
    while True:
        time.sleep(3600)


if __name__ == "__main__":
    main()