"""Tests for SecureMesh mTLS -- TrustShield Pillar 4.

Covers TLSManager (CA init, cert issuance, listing, rotation, revocation)
and MTLSClient (creation, revocation enforcement, disabled fallback).
"""

from __future__ import annotations

import json
import ssl
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING

import httpx
import pytest
from cryptography import x509
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec

from llmhosts.security.mtls import MTLSClient, MTLSConfig, create_mtls_client
from llmhosts.security.tls_manager import CertInfo, TLSManager

if TYPE_CHECKING:
    from pathlib import Path


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
def tls_dir(tmp_path: Path) -> Path:
    """Return an isolated TLS directory for a single test."""
    d = tmp_path / "tls"
    d.mkdir()
    return d


@pytest.fixture()
def manager(tls_dir: Path) -> TLSManager:
    """Return a TLSManager rooted in the test-specific TLS directory."""
    return TLSManager(tls_dir=tls_dir)


@pytest.fixture()
def initialized_manager(manager: TLSManager) -> TLSManager:
    """Return a TLSManager with the CA already initialized."""
    manager.init_ca()
    return manager


# ---------------------------------------------------------------------------
# CA initialization
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestCAInit:
    """Tests for TLSManager.init_ca()."""

    def test_creates_ca_key_and_cert(self, manager: TLSManager, tls_dir: Path) -> None:
        """init_ca() should create ca.key and ca.crt."""
        manager.init_ca()

        assert (tls_dir / "ca.key").exists()
        assert (tls_dir / "ca.crt").exists()

    def test_creates_fernet_key_file(self, manager: TLSManager, tls_dir: Path) -> None:
        """init_ca() should create ca.fernet.key for CA key encryption."""
        manager.init_ca()
        assert (tls_dir / "ca.fernet.key").exists()

    def test_ca_cert_is_valid_x509(self, manager: TLSManager, tls_dir: Path) -> None:
        """The generated CA certificate should be valid X.509."""
        manager.init_ca()

        cert = x509.load_pem_x509_certificate((tls_dir / "ca.crt").read_bytes())
        assert cert.subject.get_attributes_for_oid(x509.oid.NameOID.COMMON_NAME)[0].value == "LLMHosts SecureMesh CA"

    def test_ca_cert_has_basic_constraints(self, manager: TLSManager, tls_dir: Path) -> None:
        """The CA cert must have BasicConstraints with ca=True."""
        manager.init_ca()

        cert = x509.load_pem_x509_certificate((tls_dir / "ca.crt").read_bytes())
        bc = cert.extensions.get_extension_for_class(x509.BasicConstraints).value
        assert bc.ca is True

    def test_ca_key_is_encrypted(self, manager: TLSManager, tls_dir: Path) -> None:
        """The stored ca.key should NOT be valid PEM (it is Fernet-encrypted)."""
        manager.init_ca()

        raw = (tls_dir / "ca.key").read_bytes()
        # Fernet-encrypted bytes will not start with '-----BEGIN'
        assert not raw.startswith(b"-----BEGIN")

    def test_reinit_ca_is_noop(self, manager: TLSManager, tls_dir: Path) -> None:
        """Calling init_ca() twice should not overwrite existing CA."""
        manager.init_ca()
        first_cert_bytes = (tls_dir / "ca.crt").read_bytes()

        manager.init_ca()
        second_cert_bytes = (tls_dir / "ca.crt").read_bytes()

        assert first_cert_bytes == second_cert_bytes

    def test_ca_cert_validity_period(self, manager: TLSManager, tls_dir: Path) -> None:
        """The CA certificate should have a ~10-year validity period."""
        manager.init_ca()

        cert = x509.load_pem_x509_certificate((tls_dir / "ca.crt").read_bytes())
        delta = cert.not_valid_after_utc - cert.not_valid_before_utc
        assert delta.days >= 3649  # ~10 years


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


@pytest.mark.unit
class TestCertIssuance:
    """Tests for TLSManager.issue_cert()."""

    def test_issue_proxy_cert(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """issue_cert('proxy') should create proxy.key and proxy.crt in the TLS root."""
        key_path, crt_path = initialized_manager.issue_cert("proxy")

        assert key_path == tls_dir / "proxy.key"
        assert crt_path == tls_dir / "proxy.crt"
        assert key_path.exists()
        assert crt_path.exists()

    def test_issue_backend_cert(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """issue_cert('ollama-local') should create files in backends/."""
        key_path, crt_path = initialized_manager.issue_cert("ollama-local")

        assert key_path == tls_dir / "backends" / "ollama-local.key"
        assert crt_path == tls_dir / "backends" / "ollama-local.crt"
        assert key_path.exists()
        assert crt_path.exists()

    def test_issued_cert_is_signed_by_ca(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """The issued certificate should be signed by the local CA."""
        initialized_manager.issue_cert("test-backend")

        ca_cert = x509.load_pem_x509_certificate((tls_dir / "ca.crt").read_bytes())
        backend_cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "test-backend.crt").read_bytes())

        # Issuer of backend cert should match CA subject
        assert backend_cert.issuer == ca_cert.subject

    def test_issued_cert_has_correct_cn(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """The issued cert CN should match the requested name."""
        initialized_manager.issue_cert("my-gpu-node")

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "my-gpu-node.crt").read_bytes())
        cn = cert.subject.get_attributes_for_oid(x509.oid.NameOID.COMMON_NAME)[0].value
        assert cn == "my-gpu-node"

    def test_issued_cert_has_san(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """The issued cert should include a SubjectAlternativeName extension."""
        initialized_manager.issue_cert("backend-x")

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "backend-x.crt").read_bytes())
        san = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName).value
        dns_names = san.get_values_for_type(x509.DNSName)
        assert "backend-x" in dns_names
        assert "localhost" in dns_names

    def test_issued_cert_not_ca(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """Issued certificates must NOT have ca=True in BasicConstraints."""
        initialized_manager.issue_cert("worker")

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "worker.crt").read_bytes())
        bc = cert.extensions.get_extension_for_class(x509.BasicConstraints).value
        assert bc.ca is False

    def test_issued_cert_365_day_validity(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """Issued certs should have a 365-day validity window."""
        initialized_manager.issue_cert("node-a")

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "node-a.crt").read_bytes())
        delta = cert.not_valid_after_utc - cert.not_valid_before_utc
        assert delta.days == 365

    def test_issue_cert_without_ca_fails(self, manager: TLSManager) -> None:
        """issue_cert() should raise FileNotFoundError when CA is not initialized."""
        with pytest.raises(FileNotFoundError, match="CA not initialized"):
            manager.issue_cert("bad-backend")

    def test_issued_key_is_valid_ec(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """The issued private key should be a valid ECDSA P-256 key."""
        initialized_manager.issue_cert("gpu-1")

        pem = (tls_dir / "backends" / "gpu-1.key").read_bytes()
        key = serialization.load_pem_private_key(pem, password=None)
        assert isinstance(key, ec.EllipticCurvePrivateKey)
        assert key.curve.name == "secp256r1"


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


@pytest.mark.unit
class TestCertListing:
    """Tests for TLSManager.list_certs()."""

    def test_list_certs_empty(self, manager: TLSManager) -> None:
        """list_certs() should return empty when no certs exist."""
        assert manager.list_certs() == []

    def test_list_certs_after_init(self, initialized_manager: TLSManager) -> None:
        """After init_ca(), list_certs() should include the CA cert."""
        certs = initialized_manager.list_certs()
        names = [c.name for c in certs]
        assert "ca" in names

    def test_list_certs_includes_backends(self, initialized_manager: TLSManager) -> None:
        """Issued backend certs should appear in list_certs()."""
        initialized_manager.issue_cert("proxy")
        initialized_manager.issue_cert("backend-1")
        initialized_manager.issue_cert("backend-2")

        certs = initialized_manager.list_certs()
        names = {c.name for c in certs}
        assert {"ca", "proxy", "backend-1", "backend-2"}.issubset(names)

    def test_list_certs_returns_certinfo(self, initialized_manager: TLSManager) -> None:
        """Each item should be a CertInfo with all expected fields."""
        initialized_manager.issue_cert("proxy")
        certs = initialized_manager.list_certs()

        proxy_info = next(c for c in certs if c.name == "proxy")
        assert isinstance(proxy_info, CertInfo)
        assert isinstance(proxy_info.serial, int)
        assert isinstance(proxy_info.not_before, datetime)
        assert isinstance(proxy_info.not_after, datetime)
        assert len(proxy_info.fingerprint) == 64  # SHA-256 hex
        assert proxy_info.revoked is False

    def test_list_certs_shows_revoked_flag(self, initialized_manager: TLSManager) -> None:
        """Revoked certs should have revoked=True in list_certs()."""
        initialized_manager.issue_cert("doomed")
        initialized_manager.revoke_cert("doomed")

        certs = initialized_manager.list_certs()
        doomed = next(c for c in certs if c.name == "doomed")
        assert doomed.revoked is True


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


@pytest.mark.unit
class TestCertRotation:
    """Tests for TLSManager.rotate_cert()."""

    def test_rotate_produces_new_cert(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """rotate_cert() should produce a certificate with a different serial number."""
        initialized_manager.issue_cert("rotating-svc")
        old_cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "rotating-svc.crt").read_bytes())

        initialized_manager.rotate_cert("rotating-svc")
        new_cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "rotating-svc.crt").read_bytes())

        assert old_cert.serial_number != new_cert.serial_number

    def test_rotate_keeps_same_cn(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """Rotated cert should keep the same Common Name."""
        initialized_manager.issue_cert("my-node")
        initialized_manager.rotate_cert("my-node")

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "my-node.crt").read_bytes())
        cn = cert.subject.get_attributes_for_oid(x509.oid.NameOID.COMMON_NAME)[0].value
        assert cn == "my-node"


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


@pytest.mark.unit
class TestCertRevocation:
    """Tests for TLSManager.revoke_cert() and is_revoked()."""

    def test_revoke_existing_cert(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """revoke_cert() should return True and record the serial."""
        initialized_manager.issue_cert("victim")
        assert initialized_manager.revoke_cert("victim") is True

        cert = x509.load_pem_x509_certificate((tls_dir / "backends" / "victim.crt").read_bytes())
        assert initialized_manager.is_revoked(cert.serial_number) is True

    def test_revoke_nonexistent_cert(self, initialized_manager: TLSManager) -> None:
        """revoke_cert() should return False for a cert that doesn't exist."""
        assert initialized_manager.revoke_cert("ghost") is False

    def test_revoked_json_format(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """The revoked.json file should contain a 'serials' list."""
        initialized_manager.issue_cert("revokable")
        initialized_manager.revoke_cert("revokable")

        data = json.loads((tls_dir / "revoked.json").read_text(encoding="utf-8"))
        assert "serials" in data
        assert isinstance(data["serials"], list)
        assert len(data["serials"]) == 1

    def test_is_revoked_false_for_valid(self, initialized_manager: TLSManager) -> None:
        """is_revoked() should return False for a non-revoked serial."""
        assert initialized_manager.is_revoked(12345) is False

    def test_multiple_revocations(self, initialized_manager: TLSManager) -> None:
        """Multiple certs can be revoked independently."""
        initialized_manager.issue_cert("a")
        initialized_manager.issue_cert("b")

        initialized_manager.revoke_cert("a")
        initialized_manager.revoke_cert("b")

        certs = initialized_manager.list_certs()
        revoked = [c for c in certs if c.revoked and c.name in ("a", "b")]
        assert len(revoked) == 2


# ---------------------------------------------------------------------------
# Cert expiry detection
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestCertExpiry:
    """Tests for detecting certificate expiry via list_certs() metadata."""

    def test_cert_not_expired(self, initialized_manager: TLSManager) -> None:
        """A freshly issued cert should not be expired."""
        initialized_manager.issue_cert("fresh")

        info = next(c for c in initialized_manager.list_certs() if c.name == "fresh")
        assert info.not_after > datetime.now(timezone.utc)

    def test_cert_expiry_date_is_365_days_out(self, initialized_manager: TLSManager) -> None:
        """not_after should be ~365 days from now."""
        initialized_manager.issue_cert("timed")

        info = next(c for c in initialized_manager.list_certs() if c.name == "timed")
        expected = datetime.now(timezone.utc) + timedelta(days=365)
        # Allow 60-second tolerance for test runtime
        assert abs((info.not_after - expected).total_seconds()) < 60


# ---------------------------------------------------------------------------
# mTLS client -- creation and factory
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestMTLSClientCreation:
    """Tests for MTLSClient construction and the create_mtls_client factory."""

    def test_create_mtls_client_disabled(self) -> None:
        """When disabled, create_mtls_client returns a plain httpx.AsyncClient."""
        client = create_mtls_client(MTLSConfig(enabled=False))
        assert isinstance(client, httpx.AsyncClient)

    def test_create_mtls_client_enabled(self, initialized_manager: TLSManager) -> None:
        """When enabled, create_mtls_client returns an MTLSClient."""
        initialized_manager.issue_cert("proxy")
        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="proxy")
        client = create_mtls_client(config)
        assert isinstance(client, MTLSClient)

    def test_mtls_client_requires_tls_manager(self) -> None:
        """MTLSClient should raise if tls_manager is None."""
        config = MTLSConfig(enabled=True, tls_manager=None)
        with pytest.raises(ValueError, match="tls_manager must be set"):
            MTLSClient(config)

    def test_mtls_client_build_ssl_context(self, initialized_manager: TLSManager) -> None:
        """The internal SSL context should be configured for client auth."""
        initialized_manager.issue_cert("proxy")
        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="proxy")
        client = MTLSClient(config)

        ctx = client._build_ssl_context()
        assert isinstance(ctx, ssl.SSLContext)

    def test_mtls_client_missing_cert_raises(self, initialized_manager: TLSManager) -> None:
        """Building SSL context without issued cert should raise FileNotFoundError."""
        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="nonexistent")
        client = MTLSClient(config)

        with pytest.raises(FileNotFoundError, match="Client certificate"):
            client._build_ssl_context()

    def test_mtls_client_missing_ca_raises(self, manager: TLSManager, tls_dir: Path) -> None:
        """Building SSL context without CA cert should raise FileNotFoundError."""
        # Create a cert dir but no CA
        (tls_dir / "proxy.key").touch()
        (tls_dir / "proxy.crt").touch()

        config = MTLSConfig(enabled=True, tls_manager=manager, client_name="proxy")
        client = MTLSClient(config)

        with pytest.raises(FileNotFoundError, match="CA certificate not found"):
            client._build_ssl_context()


# ---------------------------------------------------------------------------
# mTLS client -- revocation enforcement
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestMTLSRevocation:
    """Tests for mTLS revocation checking."""

    def test_check_revocation_passes_for_valid(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """_check_revocation should not raise for a valid (non-revoked) cert."""
        initialized_manager.issue_cert("proxy")
        initialized_manager.issue_cert("backend-ok")

        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="proxy")
        client = MTLSClient(config)

        cert_pem = (tls_dir / "backends" / "backend-ok.crt").read_bytes()
        # Should not raise
        client._check_revocation(cert_pem)

    def test_check_revocation_rejects_revoked(self, initialized_manager: TLSManager, tls_dir: Path) -> None:
        """_check_revocation should raise ConnectionRefusedError for a revoked cert."""
        initialized_manager.issue_cert("proxy")
        initialized_manager.issue_cert("revoked-backend")
        initialized_manager.revoke_cert("revoked-backend")

        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="proxy")
        client = MTLSClient(config)

        cert_pem = (tls_dir / "backends" / "revoked-backend.crt").read_bytes()
        with pytest.raises(ConnectionRefusedError, match="revoked"):
            client._check_revocation(cert_pem)


# ---------------------------------------------------------------------------
# mTLS client -- async lifecycle
# ---------------------------------------------------------------------------


@pytest.mark.unit
@pytest.mark.asyncio
class TestMTLSAsyncLifecycle:
    """Async tests for MTLSClient lifecycle."""

    async def test_close_without_open(self, initialized_manager: TLSManager) -> None:
        """Calling close() before any request should not raise."""
        initialized_manager.issue_cert("proxy")
        config = MTLSConfig(enabled=True, tls_manager=initialized_manager, client_name="proxy")
        client = MTLSClient(config)
        await client.close()  # Should be safe

    async def test_disabled_client_close(self) -> None:
        """A disabled (plain httpx) client should support aclose()."""
        client = create_mtls_client(MTLSConfig(enabled=False))
        assert isinstance(client, httpx.AsyncClient)
        await client.aclose()


# ---------------------------------------------------------------------------
# TLSManager -- custom Fernet key
# ---------------------------------------------------------------------------


@pytest.mark.unit
class TestCustomFernetKey:
    """Tests for TLSManager with a user-provided Fernet key."""

    def test_custom_key_encrypts_ca(self, tls_dir: Path) -> None:
        """When a custom Fernet key is provided, it should be used for CA encryption."""
        from cryptography.fernet import Fernet

        custom_key = Fernet.generate_key()
        mgr = TLSManager(tls_dir=tls_dir, fernet_key=custom_key)
        mgr.init_ca()

        # The encrypted ca.key should be decryptable with the custom key
        encrypted = (tls_dir / "ca.key").read_bytes()
        decrypted = Fernet(custom_key).decrypt(encrypted)
        assert decrypted.startswith(b"-----BEGIN PRIVATE KEY-----")

    def test_different_key_fails_decrypt(self, tls_dir: Path) -> None:
        """Using the wrong Fernet key should fail to load the CA."""
        from cryptography.fernet import Fernet, InvalidToken

        original_key = Fernet.generate_key()
        mgr = TLSManager(tls_dir=tls_dir, fernet_key=original_key)
        mgr.init_ca()

        wrong_key = Fernet.generate_key()
        mgr2 = TLSManager(tls_dir=tls_dir, fernet_key=wrong_key)
        with pytest.raises(InvalidToken):
            mgr2._load_ca()
