#!/usr/bin/env python3
import argparse
import hashlib
import json
import os
import platform
import re
import socket
import ssl
import subprocess
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


VERSION = "0.1.0"
VM_INTERFACE_RE = re.compile(r"(?:tap|fwbr|fwln|fwpr)(\d+)")
VM_INTERFACE_DETAIL_RE = re.compile(r"(?:tap|fwbr|fwln|fwpr)(\d+)i(\d+)")


def read_text(path: str) -> str | None:
    try:
        return Path(path).read_text(encoding="utf-8").strip()
    except OSError:
        return None


def run_command(args: list[str], timeout: int = 5) -> tuple[int, str]:
    try:
        result = subprocess.run(args, capture_output=True, text=True, timeout=timeout, check=False)
        output = result.stdout.strip() or result.stderr.strip()
        return result.returncode, output
    except (OSError, subprocess.SubprocessError):
        return 127, ""


def collect_interfaces() -> list[dict[str, Any]]:
    interfaces = []
    for item in Path("/sys/class/net").iterdir() if Path("/sys/class/net").exists() else []:
        name = item.name
        if name == "lo":
            continue
        vm_match = VM_INTERFACE_RE.search(name)
        interfaces.append(
            {
                "name": name,
                "vmid": vm_match.group(1) if vm_match else None,
                "operstate": read_text(f"/sys/class/net/{name}/operstate") or "unknown",
                "rx_bytes": int(read_text(f"/sys/class/net/{name}/statistics/rx_bytes") or 0),
                "tx_bytes": int(read_text(f"/sys/class/net/{name}/statistics/tx_bytes") or 0),
                "rx_packets": int(read_text(f"/sys/class/net/{name}/statistics/rx_packets") or 0),
                "tx_packets": int(read_text(f"/sys/class/net/{name}/statistics/tx_packets") or 0),
            }
        )
    return interfaces


def interface_rank(name: str) -> int:
    if name.startswith("tap"):
        return 0
    if name.startswith("fwln"):
        return 1
    if name.startswith("fwpr"):
        return 2
    if name.startswith("fwbr"):
        return 3
    return 9


def collect_interface_traffic(interfaces: list[dict[str, Any]]) -> list[dict[str, Any]]:
    selected: dict[tuple[str, str], dict[str, Any]] = {}
    for interface in interfaces:
        vmid = interface.get("vmid")
        if not vmid:
            continue
        detail_match = VM_INTERFACE_DETAIL_RE.search(str(interface.get("name") or ""))
        nic = detail_match.group(2) if detail_match else "0"
        key = (str(vmid), nic)
        current = selected.get(key)
        if current and interface_rank(str(current.get("name") or "")) <= interface_rank(str(interface.get("name") or "")):
            continue
        selected[key] = interface

    traffic = []
    for (vmid, nic), interface in sorted(selected.items()):
        rx_bytes = int(interface.get("rx_bytes") or 0)
        tx_bytes = int(interface.get("tx_bytes") or 0)
        rx_packets = int(interface.get("rx_packets") or 0)
        tx_packets = int(interface.get("tx_packets") or 0)
        traffic.append(
            {
                "vmid": vmid,
                "nic": nic,
                "interface": interface.get("name"),
                "source": f"vm:{vmid}",
                "destination": "network",
                "protocol": "interface-counter",
                "rx_bytes": rx_bytes,
                "tx_bytes": tx_bytes,
                "bytes": rx_bytes + tx_bytes,
                "rx_packets": rx_packets,
                "tx_packets": tx_packets,
                "packets": rx_packets + tx_packets,
                "state": interface.get("operstate") or "unknown",
            }
        )
    return traffic


def parse_conntrack_line(line: str) -> dict[str, Any] | None:
    parts = line.split()
    if len(parts) < 5 or parts[0] not in {"tcp", "udp", "icmp"}:
        return None
    protocol = parts[0]
    state = None
    if protocol == "tcp" and len(parts) > 3 and "=" not in parts[3]:
        state = parts[3]

    values: dict[str, list[str]] = {}
    for part in parts:
        if "=" not in part:
            continue
        key, value = part.split("=", 1)
        values.setdefault(key, []).append(value)

    src_values = values.get("src", [])
    dst_values = values.get("dst", [])
    if not src_values or not dst_values:
        return None

    packet_values = [int(value) for value in values.get("packets", []) if value.isdigit()]
    byte_values = [int(value) for value in values.get("bytes", []) if value.isdigit()]
    sport = values.get("sport", [None])[0]
    dport = values.get("dport", [None])[0]
    return {
        "source_ip": src_values[0],
        "destination_ip": dst_values[0],
        "protocol": protocol,
        "source_port": int(sport) if sport and sport.isdigit() else None,
        "destination_port": int(dport) if dport and dport.isdigit() else None,
        "packets": sum(packet_values),
        "bytes": sum(byte_values),
        "state": state,
    }


def collect_flows(limit: int = 500) -> list[dict[str, Any]]:
    code, output = run_command(["conntrack", "-L", "-o", "extended"], timeout=10)
    if code != 0 or not output:
        return []
    flows = []
    seen = set()
    for line in output.splitlines():
        flow = parse_conntrack_line(line)
        if not flow:
            continue
        key = (
            flow["source_ip"],
            flow["destination_ip"],
            flow["protocol"],
            flow.get("source_port"),
            flow.get("destination_port"),
        )
        if key in seen:
            continue
        seen.add(key)
        flows.append(flow)
        if len(flows) >= limit:
            break
    return flows


def collect_conntrack() -> dict[str, Any]:
    code, output = run_command(["conntrack", "-C"])
    if code == 0 and output.isdigit():
        return {"count": int(output), "source": "conntrack"}
    for path in ("/proc/net/nf_conntrack", "/proc/net/ip_conntrack"):
        try:
            with open(path, "r", encoding="utf-8", errors="ignore") as handle:
                return {"count": sum(1 for _ in handle), "source": path}
        except OSError:
            continue
    return {"count": None, "source": "unavailable"}


def collect_flow_diagnostics(conntrack: dict[str, Any], flow_count: int) -> dict[str, Any]:
    code, conntrack_path = run_command(["sh", "-c", "command -v conntrack"])
    return {
        "flow_count": flow_count,
        "conntrack_binary": conntrack_path if code == 0 else None,
        "conntrack_count": conntrack.get("count"),
        "conntrack_source": conntrack.get("source"),
        "bridge_nf_call_iptables": read_text("/proc/sys/net/bridge/bridge-nf-call-iptables"),
        "bridge_nf_call_ip6tables": read_text("/proc/sys/net/bridge/bridge-nf-call-ip6tables"),
        "nf_conntrack_max": read_text("/proc/sys/net/netfilter/nf_conntrack_max"),
        "nf_conntrack_count": read_text("/proc/sys/net/netfilter/nf_conntrack_count"),
    }


def collect_firewall() -> dict[str, Any]:
    status = {}
    code, output = run_command(["systemctl", "is-active", "pve-firewall"])
    status["pve_firewall"] = output if code == 0 else "unknown"
    code, output = run_command(["nft", "-j", "list", "ruleset"], timeout=10)
    if code == 0 and output:
        status["nft_ruleset_sha256"] = hashlib.sha256(output.encode("utf-8")).hexdigest()
    else:
        status["nft_ruleset_sha256"] = None
    return status


def collect_payload(config: dict[str, Any]) -> dict[str, Any]:
    uptime = read_text("/proc/uptime")
    interfaces = collect_interfaces()
    flows = collect_flows(int(config.get("flow_limit", 500)))
    conntrack = collect_conntrack()
    return {
        "version": VERSION,
        "node_name": config.get("node_name"),
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "hostname": socket.gethostname(),
        "kernel": platform.release(),
        "uptime_seconds": float(uptime.split()[0]) if uptime else None,
        "loadavg": list(os.getloadavg()) if hasattr(os, "getloadavg") else [],
        "interfaces": interfaces,
        "interface_traffic": collect_interface_traffic(interfaces),
        "flows": flows,
        "conntrack": conntrack,
        "firewall": collect_firewall(),
        "extra": {"platform": platform.platform(), "flow_diagnostics": collect_flow_diagnostics(conntrack, len(flows))},
    }


def normalized_api_url(config: dict[str, Any]) -> str:
    api_url = str(config["api_url"]).rstrip("/")
    while api_url.endswith("/api/v1/api/v1"):
        api_url = api_url.removesuffix("/api/v1")
    if not api_url.endswith("/api/v1"):
        api_url = f"{api_url}/api/v1"
    return api_url


def post_heartbeat(config: dict[str, Any], payload: dict[str, Any]) -> None:
    api_url = normalized_api_url(config)
    heartbeat_url = f"{api_url}/agents/heartbeat"
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        heartbeat_url,
        data=data,
        headers={
            "Authorization": f"Bearer {config['token']}",
            "Content-Type": "application/json",
            "User-Agent": f"nexafabric-agent/{VERSION}",
        },
        method="POST",
    )
    context = None
    if not bool(config.get("verify_tls", True)):
        context = ssl._create_unverified_context()
    with urllib.request.urlopen(request, timeout=15, context=context) as response:
        response.read()


def load_config(path: str) -> dict[str, Any]:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def main() -> int:
    parser = argparse.ArgumentParser(description="NexaFabric Proxmox node telemetry agent")
    parser.add_argument("--config", default="/etc/nexafabric-agent/config.json")
    parser.add_argument("--once", action="store_true")
    args = parser.parse_args()
    config = load_config(args.config)
    interval = int(config.get("interval_seconds", 30))

    while True:
        payload = collect_payload(config)
        try:
            post_heartbeat(config, payload)
            print(f"heartbeat ok: {payload['collected_at']}", flush=True)
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")[:500]
            print(f"heartbeat failed: HTTP {exc.code} {exc.reason} url={exc.url} body={body}", flush=True)
        except (OSError, urllib.error.URLError) as exc:
            print(f"heartbeat failed: {exc} api_url={normalized_api_url(config)}", flush=True)
        if args.once:
            return 0
        time.sleep(interval)


if __name__ == "__main__":
    raise SystemExit(main())