#!/usr/bin/env python3

# SPDX-License-Identifier: GPL-2.0-only
# Copyright (C) 2026 Philipp H.
# Developed with assistance from ChatGPT.


import argparse
import getpass
import json
import os
import re
import sys
import time

import paramiko


PORTS = (
    ("Dark Fibre 200G L1", "0/5/255/1"),
    ("Dark Fibre 200G L2", "0/5/255/2"),
    ("C3 - 100G Core LAG", "0/5/255/5"),
    ("C5 - SAN FC32G", "0/5/255/7"),
    ("C6 - SAN FC32G", "0/5/255/8"),
    ("LIN/LOUT", "0/7/255/1"),
)

INVALID_VALUE = -2147483648

ANSI_RE = re.compile(r"\x1b\[[0-?]*[ -/]*[@-~]")
PROMPT_RE = re.compile(r"<[^>\r\n]+>\s*$")


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Collect optical power from Huawei OptiX DC908 via SSH",
    )

    parser.add_argument(
        "--host",
        required=True,
        help="Hostname or management IP of the multiplexer",
    )

    parser.add_argument(
        "--user",
        default="checkmk",
        help="SSH username, default: checkmk",
    )

    parser.add_argument(
        "--password",
        help="SSH password; if omitted, prompt interactively",
    )

    parser.add_argument(
        "--timeout",
        type=int,
        default=20,
        help="SSH and CLI timeout in seconds, default: 20",
    )

    return parser.parse_args()


def clean_output(text):
    return ANSI_RE.sub("", text).replace("\r", "")


def read_until_prompt(channel, timeout):
    output = ""
    deadline = time.monotonic() + timeout

    while time.monotonic() < deadline:
        if channel.recv_ready():
            data = channel.recv(65535)

            if not data:
                raise RuntimeError("SSH channel closed unexpectedly")

            output += data.decode("utf-8", errors="replace")
            cleaned = clean_output(output)

            if PROMPT_RE.search(cleaned):
                return cleaned
        else:
            time.sleep(0.1)

    raise RuntimeError(
        "Timeout waiting for Huawei CLI prompt. "
        f"Received: {clean_output(output)!r}"
    )


def run_command(channel, command, timeout):
    channel.send(command + "\r")
    return read_until_prompt(channel, timeout)


def parse_power(output, port):
    match = re.search(
        rf"(?m)^{re.escape(port)}[ \t]+(-?\d+)[ \t]*$",
        output,
    )

    if match is None:
        raise RuntimeError(
            f"No optical-power value found for port {port}"
        )

    raw_value = int(match.group(1))

    if raw_value == INVALID_VALUE:
        return None, raw_value

    return raw_value / 10.0, raw_value


def collect_power(host, user, password, timeout):
    client = paramiko.SSHClient()

    # Use the site user's ~/.ssh/known_hosts.
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        client.connect(
            hostname=host,
            username=user,
            password=password,
            look_for_keys=False,
            allow_agent=False,
            timeout=timeout,
            banner_timeout=timeout,
            auth_timeout=timeout,
        )

        channel = client.invoke_shell(
            term="vt100",
            width=200,
            height=1000,
        )

        read_until_prompt(channel, timeout)

        results = []

        for circuit, port in PORTS:
            rx_output = run_command(
                channel,
                f"display optical-inpower port {port}",
                timeout,
            )

            tx_output = run_command(
                channel,
                f"display optical-outpower port {port}",
                timeout,
            )

            rx_dbm, rx_raw = parse_power(rx_output, port)
            tx_dbm, tx_raw = parse_power(tx_output, port)

            results.append({
                "circuit": circuit,
                "port": port,
                "rx_dbm": rx_dbm,
                "tx_dbm": tx_dbm,
                "rx_raw": rx_raw,
                "tx_raw": tx_raw,
            })

        channel.send("quit\r")
        return results

    finally:
        client.close()


def main():
    args = parse_arguments()

    password = (
        args.password
        or os.environ.get("OPTIX_PASSWORD")
        or getpass.getpass(f"{args.user}@{args.host} password: ")
    )

    last_error = None

    for attempt in range(1, 4):
        try:
            results = collect_power(
                host=args.host,
                user=args.user,
                password=password,
                timeout=args.timeout,
            )
            break
        except Exception as exc:
            last_error = exc

            if attempt < 3:
                time.sleep(3)
    else:
        print(
            f"OptiX SSH collection failed after 3 attempts: {last_error}",
            file=sys.stderr,
        )
        return 2

    print("<<<optix_power:sep(0)>>>")

    for result in results:
        print(
            json.dumps(
                result,
                separators=(",", ":"),
                sort_keys=True,
            )
        )

    return 0


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