"""TLS certificate lifecycle manager for SecureMesh mTLS (TrustShield Pillar 4).

Manages a local Certificate Authority and per-backend client/server certificates
for mutual TLS authentication between the LLMHosts proxy and its backends.

Certificates are stored at ``~/.config/llmhosts/tls/`` with the following layout::

    tls/
        ca.key          # CA private key (Fernet-encrypted at rest)
        ca.crt          # CA certificate (PEM)
        ca.fernet.key   # Fernet key used to encrypt ca.key
        proxy.key       # Proxy client key (PEM)
        proxy.crt       # Proxy client certificate (PEM)
        backends/
            <name>.key  # Backend-specific key (PEM)
            <name>.crt  # Backend-specific certificate (PEM)
        revoked.json    # Serial numbers of revoked certificates
"""

from __future__ import annotations

import json
import logging
import os
import stat
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import TYPE_CHECKING

from cryptography import x509
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID

if TYPE_CHECKING:
    from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePrivateKey

logger = logging.getLogger(__name__)

_DEFAULT_TLS_DIR = Path.home() / ".config" / "llmhosts" / "tls"
_CERT_VALIDITY_DAYS = 365
_CA_VALIDITY_DAYS = 3650  # 10 years for the CA


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class CertInfo:
    """Public metadata for a managed certificate."""

    name: str
    serial: int
    not_before: datetime
    not_after: datetime
    fingerprint: str  # SHA-256 hex digest
    revoked: bool


# ---------------------------------------------------------------------------
# TLS Manager
# ---------------------------------------------------------------------------


class TLSManager:
    """Manages a local CA and per-backend certificates for mTLS.

    Parameters
    ----------
    tls_dir:
        Root directory for TLS material.  Defaults to
        ``~/.config/llmhosts/tls/``.
    fernet_key:
        Optional pre-existing Fernet key (url-safe base64 bytes) for
        encrypting the CA private key.  When *None*, a random key is
        generated and persisted alongside the CA material.
    """

    def __init__(
        self,
        tls_dir: Path | None = None,
        fernet_key: bytes | None = None,
    ) -> None:
        self._tls_dir = tls_dir or _DEFAULT_TLS_DIR
        self._backends_dir = self._tls_dir / "backends"
        self._revoked_path = self._tls_dir / "revoked.json"
        self._fernet_key_path = self._tls_dir / "ca.fernet.key"

        self._fernet_key = fernet_key
        self._fernet: Fernet | None = None

    # ------------------------------------------------------------------
    # Directory & key helpers
    # ------------------------------------------------------------------

    @property
    def tls_dir(self) -> Path:
        """Return the root TLS directory."""
        return self._tls_dir

    def _ensure_dirs(self) -> None:
        """Create TLS directories with restrictive permissions."""
        for d in (self._tls_dir, self._backends_dir):
            d.mkdir(parents=True, exist_ok=True)
            os.chmod(d, stat.S_IRWXU)  # 0o700

    def _get_fernet(self) -> Fernet:
        """Return (and cache) the Fernet instance for CA key encryption."""
        if self._fernet is not None:
            return self._fernet

        if self._fernet_key is not None:
            self._fernet = Fernet(self._fernet_key)
            return self._fernet

        # Try loading from disk
        if self._fernet_key_path.exists():
            raw = self._fernet_key_path.read_bytes().strip()
            self._fernet = Fernet(raw)
            return self._fernet

        # Generate a new key and persist it
        self._ensure_dirs()
        key = Fernet.generate_key()
        self._fernet_key_path.write_bytes(key)
        os.chmod(self._fernet_key_path, stat.S_IRUSR | stat.S_IWUSR)
        self._fernet = Fernet(key)
        logger.info("Generated new Fernet key for CA encryption at %s", self._fernet_key_path)
        return self._fernet

    # ------------------------------------------------------------------
    # Private-key serialization helpers
    # ------------------------------------------------------------------

    def _encrypt_private_key(self, key: EllipticCurvePrivateKey) -> bytes:
        """Serialize and Fernet-encrypt an EC private key."""
        pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        return self._get_fernet().encrypt(pem)

    def _decrypt_private_key(self, blob: bytes) -> EllipticCurvePrivateKey:
        """Decrypt and deserialize an EC private key."""
        pem = self._get_fernet().decrypt(blob)
        key = serialization.load_pem_private_key(pem, password=None)
        if not isinstance(key, ec.EllipticCurvePrivateKey):
            raise TypeError("Expected an EC private key")
        return key

    @staticmethod
    def _write_pem_key(path: Path, key: EllipticCurvePrivateKey) -> None:
        """Write an unencrypted PEM private key with 0o600 permissions."""
        pem = key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
        path.write_bytes(pem)
        os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)

    @staticmethod
    def _write_pem_cert(path: Path, cert: x509.Certificate) -> None:
        """Write a PEM certificate."""
        path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))

    # ------------------------------------------------------------------
    # Revocation list helpers
    # ------------------------------------------------------------------

    def _load_revoked(self) -> set[int]:
        """Load the set of revoked certificate serial numbers."""
        if not self._revoked_path.exists():
            return set()
        try:
            data = json.loads(self._revoked_path.read_text(encoding="utf-8"))
            return set(data.get("serials", []))
        except (json.JSONDecodeError, OSError):
            logger.warning("Corrupt revoked.json -- treating as empty")
            return set()

    def _save_revoked(self, serials: set[int]) -> None:
        """Persist the revoked serial numbers."""
        self._ensure_dirs()
        payload = {"serials": sorted(serials)}
        self._revoked_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")

    # ------------------------------------------------------------------
    # CA management
    # ------------------------------------------------------------------

    def init_ca(self) -> Path:
        """Initialize the local Certificate Authority.

        Creates a self-signed CA certificate and an ECDSA P-256 private key.
        The private key is Fernet-encrypted at rest.  If a CA already exists
        this method is a no-op and returns the existing CA cert path.

        Returns
        -------
        Path
            Path to the CA certificate (``ca.crt``).
        """
        self._ensure_dirs()
        ca_key_path = self._tls_dir / "ca.key"
        ca_crt_path = self._tls_dir / "ca.crt"

        if ca_key_path.exists() and ca_crt_path.exists():
            logger.info("CA already exists at %s -- skipping init", self._tls_dir)
            return ca_crt_path

        # Generate ECDSA P-256 key
        ca_key = ec.generate_private_key(ec.SECP256R1())

        # Build self-signed CA certificate
        subject = issuer = x509.Name(
            [
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, "LLMHosts"),
                x509.NameAttribute(NameOID.COMMON_NAME, "LLMHosts SecureMesh CA"),
            ]
        )

        now = datetime.now(timezone.utc)
        ca_cert = (
            x509.CertificateBuilder()
            .subject_name(subject)
            .issuer_name(issuer)
            .public_key(ca_key.public_key())
            .serial_number(x509.random_serial_number())
            .not_valid_before(now)
            .not_valid_after(now + timedelta(days=_CA_VALIDITY_DAYS))
            .add_extension(x509.BasicConstraints(ca=True, path_length=0), 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(ca_key, hashes.SHA256())
        )

        # Persist encrypted CA key and CA cert
        encrypted_key = self._encrypt_private_key(ca_key)
        ca_key_path.write_bytes(encrypted_key)
        os.chmod(ca_key_path, stat.S_IRUSR | stat.S_IWUSR)

        self._write_pem_cert(ca_crt_path, ca_cert)

        logger.info("Initialized SecureMesh CA at %s", self._tls_dir)
        return ca_crt_path

    def _load_ca(self) -> tuple[EllipticCurvePrivateKey, x509.Certificate]:
        """Load the CA key (decrypting) and certificate from disk.

        Raises
        ------
        FileNotFoundError
            If the CA has not been initialized.
        """
        ca_key_path = self._tls_dir / "ca.key"
        ca_crt_path = self._tls_dir / "ca.crt"

        if not ca_key_path.exists() or not ca_crt_path.exists():
            raise FileNotFoundError("CA not initialized. Call init_ca() first.")

        ca_key = self._decrypt_private_key(ca_key_path.read_bytes())
        ca_cert = x509.load_pem_x509_certificate(ca_crt_path.read_bytes())
        return ca_key, ca_cert

    # ------------------------------------------------------------------
    # Certificate issuance
    # ------------------------------------------------------------------

    def issue_cert(self, name: str) -> tuple[Path, Path]:
        """Issue a certificate signed by the local CA.

        Creates ``<name>.key`` and ``<name>.crt`` in the appropriate directory
        (``tls/`` for ``proxy``, ``tls/backends/`` for everything else).

        Parameters
        ----------
        name:
            Certificate common name.  ``"proxy"`` is treated specially.

        Returns
        -------
        tuple[Path, Path]
            ``(key_path, cert_path)`` of the newly issued certificate.
        """
        ca_key, ca_cert = self._load_ca()

        # Determine output directory
        out_dir = self._tls_dir if name == "proxy" else self._backends_dir
        out_dir.mkdir(parents=True, exist_ok=True)

        key_path = out_dir / f"{name}.key"
        crt_path = out_dir / f"{name}.crt"

        # Generate ECDSA P-256 key for the entity
        entity_key = ec.generate_private_key(ec.SECP256R1())

        subject = x509.Name(
            [
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, "LLMHosts"),
                x509.NameAttribute(NameOID.COMMON_NAME, name),
            ]
        )

        now = datetime.now(timezone.utc)
        cert = (
            x509.CertificateBuilder()
            .subject_name(subject)
            .issuer_name(ca_cert.subject)
            .public_key(entity_key.public_key())
            .serial_number(x509.random_serial_number())
            .not_valid_before(now)
            .not_valid_after(now + timedelta(days=_CERT_VALIDITY_DAYS))
            .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
            .add_extension(
                x509.KeyUsage(
                    digital_signature=True,
                    key_encipherment=True,
                    content_commitment=False,
                    data_encipherment=False,
                    key_agreement=False,
                    key_cert_sign=False,
                    crl_sign=False,
                    encipher_only=False,
                    decipher_only=False,
                ),
                critical=True,
            )
            .add_extension(
                x509.ExtendedKeyUsage(
                    [
                        x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
                        x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
                    ]
                ),
                critical=False,
            )
            .add_extension(
                x509.SubjectAlternativeName(
                    [
                        x509.DNSName(name),
                        x509.DNSName("localhost"),
                    ]
                ),
                critical=False,
            )
            .sign(ca_key, hashes.SHA256())
        )

        self._write_pem_key(key_path, entity_key)
        self._write_pem_cert(crt_path, cert)

        logger.info("Issued certificate for '%s' (serial=%d)", name, cert.serial_number)
        return key_path, crt_path

    # ------------------------------------------------------------------
    # Certificate listing
    # ------------------------------------------------------------------

    def list_certs(self) -> list[CertInfo]:
        """List all managed certificates with metadata.

        Returns
        -------
        list[CertInfo]
            Metadata for every ``.crt`` file under the TLS directory tree.
        """
        revoked_serials = self._load_revoked()
        certs: list[CertInfo] = []

        # Scan top-level dir (ca.crt, proxy.crt) and backends/
        for search_dir in (self._tls_dir, self._backends_dir):
            if not search_dir.exists():
                continue
            for crt_path in sorted(search_dir.glob("*.crt")):
                try:
                    cert = x509.load_pem_x509_certificate(crt_path.read_bytes())
                    fp = cert.fingerprint(hashes.SHA256()).hex()
                    name = crt_path.stem
                    certs.append(
                        CertInfo(
                            name=name,
                            serial=cert.serial_number,
                            not_before=cert.not_valid_before_utc,
                            not_after=cert.not_valid_after_utc,
                            fingerprint=fp,
                            revoked=cert.serial_number in revoked_serials,
                        )
                    )
                except Exception:
                    logger.warning("Failed to parse certificate: %s", crt_path)

        return certs

    # ------------------------------------------------------------------
    # Certificate rotation
    # ------------------------------------------------------------------

    def rotate_cert(self, name: str) -> tuple[Path, Path]:
        """Rotate a certificate by issuing a new one.

        The old certificate is **not** immediately revoked so that in-flight
        connections can drain.  Callers should revoke the old cert after a
        grace period if desired.

        Parameters
        ----------
        name:
            Certificate common name to rotate.

        Returns
        -------
        tuple[Path, Path]
            ``(key_path, cert_path)`` of the newly issued certificate.
        """
        logger.info("Rotating certificate for '%s'", name)
        return self.issue_cert(name)

    # ------------------------------------------------------------------
    # Certificate revocation
    # ------------------------------------------------------------------

    def revoke_cert(self, name: str) -> bool:
        """Revoke a certificate by adding its serial to the revocation list.

        Parameters
        ----------
        name:
            Certificate common name to revoke.

        Returns
        -------
        bool
            *True* if the certificate was found and revoked, *False* otherwise.
        """
        # Find the cert file
        crt_path = self._tls_dir / f"{name}.crt" if name == "proxy" else self._backends_dir / f"{name}.crt"

        if not crt_path.exists():
            logger.warning("Certificate not found for revocation: %s", name)
            return False

        cert = x509.load_pem_x509_certificate(crt_path.read_bytes())
        revoked = self._load_revoked()
        revoked.add(cert.serial_number)
        self._save_revoked(revoked)

        logger.info("Revoked certificate '%s' (serial=%d)", name, cert.serial_number)
        return True

    def is_revoked(self, serial: int) -> bool:
        """Check whether a certificate serial number is in the revocation list."""
        return serial in self._load_revoked()

    # ------------------------------------------------------------------
    # Certificate paths
    # ------------------------------------------------------------------

    def ca_cert_path(self) -> Path:
        """Return the path to the CA certificate."""
        return self._tls_dir / "ca.crt"

    def cert_paths(self, name: str) -> tuple[Path, Path]:
        """Return ``(key_path, cert_path)`` for a named certificate."""
        d = self._tls_dir if name == "proxy" else self._backends_dir
        return d / f"{name}.key", d / f"{name}.crt"
