#!/usr/bin/env python3
"""Create a privacy-conscious capability report for a prospective WhoLab worker node."""

from __future__ import annotations

import argparse
import ctypes
import hashlib
import json
import os
import platform
import shutil
import socket
import subprocess
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

PSEUDO_FILESYSTEMS = {
    "autofs",
    "binfmt_misc",
    "bpf",
    "cgroup",
    "cgroup2",
    "configfs",
    "debugfs",
    "devpts",
    "devtmpfs",
    "efivarfs",
    "fusectl",
    "hugetlbfs",
    "mqueue",
    "overlay",
    "proc",
    "pstore",
    "securityfs",
    "squashfs",
    "sysfs",
    "tmpfs",
    "tracefs",
}


def run(command: list[str], timeout: int = 10) -> dict[str, Any]:
    try:
        process = subprocess.run(
            command,
            capture_output=True,
            text=True,
            timeout=timeout,
            check=False,
        )
        return {
            "available": True,
            "exit_code": process.returncode,
            "stdout": process.stdout.strip(),
            "stderr": process.stderr.strip(),
        }
    except (FileNotFoundError, subprocess.TimeoutExpired) as error:
        return {"available": False, "error": type(error).__name__}


def os_release() -> dict[str, str]:
    path = Path("/etc/os-release")
    if not path.exists():
        return {}
    values = {}
    for line in path.read_text(errors="replace").splitlines():
        if "=" in line:
            key, value = line.split("=", 1)
            values[key.lower()] = value.strip().strip('"')
    return {key: values[key] for key in ("id", "version_id", "pretty_name") if key in values}


def linux_memory() -> dict[str, int]:
    values: dict[str, int] = {}
    path = Path("/proc/meminfo")
    if not path.exists():
        return values
    for line in path.read_text().splitlines():
        key, raw = line.split(":", 1)
        number = int(raw.strip().split()[0]) * 1024
        values[key] = number
    return {
        "total_bytes": values.get("MemTotal", 0),
        "available_bytes": values.get("MemAvailable", 0),
        "swap_total_bytes": values.get("SwapTotal", 0),
        "swap_free_bytes": values.get("SwapFree", 0),
    }


def windows_memory() -> dict[str, int]:
    class MemoryStatus(ctypes.Structure):
        _fields_ = [
            ("length", ctypes.c_ulong),
            ("memory_load", ctypes.c_ulong),
            ("total_physical", ctypes.c_ulonglong),
            ("available_physical", ctypes.c_ulonglong),
            ("total_page_file", ctypes.c_ulonglong),
            ("available_page_file", ctypes.c_ulonglong),
            ("total_virtual", ctypes.c_ulonglong),
            ("available_virtual", ctypes.c_ulonglong),
            ("available_extended_virtual", ctypes.c_ulonglong),
        ]

    status = MemoryStatus()
    status.length = ctypes.sizeof(status)
    ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status))
    return {
        "total_bytes": status.total_physical,
        "available_bytes": status.available_physical,
        "swap_total_bytes": max(0, status.total_page_file - status.total_physical),
        "swap_free_bytes": max(0, status.available_page_file - status.available_physical),
    }


def macos_memory() -> dict[str, int]:
    result = run(["sysctl", "-n", "hw.memsize"])
    total = int(result.get("stdout") or 0)
    return {
        "total_bytes": total,
        "available_bytes": 0,
        "swap_total_bytes": 0,
        "swap_free_bytes": 0,
    }


def memory() -> dict[str, int]:
    system = platform.system()
    if system == "Linux":
        return linux_memory()
    if system == "Windows":
        return windows_memory()
    if system == "Darwin":
        return macos_memory()
    return {}


def cpu_model() -> str:
    if platform.system() == "Linux":
        path = Path("/proc/cpuinfo")
        if path.exists():
            for line in path.read_text(errors="replace").splitlines():
                if line.lower().startswith("model name"):
                    return line.split(":", 1)[1].strip()
    if platform.system() == "Darwin":
        return run(["sysctl", "-n", "machdep.cpu.brand_string"]).get("stdout", "")
    return platform.processor()


def unix_mounts() -> list[tuple[str, str]]:
    mounts: list[tuple[str, str]] = []
    path = Path("/proc/self/mounts")
    if path.exists():
        for line in path.read_text(errors="replace").splitlines():
            fields = line.split()
            if len(fields) >= 3 and fields[2] not in PSEUDO_FILESYSTEMS:
                mounts.append((fields[1].replace("\\040", " "), fields[2]))
        return mounts
    result = run(["df", "-P", "-k"], timeout=15)
    for line in result.get("stdout", "").splitlines()[1:]:
        fields = line.split()
        if len(fields) >= 6:
            mounts.append((fields[-1], "unknown"))
    return mounts


def windows_mounts() -> list[tuple[str, str]]:
    bitmask = ctypes.windll.kernel32.GetLogicalDrives()
    return [(f"{chr(65 + index)}:\\", "windows") for index in range(26) if bitmask & (1 << index)]


def storage() -> list[dict[str, Any]]:
    candidates = windows_mounts() if platform.system() == "Windows" else unix_mounts()
    candidates.extend([(str(Path.home()), "home"), (str(Path.cwd()), "working-directory")])
    results = []
    seen: set[tuple[int, int]] = set()
    for mountpoint, filesystem in candidates:
        try:
            usage = shutil.disk_usage(mountpoint)
            stat = os.stat(mountpoint)
        except (FileNotFoundError, PermissionError, OSError):
            continue
        identity = (stat.st_dev, usage.total)
        if identity in seen:
            continue
        seen.add(identity)
        results.append(
            {
                "mountpoint": mountpoint,
                "filesystem": filesystem,
                "total_bytes": usage.total,
                "free_bytes": usage.free,
            }
        )
    return sorted(results, key=lambda item: item["free_bytes"], reverse=True)


def gpu() -> list[dict[str, Any]]:
    result = run(
        [
            "nvidia-smi",
            "--query-gpu=name,memory.total,driver_version",
            "--format=csv,noheader,nounits",
        ],
        timeout=15,
    )
    if result.get("exit_code") != 0:
        return []
    devices = []
    for line in result["stdout"].splitlines():
        fields = [field.strip() for field in line.split(",")]
        if len(fields) == 3:
            devices.append(
                {
                    "name": fields[0],
                    "memory_mib": int(fields[1]),
                    "driver_version": fields[2],
                }
            )
    return devices


def tool_versions() -> dict[str, Any]:
    commands = {
        "aws": ["aws", "--version"],
        "curl": ["curl", "--version"],
        "docker": ["docker", "--version"],
        "docker_compose": ["docker", "compose", "version"],
        "git": ["git", "--version"],
        "pigz": ["pigz", "--version"],
        "rclone": ["rclone", "version"],
        "wget": ["wget", "--version"],
        "zstd": ["zstd", "--version"],
    }
    tools = {}
    for name, command in commands.items():
        result = run(command)
        output = result.get("stdout") or result.get("stderr") or ""
        tools[name] = {
            "available": result.get("available", False) and result.get("exit_code") == 0,
            "version": output.splitlines()[0] if output else None,
        }
    docker_info = run(["docker", "info", "--format", "{{json .ServerVersion}}"], timeout=15)
    tools["docker"]["daemon_access"] = docker_info.get("exit_code") == 0
    return tools


def probe(url: str) -> dict[str, Any]:
    started = time.monotonic()
    try:
        request = urllib.request.Request(
            url,
            method="HEAD",
            headers={"User-Agent": "WhoLabNodeInventory/1.0"},
        )
        with urllib.request.urlopen(request, timeout=15) as response:
            return {
                "url": url,
                "status": response.status,
                "elapsed_ms": round((time.monotonic() - started) * 1000),
            }
    except Exception as error:
        return {
            "url": url,
            "error": f"{type(error).__name__}: {error}",
            "elapsed_ms": round((time.monotonic() - started) * 1000),
        }


def collect(args: argparse.Namespace) -> dict[str, Any]:
    try:
        load_average = list(os.getloadavg())
    except (AttributeError, OSError):
        load_average = []
    return {
        "schema_version": 1,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "node_label": args.label,
        "hostname": socket.gethostname(),
        "platform": {
            "system": platform.system(),
            "release": platform.release(),
            "architecture": platform.machine(),
            "python": platform.python_version(),
            "os_release": os_release(),
        },
        "cpu": {
            "logical_cores": os.cpu_count(),
            "model": cpu_model(),
            "load_average": load_average,
        },
        "memory": memory(),
        "storage": storage(),
        "gpu": gpu(),
        "tools": tool_versions(),
        "capabilities": {
            "running_as_root": hasattr(os, "geteuid") and os.geteuid() == 0,
            "systemd": shutil.which("systemctl") is not None,
        },
        "network_probes": [probe(url) for url in args.probe_url],
        "privacy": {
            "environment_variables_collected": False,
            "ip_addresses_collected": False,
            "process_list_collected": False,
            "credentials_collected": False,
        },
    }


class ApiError(RuntimeError):
    def __init__(self, status: int, detail: str):
        super().__init__(detail)
        self.status = status
        self.detail = detail


def default_credential_path() -> Path:
    if os.name == "nt":
        root = Path(os.getenv("APPDATA", Path.home()))
    else:
        root = Path(os.getenv("XDG_CONFIG_HOME", Path.home() / ".config"))
    return root / "wholab" / "node.json"


def read_credentials(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text())
    except FileNotFoundError:
        return {}


def write_credentials(path: Path, credentials: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(".tmp")
    temporary.write_text(json.dumps(credentials, indent=2) + "\n")
    if os.name != "nt":
        temporary.chmod(0o600)
    temporary.replace(path)


def api_request(
    url: str, payload: dict[str, Any] | None = None, token: str | None = None
) -> dict[str, Any]:
    headers = {
        "content-type": "application/json",
        "user-agent": "WhoLabNodeInventory/1.0",
    }
    if token:
        headers["authorization"] = f"Bearer {token}"
    request = urllib.request.Request(
        url,
        data=json.dumps(payload, ensure_ascii=False).encode() if payload is not None else None,
        headers=headers,
        method="POST" if payload is not None else "GET",
    )
    try:
        with urllib.request.urlopen(request, timeout=30) as response:
            body = response.read()
            return json.loads(body) if body else {}
    except urllib.error.HTTPError as error:
        try:
            detail = json.loads(error.read()).get("detail", str(error))
        except (json.JSONDecodeError, UnicodeDecodeError):
            detail = str(error)
        raise ApiError(error.code, detail) from error


def enrollment_code(args: argparse.Namespace) -> str | None:
    if args.enrollment_code_file:
        return args.enrollment_code_file.read_text().strip()
    value = os.getenv("WHOLAB_ENROLLMENT_CODE")
    if value:
        return value.strip()
    return None


def update_schedule(credentials: dict[str, Any], response: dict[str, Any]) -> None:
    for key in ("inventory_interval_minutes", "heartbeat_interval_minutes"):
        if key in response:
            credentials[key] = response[key]


def enroll_node(
    args: argparse.Namespace,
    credentials: dict[str, Any],
    path: Path,
) -> dict[str, Any]:
    if not args.server:
        raise RuntimeError("--server is required for first enrollment")
    report = collect(args)
    payload = {"inventory": report}
    code = enrollment_code(args)
    if code:
        payload["code"] = code
    response = api_request(
        f"{args.server.rstrip('/')}/api/backend/worker-nodes/enroll",
        payload,
    )
    credentials = {
        "server": args.server.rstrip("/"),
        "node_id": response["node_id"],
        "token": response["token"],
        "node_label": args.label,
        "last_inventory_at": int(time.time()),
    }
    update_schedule(credentials, response)
    write_credentials(path, credentials)
    print("Enrolled node: " + credentials["node_id"])
    return credentials


def sync_node(args: argparse.Namespace) -> None:
    path = args.credential_file or default_credential_path()
    credentials = read_credentials(path)
    if not credentials:
        credentials = enroll_node(args, credentials, path)
        credentials["last_heartbeat_at"] = int(time.time())
        write_credentials(path, credentials)
        print(f"Credentials: {path}")
        return

    server = credentials["server"]
    token = credentials["token"]
    policy = api_request(
        f"{server}/api/backend/worker-nodes/config",
        token=token,
    )
    update_schedule(credentials, policy)
    credentials["uploads_enabled"] = policy["uploads_enabled"]
    write_credentials(path, credentials)
    if not policy["uploads_enabled"]:
        print("Upload paused by administrator")
        return

    now = int(time.time())
    heartbeat_interval = int(credentials.get("heartbeat_interval_minutes", 5)) * 60
    inventory_interval = int(credentials.get("inventory_interval_minutes", 1440)) * 60
    heartbeat_due = now - int(credentials.get("last_heartbeat_at", 0)) >= heartbeat_interval
    inventory_due = now - int(credentials.get("last_inventory_at", 0)) >= inventory_interval
    if not heartbeat_due and not inventory_due:
        print("Sync skipped; heartbeat and inventory are not due")
        return

    uploaded = []
    if heartbeat_due:
        disks = storage()
        response = api_request(
            f"{server}/api/backend/worker-nodes/heartbeat",
            {
                "status": "idle",
                "free_bytes": max((disk["free_bytes"] for disk in disks), default=0),
            },
            token,
        )
        update_schedule(credentials, response)
        credentials["last_heartbeat_at"] = now
        uploaded.append("heartbeat")

    if inventory_due:
        response = api_request(
            f"{server}/api/backend/worker-nodes/inventory",
            collect(args),
            token,
        )
        update_schedule(credentials, response)
        credentials["last_inventory_at"] = now
        uploaded.append("inventory")

    write_credentials(path, credentials)
    print("Uploaded: " + ", ".join(uploaded))


def write_report(args: argparse.Namespace) -> None:
    report = collect(args)
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    safe_label = "".join(
        character if character.isalnum() or character in "-_" else "-" for character in args.label
    )
    output = args.output or Path(f"wholab-node-{safe_label}-{timestamp}.json")
    output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n")
    if os.name != "nt":
        output.chmod(0o600)
    digest = hashlib.sha256(output.read_bytes()).hexdigest()
    print(f"Report: {output.resolve()}")
    print(f"SHA256: {digest}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--label", default=socket.gethostname(), help="Human-readable node label")
    parser.add_argument("--output", type=Path, help="Output JSON path")
    parser.add_argument(
        "--probe-url",
        action="append",
        default=[],
        help="Optional HTTPS URL to test with one HEAD request; may be repeated",
    )
    parser.add_argument("--sync", action="store_true", help="Enroll or sync with WhoLab")
    parser.add_argument("--server", help="WhoLab base URL, required for first enrollment")
    parser.add_argument("--credential-file", type=Path, help="Override local credential path")
    parser.add_argument(
        "--enrollment-code-file",
        type=Path,
        help="Read an optional permanent invitation code from a file",
    )
    args = parser.parse_args()
    try:
        if args.sync:
            sync_node(args)
        else:
            write_report(args)
    except ApiError as error:
        if error.status == 403:
            print(f"Upload paused by administrator: {error.detail}")
            return
        raise SystemExit(f"WhoLab API error ({error.status}): {error.detail}") from error
    except (OSError, RuntimeError, urllib.error.URLError) as error:
        raise SystemExit(f"Node sync failed: {error}") from error


if __name__ == "__main__":
    main()