#!/usr/bin/python3
# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.

import glob
import os
import struct
import re # For extracting firmware prefix
import sys

# Define the C struct format for Python's struct module
# struct common_firmware_header {
#   uint32_t size_bytes;                // Offset 0, 4 bytes
#   uint32_t header_size_bytes;         // Offset 4, 4 bytes
#   uint16_t header_version_major;      // Offset 8, 2 bytes
#   uint16_t header_version_minor;      // Offset 10, 2 bytes
#   uint16_t ip_version_major;          // Offset 12, 2 bytes
#   uint16_t ip_version_minor;          // Offset 14, 2 bytes
#   uint32_t ucode_version;             // Offset 16, 4 bytes
#   uint32_t ucode_size_bytes;          // Offset 20, 4 bytes
#   uint32_t ucode_array_offset_bytes;  // Offset 24, 4 bytes
#   uint32_t crc32;                     // Offset 28, 4 bytes
# }; Total size = 32 bytes

# '<' indicates little-endian
# 'I' for uint32_t (4 bytes)
# 'H' for uint16_t (2 bytes)
FIRMWARE_HEADER_FORMAT = "<IIHHHHIIII"
FIRMWARE_HEADER_SIZE = struct.calcsize(FIRMWARE_HEADER_FORMAT) # Should be 32

FIRMWARE_HEADER_FIELD_NAMES = [
    "Total Size (bytes)",
    "Header Size (bytes)",
    "Header Version Major",
    "Header Version Minor",
    "IP Version Major",
    "IP Version Minor",
    "Microcode Version (raw)", # This will be further parsed
    "Microcode Size (bytes)",
    "Microcode Array Offset (bytes)",
    "CRC32 Checksum"
]

# Defines for parsing ucode_version
VCN_FW_VERSION_VEP_IF_MASK        = 0XF0000000
VCN_FW_VERSION_VEP_IF_SHIFT       = 28
VCN_FW_VERSION_DEC_IF_MASK        = 0x0F000000
VCN_FW_VERSION_DEC_IF_SHIFT       = 24
VCN_FW_VERSION_ENC_IF_MAJOR_MASK  = 0x00F00000
VCN_FW_VERSION_ENC_IF_MAJOR_SHIFT = 20
VCN_FW_VERSION_ENC_IF_MINOR_MASK  = 0x000FF000
VCN_FW_VERSION_ENC_IF_MINOR_SHIFT = 12
VCN_FW_VERSION_REVISION_MASK      = 0x00000FFF
VCN_FW_VERSION_REVISION_SHIFT     = 0

def parse_ucode_version(ucode_version_raw):
    """Parses the raw ucode_version uint32_t into its components."""
    vep_if = (ucode_version_raw & VCN_FW_VERSION_VEP_IF_MASK) >> VCN_FW_VERSION_VEP_IF_SHIFT
    dec_if = (ucode_version_raw & VCN_FW_VERSION_DEC_IF_MASK) >> VCN_FW_VERSION_DEC_IF_SHIFT
    enc_if_major = (ucode_version_raw & VCN_FW_VERSION_ENC_IF_MAJOR_MASK) >> VCN_FW_VERSION_ENC_IF_MAJOR_SHIFT
    enc_if_minor = (ucode_version_raw & VCN_FW_VERSION_ENC_IF_MINOR_MASK) >> VCN_FW_VERSION_ENC_IF_MINOR_SHIFT
    revision = (ucode_version_raw & VCN_FW_VERSION_REVISION_MASK) >> VCN_FW_VERSION_REVISION_SHIFT
    
    return {
        "VEP Interface Version": vep_if,
        "Decoder Interface Version": dec_if,
        "Encoder Interface Major Version": enc_if_major,
        "Encoder Interface Minor Version": enc_if_minor,
        "Revision": revision,
    }

def parse_and_print_firmware_header(file_path, expected_enc_minor_version, file_prefix, error_log):
    """Reads the firmware file, parses the common_firmware_header, and checks encoder minor version."""
    try:
        with open(file_path, 'rb') as f:
            header_bytes = f.read(FIRMWARE_HEADER_SIZE)

        if len(header_bytes) < FIRMWARE_HEADER_SIZE:
            return

        unpacked_data = struct.unpack(FIRMWARE_HEADER_FORMAT, header_bytes)
        parsed_header_dict = dict(zip(FIRMWARE_HEADER_FIELD_NAMES, unpacked_data))

        # Parse microcode version to get encoder minor version
        ucode_version_raw = parsed_header_dict["Microcode Version (raw)"]
        parsed_ucode_components = parse_ucode_version(ucode_version_raw)
        actual_enc_minor = parsed_ucode_components["Encoder Interface Minor Version"]
        
        if expected_enc_minor_version is not None and actual_enc_minor != expected_enc_minor_version:
            error_msg = {
                'file': file_path,
                'prefix': file_prefix,
                'actual': actual_enc_minor,
                'expected': expected_enc_minor_version
            }
            error_log.append(error_msg)

    except Exception:
        # Silently skip files that can't be parsed
        pass

def get_firmware_prefix(filename):
    """Extracts a prefix like 'vcn_3' from a filename like 'vcn_3.0.1.bin' or 'vcn_3_0_1.bin'."""
    basename = os.path.basename(filename)
    # Match 'vcn_NUMBER' where NUMBER is one or more digits (e.g., vcn_3, vcn_40)
    match = re.match(r"(vcn_[0-9]+)", basename, re.IGNORECASE) # Ignore case for VCN prefix
    if match:
        return match.group(1).lower() # Return in lowercase for consistent map keys
    return None

def find_and_check_firmware(firmware_dir, file_pattern, firmware_version_map):
    """Finds firmware files and checks encoder minor versions, returning error details."""
    search_path = os.path.join(firmware_dir, file_pattern)
    error_log = []

    firmware_files = glob.glob(search_path)

    if not firmware_files:
        return error_log

    for fw_file in firmware_files:
        prefix = get_firmware_prefix(fw_file)
        expected_minor_version = None
        display_prefix_for_error = "unknown"

        if prefix:
            display_prefix_for_error = prefix
            if prefix in firmware_version_map:
                expected_minor_version = firmware_version_map[prefix]
        
        parse_and_print_firmware_header(fw_file, expected_minor_version, display_prefix_for_error, error_log)
    
    return error_log

def main():
    # Map of firmware filename prefix to expected Encoder Interface Minor Version
    # These values should be updated when building the package
    FIRMWARE_PREFIX_TO_ENC_MINOR_EXPECTED = {
        'vcn_3': 33,
        'vcn_4': 24,
        'vcn_5': 9,
        # Add more prefixes and expected versions as needed
    }

    # For trigger scripts, use fixed paths
    firmware_dir = "/lib/firmware/amdgpu/"
    file_pattern = "vcn_*.bin"

    # Check if directory exists
    if not os.path.exists(firmware_dir):
        sys.exit(0)

    # Check firmware versions
    errors = find_and_check_firmware(firmware_dir, file_pattern, FIRMWARE_PREFIX_TO_ENC_MINOR_EXPECTED)
    
    # Get list of firmware files found
    firmware_files = glob.glob(os.path.join(firmware_dir, file_pattern))
    
    if not firmware_files:
        sys.exit(0)
    
    if errors:
        print("VCN firmware version mismatch detected.")
        print("Please upgrade your libamdenc package to ensure compatibility.")
        sys.exit(0)
    else:
        print("All VCN firmware versions are compatible.")
        sys.exit(0)

if __name__ == "__main__":
    main()