"""
Byte-level validation tests for the .intunewin file format.

These tests serve as the format specification: if they pass, the byte layout is correct.
The only remaining validation is a real Intune upload (manual, Phase 5 gate).
"""
import base64
import hashlib
import hmac
import io
import os
import xml.etree.ElementTree as ET
import zipfile

import pytest
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad

from imptune.generators.intunewin_builder import build_intunewin


NAMESPACE = "http://schemas.microsoft.com/IntuneWin"


@pytest.fixture
def source_dir(tmp_path):
    """Create a small test source directory with a few files."""
    src = tmp_path / "source"
    src.mkdir()
    (src / "install.ps1").write_text("Write-Host 'Installing printer...'")
    (src / "config.json").write_text('{"printer": "HP LaserJet"}')
    (src / "readme.txt").write_text("Printer deployment package")
    return str(src)


@pytest.fixture
def built_package(source_dir, tmp_path):
    """Build an .intunewin package and return the path."""
    output = str(tmp_path / "package.intunewin")
    build_intunewin(source_dir, "install.ps1", output)
    return output


@pytest.fixture
def package_contents(built_package):
    """Extract outer ZIP contents and parsed Detection.xml."""
    with zipfile.ZipFile(built_package, "r") as outer:
        names = outer.namelist()
        blob = outer.read("IntuneWinPackage/Contents/IntunePackage.intunewin")
        detection_xml_bytes = outer.read("IntuneWinPackage/Metadata/Detection.xml")

    tree = ET.fromstring(detection_xml_bytes.decode("utf-8"))
    return {
        "names": names,
        "blob": blob,
        "detection_xml_bytes": detection_xml_bytes,
        "tree": tree,
    }


def _get_xml_text(tree, tag):
    """Get text of a direct child element (with namespace)."""
    elem = tree.find(f"{{{NAMESPACE}}}{tag}")
    if elem is None:
        # Try without namespace as fallback
        elem = tree.find(tag)
    return elem.text if elem is not None else None


def _get_encryption_info(tree):
    """Get EncryptionInfo sub-element."""
    enc = tree.find(f"{{{NAMESPACE}}}EncryptionInfo")
    if enc is None:
        enc = tree.find("EncryptionInfo")
    return enc


def _get_enc_text(tree, tag):
    """Get text of a child of EncryptionInfo."""
    enc = _get_encryption_info(tree)
    if enc is None:
        return None
    elem = enc.find(f"{{{NAMESPACE}}}{tag}")
    if elem is None:
        elem = enc.find(tag)
    return elem.text if elem is not None else None


class TestOuterZipStructure:
    def test_output_is_valid_zip(self, built_package):
        """build_intunewin() output file is a valid ZIP archive."""
        assert zipfile.is_zipfile(built_package), "Output file must be a valid ZIP archive"

    def test_outer_zip_structure(self, package_contents):
        """Outer ZIP contains exactly the two required entries."""
        names = set(package_contents["names"])
        assert "IntuneWinPackage/Contents/IntunePackage.intunewin" in names
        assert "IntuneWinPackage/Metadata/Detection.xml" in names

    def test_outer_zip_stored(self, built_package):
        """Outer ZIP entries use ZIP_STORED compression (no extra compression on encrypted content)."""
        with zipfile.ZipFile(built_package, "r") as outer:
            for info in outer.infolist():
                assert info.compress_type == zipfile.ZIP_STORED, (
                    f"Entry {info.filename} must use ZIP_STORED, got compress_type={info.compress_type}"
                )


class TestDetectionXml:
    def test_detection_xml_valid(self, package_contents):
        """Detection.xml is valid XML with ApplicationInfo root element in the correct namespace."""
        tree = package_contents["tree"]
        assert tree.tag == f"{{{NAMESPACE}}}ApplicationInfo", (
            f"Root element must be ApplicationInfo with namespace {NAMESPACE}, got {tree.tag}"
        )

    def test_detection_xml_fields(self, package_contents):
        """Detection.xml contains all required elements including all 8 EncryptionInfo sub-elements."""
        tree = package_contents["tree"]
        # Direct children
        for field in ("Name", "UnencryptedContentSize", "FileName", "SetupFile"):
            assert _get_xml_text(tree, field) is not None, f"Missing field: {field}"

        # EncryptionInfo sub-elements (all 8 required)
        for field in (
            "EncryptionKey",
            "MacKey",
            "InitializationVector",
            "Mac",
            "MacAlgorithm",
            "ProfileIdentifier",
            "FileDigest",
            "FileDigestAlgorithm",
        ):
            assert _get_enc_text(tree, field) is not None, f"Missing EncryptionInfo/{field}"

    def test_setup_file_in_detection_xml(self, package_contents):
        """SetupFile element matches the setup_file argument passed to build_intunewin."""
        tree = package_contents["tree"]
        assert _get_xml_text(tree, "SetupFile") == "install.ps1"


class TestEncryptedBlobLayout:
    def test_encrypted_blob_layout(self, package_contents):
        """Encrypted blob starts with 32 bytes (HMAC) + 16 bytes (IV) + remainder (ciphertext)."""
        blob = package_contents["blob"]
        # Must be at least 48 bytes (HMAC + IV) plus at least one AES block (16 bytes)
        assert len(blob) >= 64, f"Blob too short: {len(blob)} bytes"
        # Total length = 48 header + ciphertext length; ciphertext length is a multiple of 16
        ciphertext_len = len(blob) - 48
        assert ciphertext_len > 0, "Blob has no ciphertext after header"
        assert ciphertext_len % 16 == 0, (
            f"Ciphertext length {ciphertext_len} must be a multiple of AES block size 16"
        )

    def test_iv_is_16_bytes(self, package_contents):
        """IV extracted from Detection.xml base64-decodes to exactly 16 bytes (NOT 32 — critical per RESEARCH.md)."""
        tree = package_contents["tree"]
        iv_b64 = _get_enc_text(tree, "InitializationVector")
        assert iv_b64 is not None, "InitializationVector missing from Detection.xml"
        iv = base64.b64decode(iv_b64)
        assert len(iv) == 16, f"IV must be exactly 16 bytes, got {len(iv)}"

    def test_encryption_key_is_32_bytes(self, package_contents):
        """EncryptionKey from Detection.xml base64-decodes to exactly 32 bytes."""
        tree = package_contents["tree"]
        key_b64 = _get_enc_text(tree, "EncryptionKey")
        assert key_b64 is not None, "EncryptionKey missing from Detection.xml"
        key = base64.b64decode(key_b64)
        assert len(key) == 32, f"EncryptionKey must be exactly 32 bytes, got {len(key)}"

    def test_mac_key_is_32_bytes(self, package_contents):
        """MacKey from Detection.xml base64-decodes to exactly 32 bytes."""
        tree = package_contents["tree"]
        mac_key_b64 = _get_enc_text(tree, "MacKey")
        assert mac_key_b64 is not None, "MacKey missing from Detection.xml"
        mac_key = base64.b64decode(mac_key_b64)
        assert len(mac_key) == 32, f"MacKey must be exactly 32 bytes, got {len(mac_key)}"


class TestCryptographicVerification:
    def test_hmac_matches(self, package_contents):
        """HMAC-SHA256 of ciphertext matches first 32 bytes of blob AND Mac in Detection.xml."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        blob_hmac = blob[:32]
        ciphertext = blob[48:]

        mac_key_b64 = _get_enc_text(tree, "MacKey")
        mac_key = base64.b64decode(mac_key_b64)

        computed_hmac = hmac.new(mac_key, ciphertext, hashlib.sha256).digest()

        # Must match the blob header
        assert computed_hmac == blob_hmac, (
            "HMAC-SHA256 of ciphertext does not match the first 32 bytes of the blob"
        )

        # Must also match Detection.xml Mac field
        xml_mac = base64.b64decode(_get_enc_text(tree, "Mac"))
        assert computed_hmac == xml_mac, (
            "HMAC-SHA256 of ciphertext does not match the Mac value in Detection.xml"
        )

    def test_decryption_roundtrip(self, source_dir, package_contents):
        """Decrypt the ciphertext and verify it is a valid ZIP containing the original source files."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        # Must be a valid ZIP
        assert zipfile.is_zipfile(io.BytesIO(plaintext)), (
            "Decrypted plaintext is not a valid ZIP file"
        )

        # Must contain the original source files
        with zipfile.ZipFile(io.BytesIO(plaintext), "r") as inner_zip:
            inner_names = set(inner_zip.namelist())

        for filename in ("install.ps1", "config.json", "readme.txt"):
            assert filename in inner_names, (
                f"Original file {filename} not found in decrypted inner ZIP. Found: {inner_names}"
            )

    def test_file_digest_matches(self, package_contents):
        """FileDigest in Detection.xml matches SHA256 of the decrypted plaintext ZIP."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        computed_digest = hashlib.sha256(plaintext).digest()
        xml_digest = base64.b64decode(_get_enc_text(tree, "FileDigest"))

        assert computed_digest == xml_digest, (
            "FileDigest in Detection.xml does not match SHA256 of decrypted plaintext"
        )

    def test_unencrypted_content_size(self, package_contents):
        """UnencryptedContentSize in Detection.xml matches byte length of decrypted plaintext ZIP."""
        blob = package_contents["blob"]
        tree = package_contents["tree"]

        aes_key = base64.b64decode(_get_enc_text(tree, "EncryptionKey"))
        iv = base64.b64decode(_get_enc_text(tree, "InitializationVector"))
        ciphertext = blob[48:]

        cipher = AES.new(aes_key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)

        xml_size = int(_get_xml_text(tree, "UnencryptedContentSize"))
        assert xml_size == len(plaintext), (
            f"UnencryptedContentSize {xml_size} does not match actual plaintext size {len(plaintext)}"
        )