"""Layer 2: TLS everywhere — pqx hybrid certificate management.

Uses pqx CLI to generate post-quantum hybrid certificates
(ECDSA + ML-DSA-65) for node-to-node mTLS.

Certificate hierarchy:
    CA (hybrid) → Node certs (per worker) → Client certs (per service)

The generated certs work with:
- Python's ssl module (classical component)
- CockroachDB (classical component)
- Go services with pqx library (full hybrid verification)
"""
from __future__ import annotations

import json
import logging
import os
import shutil
import ssl
import subprocess
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# Default cert directory
DEFAULT_CERTS_DIR = Path.home() / ".greymatter" / "certs"


def _find_pqx() -> Optional[str]:
    """Find the pqx binary."""
    # Check common locations
    candidates = [
        shutil.which("pqx"),
        str(Path.home() / "bin" / "pqx"),
        str(Path.home() / "go" / "bin" / "pqx"),
        "/usr/local/bin/pqx",
    ]
    # Check env override
    env_path = os.environ.get("PQX_BIN")
    if env_path:
        candidates.insert(0, env_path)

    for path in candidates:
        if path and os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    return None


def _run_pqx(args: list[str], certs_dir: Optional[Path] = None) -> subprocess.CompletedProcess:
    """Run pqx CLI command."""
    pqx = _find_pqx()
    if not pqx:
        raise FileNotFoundError(
            "pqx binary not found. Install: go install github.com/hb69dpsgw4-source/pqx/cmd/pqx@latest"
        )

    cmd = [pqx] + args
    if certs_dir:
        cmd.extend(["--certs-dir", str(certs_dir)])

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    if result.returncode != 0:
        raise RuntimeError(f"pqx failed: {result.stderr}")
    return result


def generate_ca(
    certs_dir: Optional[Path] = None,
    algo: str = "hybrid",
    level: int = 65,
) -> Path:
    """Generate a hybrid CA certificate using pqx.

    Creates:
        certs_dir/ca.crt  — CA certificate (ECDSA + ML-DSA-65)
        certs_dir/ca.key  — CA private key

    Args:
        certs_dir: Directory for cert files. Default ~/.greymatter/certs/
        algo: Certificate algorithm (hybrid, composite, classical).
        level: ML-DSA security level (44, 65, 87).

    Returns:
        Path to the CA cert file.
    """
    certs_dir = certs_dir or DEFAULT_CERTS_DIR
    certs_dir.mkdir(parents=True, exist_ok=True)

    _run_pqx([
        "cert", "create-ca",
        "--algo", algo,
        "--level", str(level),
    ], certs_dir=certs_dir)

    ca_cert = certs_dir / "ca.crt"
    ca_key = certs_dir / "ca.key"

    # Restrict permissions
    if ca_key.exists():
        ca_key.chmod(0o600)

    logger.info("Generated hybrid CA cert at %s", ca_cert)
    return ca_cert


def generate_node_cert(
    hostnames: list[str],
    certs_dir: Optional[Path] = None,
) -> Path:
    """Generate a node certificate signed by the CA.

    Creates:
        certs_dir/node.crt  — Node certificate
        certs_dir/node.key  — Node private key

    Args:
        hostnames: List of hostnames/IPs for the SAN field.
        certs_dir: Directory containing ca.crt and ca.key.

    Returns:
        Path to the node cert file.
    """
    certs_dir = certs_dir or DEFAULT_CERTS_DIR

    _run_pqx([
        "cert", "create-node",
    ] + hostnames, certs_dir=certs_dir)

    node_key = certs_dir / "node.key"
    if node_key.exists():
        node_key.chmod(0o600)

    logger.info("Generated node cert for %s", hostnames)
    return certs_dir / "node.crt"


def generate_client_cert(
    username: str,
    certs_dir: Optional[Path] = None,
) -> Path:
    """Generate a client certificate for service-to-service auth.

    Creates:
        certs_dir/client.<username>.crt
        certs_dir/client.<username>.key

    Args:
        username: Client identity (e.g., "root", "worker-2").
        certs_dir: Directory containing ca.crt and ca.key.

    Returns:
        Path to the client cert file.
    """
    certs_dir = certs_dir or DEFAULT_CERTS_DIR

    _run_pqx([
        "cert", "create-client", username,
    ], certs_dir=certs_dir)

    client_key = certs_dir / f"client.{username}.key"
    if client_key.exists():
        client_key.chmod(0o600)

    logger.info("Generated client cert for %s", username)
    return certs_dir / f"client.{username}.crt"


def create_ssl_context(
    certs_dir: Optional[Path] = None,
    server: bool = True,
    client_cert: Optional[str] = None,
) -> ssl.SSLContext:
    """Create an SSL context from pqx-generated certificates.

    The classical component (ECDSA) of the hybrid certs is used
    by Python's ssl module. The post-quantum component (ML-DSA) is
    verified by Go services using the pqx library.

    Args:
        certs_dir: Directory containing ca.crt, node.crt, node.key.
        server: If True, create server context. If False, client context.
        client_cert: Username for client cert (e.g., "root").

    Returns:
        Configured ssl.SSLContext.
    """
    certs_dir = certs_dir or DEFAULT_CERTS_DIR

    if server:
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ctx.load_cert_chain(
            certfile=str(certs_dir / "node.crt"),
            keyfile=str(certs_dir / "node.key"),
        )
    else:
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
        if client_cert:
            ctx.load_cert_chain(
                certfile=str(certs_dir / f"client.{client_cert}.crt"),
                keyfile=str(certs_dir / f"client.{client_cert}.key"),
            )

    # Load CA for verification
    ca_cert = certs_dir / "ca.crt"
    if ca_cert.exists():
        ctx.load_verify_locations(str(ca_cert))
        ctx.verify_mode = ssl.CERT_REQUIRED

    # Enforce TLS 1.3 (required for ML-KEM key exchange)
    ctx.minimum_version = ssl.TLSVersion.TLSv1_3

    return ctx


def verify_cert(cert_path: Path) -> dict:
    """Verify a certificate using pqx and return cert info.

    Returns dict with subject, issuer, validity, PQC parameters.
    """
    result = _run_pqx(["cert", "info", str(cert_path)])
    return {"raw": result.stdout, "valid": True}


def setup_cluster_certs(
    nodes: list[dict],
    certs_dir: Optional[Path] = None,
) -> dict:
    """Generate a complete certificate set for the cluster.

    Creates CA + node certs + client certs for all nodes.

    Args:
        nodes: List of dicts with 'id', 'hostnames' keys.
        certs_dir: Base directory for certificates.

    Returns:
        Dict mapping node_id → cert paths.
    """
    certs_dir = certs_dir or DEFAULT_CERTS_DIR
    result = {}

    # Generate CA
    generate_ca(certs_dir)

    # Generate per-node certs
    for node in nodes:
        node_id = node["id"]
        hostnames = node.get("hostnames", [])
        node_dir = certs_dir / node_id
        node_dir.mkdir(parents=True, exist_ok=True)

        # Copy CA cert to node dir
        shutil.copy2(certs_dir / "ca.crt", node_dir / "ca.crt")
        shutil.copy2(certs_dir / "ca.key", node_dir / "ca.key")

        generate_node_cert(hostnames, certs_dir=node_dir)
        generate_client_cert(node_id, certs_dir=node_dir)

        result[node_id] = {
            "ca": str(node_dir / "ca.crt"),
            "cert": str(node_dir / "node.crt"),
            "key": str(node_dir / "node.key"),
            "client_cert": str(node_dir / f"client.{node_id}.crt"),
            "client_key": str(node_dir / f"client.{node_id}.key"),
        }

    logger.info("Generated cluster certs for %d nodes", len(nodes))
    return result
