"""
BIBIVIDI — Certificate management
Generates a local CA and per-domain server certificates for HTTPS interception.
Certs are stored in ~/.bibividi/.
"""
from __future__ import annotations

import datetime
import ipaddress
import logging
from pathlib import Path

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.x509 import Certificate
from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID

from bvd.config import settings

logger = logging.getLogger(__name__)

# In-memory cache of per-domain certs to avoid regenerating on every request
_domain_cert_cache: dict[str, tuple[bytes, bytes]] = {}


# ── Path helpers ──────────────────────────────────────────────────────────────

def ca_cert_path() -> Path:
    return settings.data_dir / "ca.crt"


def ca_key_path() -> Path:
    return settings.data_dir / "ca.key"


def server_cert_path() -> Path:
    return settings.data_dir / "server.crt"


def server_key_path() -> Path:
    return settings.data_dir / "server.key"


def ca_exists() -> bool:
    return ca_cert_path().exists() and ca_key_path().exists()


def server_cert_exists() -> bool:
    return server_cert_path().exists() and server_key_path().exists()


# ── Internal helpers ──────────────────────────────────────────────────────────

def _now() -> datetime.datetime:
    return datetime.datetime.now(datetime.UTC)


def _gen_key() -> RSAPrivateKey:
    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


def _key_pem(key: RSAPrivateKey) -> bytes:
    return key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption(),
    )


def _save_key(key: RSAPrivateKey, path: Path) -> None:
    path.write_bytes(_key_pem(key))
    path.chmod(0o600)


def _save_cert(cert: Certificate, path: Path) -> None:
    path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))


# ── CA generation ─────────────────────────────────────────────────────────────

def generate_ca() -> tuple[Certificate, RSAPrivateKey]:
    """Generate a local CA cert+key. Saved to ~/.bibividi/ca.{crt,key}."""
    key = _gen_key()
    now = _now()
    name = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, "BIBIVIDI Local CA"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "BIBIVIDI"),
    ])
    cert = (
        x509.CertificateBuilder()
        .subject_name(name)
        .issuer_name(name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now)
        .not_valid_after(now + datetime.timedelta(days=3650))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(
            x509.KeyUsage(
                digital_signature=True, key_cert_sign=True, crl_sign=True,
                content_commitment=False, key_encipherment=False,
                data_encipherment=False, key_agreement=False,
                encipher_only=False, decipher_only=False,
            ),
            critical=True,
        )
        .sign(key, hashes.SHA256())
    )
    _save_cert(cert, ca_cert_path())
    _save_key(key, ca_key_path())
    logger.info("CA cert generated: %s", ca_cert_path())
    return cert, key


def generate_server_cert(
    ca_cert: Certificate,
    ca_key: RSAPrivateKey,
    hostnames: list[str] | None = None,
) -> None:
    """Generate a localhost server cert signed by our CA."""
    if hostnames is None:
        hostnames = ["localhost", "127.0.0.1", "::1"]

    key = _gen_key()
    now = _now()
    san: list[x509.GeneralName] = []
    for h in hostnames:
        try:
            san.append(x509.IPAddress(ipaddress.ip_address(h)))
        except ValueError:
            san.append(x509.DNSName(h))

    cert = (
        x509.CertificateBuilder()
        .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]))
        .issuer_name(ca_cert.subject)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now)
        .not_valid_after(now + datetime.timedelta(days=825))
        .add_extension(x509.SubjectAlternativeName(san), critical=False)
        .add_extension(
            x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]),
            critical=False,
        )
        .sign(ca_key, hashes.SHA256())
    )
    _save_cert(cert, server_cert_path())
    _save_key(key, server_key_path())
    logger.info("Server cert generated: %s", server_cert_path())


def get_or_generate_certs() -> None:
    """Ensure CA + server cert exist, generating both if needed."""
    if not ca_exists():
        ca_cert, ca_key = generate_ca()
    else:
        ca_cert, ca_key = _load_ca()
    if not server_cert_exists():
        generate_server_cert(ca_cert, ca_key)


# ── Per-domain cert for MITM ──────────────────────────────────────────────────

def cert_pem_for_domain(hostname: str) -> tuple[bytes, bytes]:
    """Return (cert_pem, key_pem) for a hostname, generating if not cached.

    Used by the CONNECT tunnel to impersonate AI API hosts.
    """
    if hostname not in _domain_cert_cache:
        _domain_cert_cache[hostname] = _generate_domain_cert(hostname)
    return _domain_cert_cache[hostname]


def _generate_domain_cert(hostname: str) -> tuple[bytes, bytes]:
    if not ca_exists():
        generate_ca()
    ca_cert, ca_key = _load_ca()

    key = _gen_key()
    now = _now()

    try:
        san_entry: x509.GeneralName = x509.IPAddress(ipaddress.ip_address(hostname))
    except ValueError:
        san_entry = x509.DNSName(hostname)

    cert = (
        x509.CertificateBuilder()
        .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, hostname)]))
        .issuer_name(ca_cert.subject)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now)
        .not_valid_after(now + datetime.timedelta(days=825))
        .add_extension(x509.SubjectAlternativeName([san_entry]), critical=False)
        .add_extension(
            x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]),
            critical=False,
        )
        .sign(ca_key, hashes.SHA256())
    )
    return (
        cert.public_bytes(serialization.Encoding.PEM),
        _key_pem(key),
    )


def _load_ca() -> tuple[Certificate, RSAPrivateKey]:
    from cryptography.hazmat.primitives.serialization import load_pem_private_key
    from cryptography.x509 import load_pem_x509_certificate

    ca_cert = load_pem_x509_certificate(ca_cert_path().read_bytes())
    ca_key = load_pem_private_key(ca_key_path().read_bytes(), password=None)
    assert isinstance(ca_key, RSAPrivateKey)
    return ca_cert, ca_key


# ── OS trust store installation ───────────────────────────────────────────────

def install_ca_system() -> str:
    """Install our CA into the OS trust store.

    Returns a string indicating the method used.
    Raises on failure.
    """
    import platform
    import shutil
    import subprocess

    ca_path = ca_cert_path()
    if not ca_path.exists():
        raise FileNotFoundError(
            f"CA cert not found: {ca_path}\nRun 'bibividi cert generate' first."
        )

    system = platform.system()

    if system == "Darwin":
        subprocess.run(
            [
                "security", "add-trusted-cert",
                "-d", "-r", "trustRoot",
                "-k", "/Library/Keychains/System.keychain",
                str(ca_path),
            ],
            check=True,
        )
        return "macOS system keychain"

    if system == "Linux":
        if Path("/usr/local/share/ca-certificates").exists():
            dest = Path("/usr/local/share/ca-certificates/bibividi-ca.crt")
            shutil.copy(ca_path, dest)
            subprocess.run(["update-ca-certificates"], check=True)
            return "Linux (Debian/Ubuntu)"
        if Path("/etc/pki/ca-trust/source/anchors").exists():
            dest = Path("/etc/pki/ca-trust/source/anchors/bibividi-ca.crt")
            shutil.copy(ca_path, dest)
            subprocess.run(["update-ca-trust"], check=True)
            return "Linux (RHEL/Fedora)"
        return "Linux (unknown — copy manually)"

    if system == "Windows":
        import subprocess
        subprocess.run(["certutil", "-addstore", "Root", str(ca_path)], check=True)
        return "Windows certificate store"

    return "unknown OS — install manually"


def ssl_context_for_server():  # type: ignore[return]
    """Return an ssl.SSLContext for uvicorn HTTPS mode."""
    import ssl
    get_or_generate_certs()
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(server_cert_path(), server_key_path())
    return ctx
