from __future__ import annotations
import binascii
import requests
import warnings
from cryptography import x509
from typing import Optional, Dict, Any, Tuple
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, rsa, padding
from cryptography.x509.oid import ObjectIdentifier, SignatureAlgorithmOID
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature

from importlib import import_module
tee_amd_sev_snp_rust_py_wrap = import_module(__package__ + ".tee_amd_sev_snp_rust_py_wrap")

from .tee_attestation_dataclasses import ParsedAttestationReport, ReportedTcb, CpuId
from .AttestationExceptions import (
    AttestationError,
    ReportDataError,
    KdsFetchError,
    CertChainError,
    MetadataMismatchError,
    SignatureError,
    UnsupportedChipFamilyError,
    FirmwareOpenErrorAmdSevSnp,
    AttestationParseErrorAmdSevSnp, 
)


warnings.filterwarnings(
    "ignore",
    message="Parsed a serial number which wasn't positive",
)


KDS_CERT_SITE = "https://kdsintf.amd.com"
KDS_VCEK_PATH = "/vcek/v1"
KDS_CERT_CHAIN_PATH = "cert_chain"


# AMD SEV-SNP X.509 extension OIDs
OID_BOOTLOADER = ObjectIdentifier("1.3.6.1.4.1.3704.1.3.1")
OID_TEE        = ObjectIdentifier("1.3.6.1.4.1.3704.1.3.2")
OID_SNP        = ObjectIdentifier("1.3.6.1.4.1.3704.1.3.3")
OID_UCODE      = ObjectIdentifier("1.3.6.1.4.1.3704.1.3.8")
OID_HWID       = ObjectIdentifier("1.3.6.1.4.1.3704.1.4")


class AmdSevSnpVerifier:
    """
    Verify AMD SEV-SNP attestation reports using AMD KDS certificates.

    The verifier will:
    - Parse SEV-SNP report bytes into a structured ParsedAttestationReport
    - Recompute and validate report_data binding to (pubkey_der, client_id, proof)
    - Fetch ARK/ASK/VCEK from AMD KDS and validate the certificate chain
    - Validate VCEK metadata (TCB + HWID) matches the attestation report
    - Verify the report signature using the VCEK public key
    """


    @staticmethod
    def _derive_chip_family(cpu: CpuId) -> str:
        """
        Derive the AMD product family name from CPUID family and model.

        The mapping will:
        - Convert CPUID family/model to a supported product string (e.g., Milan/Genoa/Turin)
        - Raise if the family/model is not supported

        Args:
        - cpu: CpuID object containing identification values extracted from the attestation report.

        Returns:
        - AMD product family name used by KDS endpoints

        Raises:
        - UnsupportedChipFamilyError if the CPUID family/model is unknown
        """
        fam = cpu.fam_id
        mod = cpu.mod_id

        # Family 19h (Zen 3 / Zen 4)
        if fam == 0x19:
            if mod < 0x10:
                return "Milan"
            else:
                return "Genoa"

        # Family 1Ah (Zen 5)
        if fam == 0x1A:
            return "Turin"

        raise UnsupportedChipFamilyError(
            f"Unsupported CPUID family/model: fam=0x{fam:x}, mod=0x{mod:x}, step=0x{cpu.step:x}"
        )


    @staticmethod
    def _verify_cert_signed_by(child, issuer) -> None:
        """
        Verify an X.509 certificate signature using the issuer public key.

        The verifier will:
        - Validate child.signature over child.tbs_certificate_bytes using issuer.public_key()
        - Support RSA (PKCS#1 v1.5 and PSS) and ECDSA issuer keys
        - Raise if the signature is invalid or the key type is unsupported

        Args:
        - child: X.509 certificate whose signature must be verified.
        - issuer: X.509 certificate providing the public key used for verification.

        Raises:
        - CertChainError if verification fails
        """
        pub = issuer.public_key()
        try:
            if isinstance(pub, rsa.RSAPublicKey):
                if child.signature_algorithm_oid == SignatureAlgorithmOID.RSASSA_PSS:
                    params = child.signature_algorithm_parameters

                    if params is None:
                        mgf = padding.MGF1(child.signature_hash_algorithm)
                        salt_len = child.signature_hash_algorithm.digest_size
                    else:
                        mgf = getattr(params, "mgf", padding.MGF1(child.signature_hash_algorithm))
                        salt_len = getattr(params, "salt_length", child.signature_hash_algorithm.digest_size)

                    pub.verify(
                        child.signature,
                        child.tbs_certificate_bytes,
                        padding.PSS(mgf=mgf, salt_length=salt_len),
                        child.signature_hash_algorithm,
                    )
                else:
                    pub.verify(
                        child.signature,
                        child.tbs_certificate_bytes,
                        padding.PKCS1v15(),
                        child.signature_hash_algorithm,
                    )
            elif isinstance(pub, ec.EllipticCurvePublicKey):
                pub.verify(
                    child.signature,
                    child.tbs_certificate_bytes,
                    ec.ECDSA(child.signature_hash_algorithm),
                )
            else:
                raise CertChainError(f"Unsupported issuer key type: {type(pub).__name__}")

        except Exception as e:
            raise CertChainError(f"Certificate signature verification failed: {e!r}")


    @staticmethod
    def _get_unrecognized_ext_value(cert: x509.Certificate, oid: ObjectIdentifier) -> Optional[bytes]:
        """
        Extract the raw value of a vendor X.509 extension.

        The helper will:
        - Return the raw bytes for UnrecognizedExtension values
        - Return None if the extension is not present

        Args:
        - cert: X.509 certificate containing the extension.
        - oid: Object identifier of the extension to retrieve.

        Returns:
        - Raw extension bytes if present
        - None if not present
        """
        try:
            ext = cert.extensions.get_extension_for_oid(oid).value
        except x509.ExtensionNotFound:
            return None
        # SNP uses vendor extensions; in cryptography these often appear as UnrecognizedExtension
        if isinstance(ext, x509.UnrecognizedExtension):
            return ext.value
        # If it’s some recognized type, serialize value conservatively
        return getattr(ext, "value", None)


    @staticmethod
    def _chip_id_hex(chip_id_64: bytes) -> str:
        """
        Convert a 64-byte chip_id into lowercase hex.

        Args:
        - chip_id_64: Raw chip identifier bytes (must be exactly 64 bytes).

        Returns:
        - Hex string encoding of chip_id

        Raises:
        - AttestationError if chip_id is not 64 bytes
        """
        if len(chip_id_64) != 64:
            raise AttestationError(f"Chip_id must be 64 bytes, got {len(chip_id_64)}")
        return binascii.hexlify(chip_id_64).decode("ascii")


    @staticmethod
    def _check_cert_ext_byte(ext_value: bytes, expected: int) -> None:
        """
        Validate an INTEGER-like X.509 extension encoding a single expected byte.

        The checker will:
        - Enforce basic ASN.1-like framing (0x02, length, value)
        - Compare the encoded value with the expected byte

        Args:
        - ext_value: Raw extension value bytes extracted from the certificate.
        - expected: Expected single-byte value to compare against.

        Raises:
        - MetadataMismatchError if the extension value does not match
        """
        if not ext_value:
            raise MetadataMismatchError("Empty extension value")
        if ext_value[0] != 0x02:
            raise MetadataMismatchError("Unexpected ASN.1 type in extension value (expected INTEGER-like encoding)")
        if len(ext_value) < 3:
            raise MetadataMismatchError("Extension value too short")
        if ext_value[1] not in (0x01, 0x02):
            raise MetadataMismatchError("Unexpected length in extension value")
        if ext_value[-1] != expected:
            raise MetadataMismatchError(f"Extension byte mismatch: got {ext_value[-1]}, expected {expected}")


    @staticmethod
    def _check_cert_ext_bytes(ext_value: bytes, expected: bytes) -> None:
        """
        Validate an X.509 extension value matches expected bytes exactly.

        The checker will:
        - Compare ext_value against expected without transformation

        Args:
        - ext_value: Raw extension value bytes extracted from the certificate.
        - expected: Expected bytes sequence to match exactly.

        Raises:
        - MetadataMismatchError if the bytes do not match
        """
        if ext_value != expected:
            raise MetadataMismatchError("Extension bytes do not match expected value")


    @staticmethod
    def _validate_cert_metadata(vcek_cert: x509.Certificate, report: ParsedAttestationReport) -> None:
        """
        Validate VCEK metadata matches the attestation report metadata.

        The validator will:
        - Compare VCEK TCB extension values against report.reported_tcb
        - Compare VCEK HWID extension bytes against report.chip_id

        Args:
        - vcek_cert: VCEK X.509 certificate fetched from AMD KDS.
        - report: ParsedAttestationReport object containing TCB and chip metadata.

        Raises:
        - MetadataMismatchError if any metadata check fails
        """
        bl = AmdSevSnpVerifier._get_unrecognized_ext_value(vcek_cert, OID_BOOTLOADER)
        tee = AmdSevSnpVerifier._get_unrecognized_ext_value(vcek_cert, OID_TEE)
        snp = AmdSevSnpVerifier._get_unrecognized_ext_value(vcek_cert, OID_SNP)
        ucode = AmdSevSnpVerifier._get_unrecognized_ext_value(vcek_cert, OID_UCODE)
        hwid = AmdSevSnpVerifier._get_unrecognized_ext_value(vcek_cert, OID_HWID)

        tcb = report.reported_tcb

        if bl is not None:
            AmdSevSnpVerifier._check_cert_ext_byte(bl, tcb.bootloader)
        if tee is not None:
            AmdSevSnpVerifier._check_cert_ext_byte(tee, tcb.tee)
        if snp is not None:
            AmdSevSnpVerifier._check_cert_ext_byte(snp, tcb.snp)
        if ucode is not None:
            AmdSevSnpVerifier._check_cert_ext_byte(ucode, tcb.microcode)
        if hwid is not None:
            AmdSevSnpVerifier._check_cert_ext_bytes(hwid, report.chip_id)


    @staticmethod
    def _request_cert_chain(sev_prod_name: str) -> Tuple[x509.Certificate, x509.Certificate]:
        """
        Fetch the ARK and ASK certificates for a given AMD product family.

        The request will:
        - Download the PEM chain from AMD KDS
        - Parse certificates and select the self-signed ARK and the ASK signed by ARK

        Args:
        - sev_prod_name: AMD product family name (e.g., Milan, Genoa, Turin)

        Returns:
        - (ark, ask) as X.509 certificates

        Raises:
        - KdsFetchError if the HTTP request fails
        - CertChainError if parsing or selection fails
        """
        url = f"{KDS_CERT_SITE}{KDS_VCEK_PATH}/{sev_prod_name}/{KDS_CERT_CHAIN_PATH}"
        try:
            rsp = requests.get(url, timeout=20)
            rsp.raise_for_status()
        except Exception as e:
            raise KdsFetchError(f"Failed to fetch cert chain: {e!r}")

        # PEM containing ARK + ASK
        pem = rsp.content
        certs = list(x509.load_pem_x509_certificates(pem))
        if len(certs) < 2:
            raise CertChainError("Expected at least ARK + ASK in chain PEM")

        # ARK is the self-signed root (subject == issuer)
        roots = [c for c in certs if c.subject == c.issuer]
        if not roots:
            raise CertChainError("Could not find self-signed ARK in returned chain")
        ark = roots[0]

        # ASK is the cert whose issuer is ARK's subject (signed by ARK)
        asks = [c for c in certs if c.issuer == ark.subject and c.subject != ark.subject]
        if not asks:
            raise CertChainError("Could not find ASK signed by ARK in returned chain")
        ask = asks[0]

        return ark, ask


    @staticmethod
    def _request_vcek(
        sev_prod_name: str,
        chip_id_64: bytes,
        bl_spl: int,
        tee_spl: int,
        snp_spl: int,
        ucode_spl: int,
    ) -> x509.Certificate:
        """
        Fetch the VCEK certificate for a chip and TCB from AMD KDS.

        The request will:
        - Build the KDS VCEK URL using hwid (chip_id) and SPL values (bootloader/tee/snp/ucode)
        - Download the DER certificate and parse it as X.509

        Args:
        - sev_prod_name: AMD product family name used by the KDS endpoint (e.g., Milan, Genoa, Turin).
        - chip_id_64: Raw chip identifier bytes (64 bytes).
        - bl_spl: Bootloader security patch level (SPL).
        - tee_spl: TEE firmware security patch level (SPL).
        - snp_spl: SNP firmware security patch level (SPL).
        - ucode_spl: CPU microcode security patch level (SPL).

        Returns:
        - VCEK X.509 certificate

        Raises:
        - KdsFetchError if the HTTP request fails
        - CertChainError if the certificate cannot be parsed
        """
        hwid = AmdSevSnpVerifier._chip_id_hex(chip_id_64)
        url = (
            f"{KDS_CERT_SITE}{KDS_VCEK_PATH}/{sev_prod_name}/{hwid}"
            f"?blSPL={bl_spl:02}&teeSPL={tee_spl:02}&snpSPL={snp_spl:02}&ucodeSPL={ucode_spl:02}"
        )

        try:
            rsp = requests.get(url, timeout=20)
            rsp.raise_for_status()
        except Exception as e:
            raise KdsFetchError(f"Failed to fetch VCEK: {e!r}")

        # DER containing VCEK
        der = rsp.content
        try:
            return x509.load_der_x509_certificate(der)
        except Exception as e:
            raise CertChainError(f"Failed to parse VCEK DER: {e!r}")


    @staticmethod
    def _verify_report_with_vcek(report: ParsedAttestationReport, vcek: x509.Certificate) -> None:
        """
        Verify the SEV-SNP report signature using the VCEK public key.

        The verifier will:
        - Convert little-endian r/s components into a DER-encoded ECDSA signature
        - Verify ECDSA P-384 over SHA-384 using the VCEK public key

        Args:
        - report: ParsedAttestationReport object containing signed bytes and signature.
        - vcek: VCEK X.509 certificate providing the verification public key.

        Raises:
        - SignatureError if signature validation fails
        """
        # Check ECDSA signature parameters lenght
        if len(report.sig_r) != 72:
            raise SignatureError(f"Error: sig_r component must be 72 bytes, got {len(report.sig_r)}")
        if len(report.sig_s) != 72:
            raise SignatureError(f"Error: sig_s component must be 72 bytes, got {len(report.sig_s)}")

        # Build DER ECDSA signature from r and s values extracted from the report
        r_int = int.from_bytes(report.sig_r, "little")
        s_int = int.from_bytes(report.sig_s, "little")

        # Sanity: P-384 <= 384 bits
        if r_int.bit_length() > 384 or s_int.bit_length() > 384:
            raise SignatureError(f"Unexpected r/s size: r={r_int.bit_length()} bits, s={s_int.bit_length()} bits")

        # Build signature
        report_signature_der = encode_dss_signature(r_int, s_int)

        # Verify ECDSA P-384 over SHA-384
        try:
            vcek.public_key().verify(
                report_signature_der,
                report.signed_bytes,
                ec.ECDSA(hashes.SHA384()),
            )
        except Exception as e:
            raise SignatureError(f"VCEK did NOT sign the attestation report: {e!r}")


    @staticmethod
    def _verify_amd_chain_and_tcb(report: ParsedAttestationReport, chip_family: str) -> None:
        """
        Verify AMD certificate chain, metadata binding, and report signature.

        The verifier will:
        - Fetch ARK + ASK for the chip family
        - Fetch VCEK for (chip_id, reported_tcb)
        - Verify certificate chain signatures (ARK self-signed, ASK by ARK, VCEK by ASK)
        - Validate VCEK metadata matches report (TCB + HWID)
        - Verify the report signature using the VCEK public key

        Args:
        - report: ParsedAttestationReport object to be verified.
        - chip_family: AMD product family name derived from CPUID (e.g., Milan, Genoa, Turin).

        Raises:
        - KdsFetchError, CertChainError, MetadataMismatchError, SignatureError on failure
        """
        # Fetch ARK + ASK
        ark, ask = AmdSevSnpVerifier._request_cert_chain(chip_family)

        # Fetch VCEK for this chip + TCB
        tcb = report.reported_tcb
        vcek = AmdSevSnpVerifier._request_vcek(
            chip_family,
            report.chip_id,
            tcb.bootloader,
            tcb.tee,
            tcb.snp,
            tcb.microcode,
        )

        # Chain checks: Verify ARK self-signed, ASK signed by ARK, VCEK signed by ASK
        AmdSevSnpVerifier._verify_cert_signed_by(ark, ark) 
        AmdSevSnpVerifier._verify_cert_signed_by(ask, ark)
        AmdSevSnpVerifier._verify_cert_signed_by(vcek, ask)

        # Verify VCEK's metadata match the attestation report's metadata
        AmdSevSnpVerifier._validate_cert_metadata(vcek, report)

        # Verify the report's ECDSA signature using the VCEK public key (ECDSA P-384 over SHA-384)
        AmdSevSnpVerifier._verify_report_with_vcek(report, vcek)


    @staticmethod
    def verify_amd_sev_snp_attestation_report(att_report: bytes, expected_report_data: bytes) -> None:
        """
        Verify an attestation report produced from an AMD SEV-SNP platform.

        The verifier will:
        - Parse att_report into a structured ParsedAttestationReport
        - Enforce exact 64-byte size for report_data and expected_report_data
        - Compare parsed report_data to the expected_report_data
        - Derive chip family from CPUID and run AMD chain + metadata + signature checks

        Args:
        - att_report: the attestation report bytes.
        - expected_report_data: the bytes expected to be found inside the report_data field of the report

        Raises:
        - AttestationParseErrorAmdSevSnp if the report cannot be properly parsed
        - ReportDataError if report_data does not match expectations
        - UnsupportedChipFamilyError if CPUID is not supported
        - KdsFetchError, CertChainError, MetadataMismatchError, SignatureError on verification failure
        """
        # Parse the TeeAttestation.report bytes using the rust library and returns a dictionary
        try:
            parsed_report = tee_amd_sev_snp_rust_py_wrap.parse_report(att_report)

        except ValueError as e:
            raise AttestationParseErrorAmdSevSnp("Failed to parse AMD SEV-SNP attestation report.") from e

        # Convert parsed report dictionary to a dataclass object
        report = ParsedAttestationReport(
            chip_id=parsed_report["chip_id"],
            report_data=parsed_report["report_data"],
            reported_tcb=ReportedTcb(**parsed_report["reported_tcb"]),
            cpuid=CpuId(**parsed_report["cpuid"]),
            signed_bytes=parsed_report["signed_bytes"],
            sig_r=parsed_report["sig_r"],
            sig_s=parsed_report["sig_s"],
        )
        
        # Enforce report_data size exactly like SNP expects
        if len(report.report_data) != 64:
            raise ReportDataError(f"Report data must be exactly 64 bytes, got {len(report.report_data)}")

        # Enforce expected_report_data size exactly like SNP expects
        if len(expected_report_data) != 64:
            raise ReportDataError(f"Expected report data must be exactly 64 bytes, got {len(expected_report_data)}")
        
        # Compare report.report_data against expected_report_data
        if report.report_data != expected_report_data:
            raise ReportDataError("Report_data does not match expected report data")

        # Derive the chip family commercial name from the cpuid informations in the report
        chip_family = AmdSevSnpVerifier._derive_chip_family(report.cpuid)

        # Run AMD chain + metadata + report signature verification
        AmdSevSnpVerifier._verify_amd_chain_and_tcb(report, chip_family)


    @staticmethod
    def get_amd_sev_snp_attestation_report(report_data: bytes, msg_ver: int, vmpl: int) -> bytes:
        """
        Fetch an AMD SEV-SNP attestation report from the local platform.

        The report call will:
        - Invoke the SEV-SNP firmware interface via the platform Rust wrapper
        - Request a report using report_data as the 64-byte nonce input
        - Return the raw report bytes as produced by the platform Rust library

        Args:
        - report_data: Nonce/report_data buffer (must be 64 bytes for SNP).
        - msg_ver: Message/report version selector passed to the platform wrapper.
        - vmpl: VM privilege level selector passed to the platform wrapper.

        Returns:
        - Raw SEV-SNP attestation report bytes

        Raises:
        - FirmwareOpenError if the SEV-SNP firmware interface cannot be opened.
        """
        try:
            return tee_amd_sev_snp_rust_py_wrap.get_report(report_data, msg_ver, vmpl)

        except ValueError as e:
            msg = str(e)

            # Firmware::open failure propagated from Rust
            if "Firmware::open failed" in msg:
                raise FirmwareOpenErrorAmdSevSnp("Unable to open SEV-SNP firmware interface") from e

            # Any other ValueError from Rust stays meaningful
            raise