//! AACS Media Key Block — `[C]` Chapter 3.
//!
//! The MKB record format (framing walker, the `MkbRecord` view, record-body
//! finders), the MKBType / AACS-generation classification, and MKB-file
//! utilities (content length, trimming, version). Consolidated here so the one
//! place that understands MKB bytes is `mkb`. Some duplicate record finders
//! still live side by side pending a follow-up that collapses them.

// ── MKB record types ([C] Chapter 3) ──────────────────────────────────────
// The ONE canonical set; every record-type comparison in `aacs` references these
// so a type byte is never a bare literal scattered across files.

/// Type-and-Version — carries the 32-bit MKBType / AACS generation.
pub(crate) const REC_TYPE_AND_VERSION: u8 = 0x10;
/// Explicit Subset-Difference — the per-slot `(u_mask_shift, uv)` table.
pub(crate) const REC_SUBSET_DIFFERENCE: u8 = 0x04;
/// Media Key Data — the classical (1.0 / 2.0) per-subset cvalue table.
pub(crate) const REC_MEDIA_KEY_DATA: u8 = 0x05;
/// Subset-Difference Index — a search index into `0x04`, NOT a cvalue table.
#[cfg(test)]
pub(crate) const REC_SUBSET_DIFFERENCE_INDEX: u8 = 0x07;
/// Media Key Variant Data (AACS 2.1) — the per-subset-difference `C` table
/// (one 16-byte C per slot); the `Kmp` step reads C from HERE, not `0x2d`.
pub(crate) const REC_MEDIA_KEY_VARIANT_DATA: u8 = 0x0c;
/// Variant Data + Nonce (AACS 2.1) — the `VARIANTS[uv]` table (leading bytes)
/// with the 16-byte `Kvn` Nonce at the tail.
pub(crate) const REC_VARIANT_DATA_AND_NONCE: u8 = 0x2d;
/// Variant Key Data table (AACS 2.1) — 65,535×16, indexed by the resolved VKD index.
pub(crate) const REC_VKD_TABLE: u8 = 0x2f;
/// Verify-Media-Key — AACS 1.0.
pub(crate) const REC_VERIFY_MEDIA_KEY_V1: u8 = 0x81;
/// Verify-Media-Key — AACS 2.x.
pub(crate) const REC_VERIFY_MEDIA_KEY_V2: u8 = 0x86;

/// A single MKB record produced by [`walk_mkb`].
#[derive(Debug, Clone)]
pub struct MkbRecord {
    /// Byte offset of the record within the MKB.
    pub offset: usize,
    /// Record type byte.
    pub rec_type: u8,
    /// Record length in bytes (includes the 4-byte header).
    pub rec_len: usize,
    /// Record body (the bytes after the 4-byte header).
    pub body: Vec<u8>,
}

/// Walk an MKB into a flat list of records.
///
/// MKB record framing per AACS: 1 byte type, 3 bytes BE length
/// INCLUDING the 4-byte header, followed by payload. The walker stops
/// at the first `(type=0, len=0)` end marker or at end of buffer.
pub fn walk_mkb(mkb: &[u8]) -> Vec<MkbRecord> {
    mkb_records(mkb)
        .map(|(offset, rec_type, rec_len)| MkbRecord {
            offset,
            rec_type,
            rec_len,
            body: mkb[offset + 4..offset + rec_len].to_vec(),
        })
        .collect()
}

// THE single MKB record-framing walker: 4-byte header (type + BE24 length) then body, stopping
// at the `00 000000` end marker. `walk_mkb` and every other MKB walk are built on this.
pub(crate) fn mkb_records(mkb: &[u8]) -> impl Iterator<Item = (usize, u8, usize)> + '_ {
    let mut pos = 0usize;
    std::iter::from_fn(move || {
        if pos + 4 > mkb.len() {
            return None;
        }
        let rec_type = mkb[pos];
        let rec_len = ((mkb[pos + 1] as usize) << 16)
            | ((mkb[pos + 2] as usize) << 8)
            | (mkb[pos + 3] as usize);
        if rec_type == 0 && rec_len == 0 {
            return None;
        }
        if rec_len < 4 || pos + rec_len > mkb.len() {
            return None;
        }
        let here = pos;
        pos += rec_len;
        Some((here, rec_type, rec_len))
    })
}

pub(crate) fn mkb_find_body(records: &[MkbRecord], rec_type: u8) -> Option<&[u8]> {
    records
        .iter()
        .find(|r| r.rec_type == rec_type && !r.body.is_empty())
        .map(|r| r.body.as_slice())
}

/// AACS protection generation a disc carries.
///
/// The content cert type byte distinguishes V10 (`0x00`) from V20 (`0x10`; any other
/// byte is rejected, see [`super::inf::parse_content_cert`]). A V21 disc carries a V20
/// cert and is upgraded to `V21` only when the MKB has Variant records `0x2d`/`0x2f`.
///
/// Key-storage stride in `Unit_Key_RO.inf` is 48 bytes for V10 and 64
/// bytes for V20 / V21.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AacsVersion {
    /// AACS 1.0 — original BD-ROM.
    V10,
    /// AACS 2.0 — UHD-BD, classical Media Key derivation.
    V20,
    /// AACS 2.1 — UHD-BD with Media Key Variant chain on top of V20.
    V21,
}

/// AACS major version as the small integer threaded through the scan / key
/// paths (`AacsState.version`, `DiscInputs.version`, `DiscInputsCtx::new`):
/// 1 = AACS 1.0 (BD), 2 = AACS 2.x (UHD). Centralised so the bare `1`/`2` — and
/// the V10-vs-else stride choice it drives — lives in exactly one place.
pub const AACS_MAJOR_BD: u8 = 1;

/// AACS major version 2 (UHD-BD, 2.0 / 2.1): see [`AACS_MAJOR_BD`].
pub const AACS_MAJOR_UHD: u8 = 2;

impl AacsVersion {
    /// Stride (in bytes) between successive encrypted unit keys in
    /// `Unit_Key_RO.inf`.
    pub(crate) fn unit_key_stride(self) -> usize {
        match self {
            AacsVersion::V10 => 48,
            AacsVersion::V20 | AacsVersion::V21 => 64,
        }
    }

    /// This version as the major integer ([`AACS_MAJOR_BD`] / [`AACS_MAJOR_UHD`]).
    pub fn major(self) -> u8 {
        match self {
            AacsVersion::V10 => AACS_MAJOR_BD,
            AacsVersion::V20 | AacsVersion::V21 => AACS_MAJOR_UHD,
        }
    }

    /// The version a bare major integer selects for stride purposes: only the
    /// BD major is V10; every other value takes the V20/V21 64-byte stride.
    pub fn from_major(major: u8) -> Self {
        if major == AACS_MAJOR_BD {
            AacsVersion::V10
        } else {
            AacsVersion::V20
        }
    }
}

/// The AACS generation that sets the `Unit_Key_RO.inf` stride, from the strongest signal
/// present: the content certificate's major version, then a pre-recorded MKB's type, then
/// `index.bdmv` (`INDX0300` is UHD). With none of them, UHD's 64-byte stride. The one
/// resolver for every path that parses the unit-key file, so they cannot disagree.
pub fn resolve_aacs_version(
    cert_major: Option<u8>,
    mkb: &[u8],
    index_is_uhd: Option<bool>,
) -> AacsVersion {
    if let Some(major) = cert_major {
        return AacsVersion::from_major(major);
    }
    if let Some(t @ (MkbType::Prerecorded | MkbType::CategoryC20 | MkbType::CategoryC21)) =
        mkb_type(mkb)
    {
        return t.generation();
    }
    match index_is_uhd {
        Some(false) => AacsVersion::V10,
        _ => AacsVersion::V20,
    }
}

/// Find Verify Media Key Record (type 0x81 for AACS 1.0, 0x86 for AACS 2.0/2.1) in MKB.
/// 0x81: `[C]` §3.2.5.1.4. 0x86 (AACS 2.x): `[RE]` — not in the public spec (from real 2.x MKBs).
pub(crate) fn mkb_find_mk_dv(mkb: &[u8]) -> Option<[u8; 16]> {
    // Verify-Media-Key record (0x81 for AACS 1.0, 0x86 for AACS 2.x): mk_dv is
    // the 16 bytes at record offset 4 (body offset 0). Needs rec_len >= 20.
    let found = mkb_records(mkb).find(|&(_, rt, len)| {
        (rt == REC_VERIFY_MEDIA_KEY_V1 || rt == REC_VERIFY_MEDIA_KEY_V2) && len >= 20
    });
    match found {
        Some((o, rec_type, rec_len)) => {
            let mut dv = [0u8; 16];
            dv.copy_from_slice(&mkb[o + 4..o + 20]);
            tracing::debug!(
                target: "freemkv::disc",
                phase = "mkb_mk_dv_found",
                rec_type,
                pos = o,
                rec_len,
                "mk_dv extracted from MKB"
            );
            Some(dv)
        }
        None => {
            tracing::warn!(
                target: "freemkv::disc",
                phase = "mkb_mk_dv_not_found",
                "no 0x81/0x86 record with rec_len>=20 found"
            );
            None
        }
    }
}

/// Find Subset-Difference records (type 0x04) in MKB. `[C]` §3.2.5.1.5.
pub(crate) fn mkb_find_subdiff_records(mkb: &[u8]) -> Option<Vec<u8>> {
    find_record_body(mkb, REC_SUBSET_DIFFERENCE)
}

// Find the Media Key Data Record (cvalues table) in an MKB. `[C]` §3.2.4 / §3.2.5.1.7. 0x05 is
// the pre-recorded / Category C table (libaacs `mkb_cvalues`); a Class II MKB carries 0x0c
// instead (see `check_mkb_class`). 0x07 is the Subset-Difference Index.
pub(crate) fn mkb_find_cvalues(mkb: &[u8]) -> Option<Vec<u8>> {
    find_record_body(mkb, REC_MEDIA_KEY_DATA)
}

/// Walk an MKB and return the payload (header stripped) of the first
/// record matching `rec_type`. Returns `None` if no such record exists or
/// the record is empty.
pub(crate) fn find_record_body(mkb: &[u8], rec_type_wanted: u8) -> Option<Vec<u8>> {
    find_record_slice(mkb, rec_type_wanted).map(<[u8]>::to_vec)
}

// Borrowing form of [`find_record_body`]: the body in place, no copy.
pub(crate) fn find_record_slice(mkb: &[u8], rec_type_wanted: u8) -> Option<&[u8]> {
    mkb_records(mkb)
        .find(|&(_, rt, len)| rt == rec_type_wanted && len > 4)
        .map(|(o, _, len)| &mkb[o + 4..o + len])
}

/// Real content length of an MKB: the byte offset where the record stream
/// ends. MKB files (especially `MKB_RW.inf`, but `MKB_RO.inf` too on some
/// discs) are allocated to a fixed size — often ~128 MiB — with the records at
/// the front and the rest zero padding. Walking records (type+len) and stopping
/// at the first padding byte (`type == 0` / zero-length / overrun) gives the
/// actual size so callers can trim off megabytes of zeros before sending or
/// archiving. Returns `mkb.len()` only if the whole buffer parsed as records.
pub fn mkb_content_len(mkb: &[u8]) -> usize {
    // End of the last framed record = where the fixed-region zero padding begins.
    // (The `00 000000` terminator / overrun stops the walk; real MKBs pad with
    // zeros, so this matches the prior "stop at the first padding byte".)
    mkb_records(mkb)
        .last()
        .map(|(o, _, len)| o + len)
        .unwrap_or(0)
}

/// How far a prefix `buf` of an MKB file shows its record stream, walking the record
/// headers (4 bytes: type + BE24 length) as [`mkb_content_len`] does. `Ok(end)`: the walk
/// stopped on a header inside `buf` (the `00 000000` end marker, padding, or a header that
/// frames no record), so the stream ends at `end`. `Err(need)`: a record, or the header
/// after it, runs past `buf`; the stream needs at least `need` bytes to be seen.
pub(crate) fn mkb_prefix_end(buf: &[u8]) -> std::result::Result<usize, usize> {
    let mut pos = 0usize;
    loop {
        let Some(head) = buf.get(pos..pos + 4) else {
            return Err(pos + 4);
        };
        let rec_len = ((head[1] as usize) << 16) | ((head[2] as usize) << 8) | (head[3] as usize);
        if (head[0] == 0 && rec_len == 0) || rec_len < 4 {
            return Ok(pos);
        }
        if pos + rec_len > buf.len() {
            return Err(pos + rec_len + 4);
        }
        pos += rec_len;
    }
}

/// Trim an MKB's trailing fixed-region padding to its real content length —
/// but ONLY when [`mkb_content_len`] actually found one. It returns 0 for an
/// MKB whose first record cannot be parsed; truncating to 0 in that case would
/// hand downstream consumers (and the online key service) an EMPTY MKB that can
/// never resolve. So a 0 (or a length that isn't strictly inside the buffer)
/// leaves the MKB untouched. A 0.31.0 regression dropped this guard and
/// `truncate`-d unconditionally, zeroing unrecognised MKBs.
pub fn trim_mkb(mut mkb: Vec<u8>) -> Vec<u8> {
    let n = mkb_content_len(&mkb);
    if n > 0 && n < mkb.len() {
        mkb.truncate(n);
    }
    mkb
}

/// Get MKB version from Type and Version Record (type 0x10).
/// Layout: 4-byte record header at `pos` (type + BE24 length), then the
/// record body starts at `pos + 4`. The body holds the BE u32 Type field at
/// body offset 0 (`pos + 4`), then the BE u32 version at body offset 4
/// (`pos + 8`).
pub fn mkb_version(mkb: &[u8]) -> Option<u32> {
    // Type-and-Version record (0x10): version is the BE u32 at body offset 4
    // (record offset 8). Needs rec_len >= 12 (4 header + 4 type + 4 version).
    mkb_records(mkb)
        .find(|&(_, rt, len)| rt == REC_TYPE_AND_VERSION && len >= 12)
        .map(|(o, _, _)| u32::from_be_bytes([mkb[o + 8], mkb[o + 9], mkb[o + 10], mkb[o + 11]]))
}

/// `0x00031003` — recordable media MKB (Class I & II compute Km directly).
pub const MKB_TYPE_3_RECORDABLE: u32 = 0x0003_1003;

/// `0x00041003` — AACS 1.0 pre-recorded content MKB (KCD-based). Standard BD.
pub const MKB_TYPE_4_PRERECORDED: u32 = 0x0004_1003;

/// `0x000A1003` — Class II / Unified MKB (Sequence-Key-Block functionality).
pub const MKB_TYPE_10_CLASS_II: u32 = 0x000A_1003;

/// `0x48141003` — AACS 2.0 Category C (UHD content) MKB type value.
pub const MKB_20_CATEGORY_C: u32 = 0x4814_1003;

/// `0x48151003` — AACS 2.1 Category C (UHD content) MKB type value.
pub const MKB_21_CATEGORY_C: u32 = 0x4815_1003;

/// The AACS MKB Type field, decoded.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MkbType {
    /// Type 3 — recordable media.
    Recordable,
    /// Type 4 — AACS 1.0 pre-recorded content (KCD). Standard Blu-ray.
    Prerecorded,
    /// Type 10 — Class II / Unified (SKB).
    ClassII,
    /// AACS 2.0 Category C — UHD content.
    CategoryC20,
    /// AACS 2.1 Category C — UHD content.
    CategoryC21,
    /// Unrecognized MKBType value (raw field preserved).
    Other(u32),
}

impl MkbType {
    pub(crate) fn from_raw(raw: u32) -> Self {
        match raw {
            MKB_TYPE_3_RECORDABLE => MkbType::Recordable,
            MKB_TYPE_4_PRERECORDED => MkbType::Prerecorded,
            MKB_TYPE_10_CLASS_II => MkbType::ClassII,
            MKB_20_CATEGORY_C => MkbType::CategoryC20,
            MKB_21_CATEGORY_C => MkbType::CategoryC21,
            other => MkbType::Other(other),
        }
    }

    /// AACS generation this MKB belongs to (Category C → 2.0/2.1, else 1.0).
    pub fn generation(self) -> AacsVersion {
        match self {
            MkbType::CategoryC21 => AacsVersion::V21,
            MkbType::CategoryC20 => AacsVersion::V20,
            _ => AacsVersion::V10,
        }
    }

    /// `true` for UHD (AACS 2.x Category C); `false` for Blu-ray (AACS 1.x).
    pub fn is_uhd(self) -> bool {
        matches!(self, MkbType::CategoryC20 | MkbType::CategoryC21)
    }
}

/// The raw 32-bit MKBType field from the Type-and-Version record (0x10), bytes
/// 4-7. `None` if no 0x10 record is present. `[C]` §3.2.5.1.1 Table 3-2.
pub fn mkb_type_raw(mkb: &[u8]) -> Option<u32> {
    // Type-and-Version record (0x10): the 32-bit MKBType is bytes 4-7 (body
    // offset 0). Needs rec_len >= 8 (4 header + 4 type).
    mkb_records(mkb)
        .find(|&(_, rt, len)| rt == REC_TYPE_AND_VERSION && len >= 8)
        .map(|(o, _, _)| u32::from_be_bytes([mkb[o + 4], mkb[o + 5], mkb[o + 6], mkb[o + 7]]))
}

/// Decode an MKB's Type field. `None` if no Type-and-Version record is present.
pub fn mkb_type(mkb: &[u8]) -> Option<MkbType> {
    mkb_type_raw(mkb).map(MkbType::from_raw)
}

/// `Some(true)` if this MKB is a UHD (AACS 2.x Category C) block, `Some(false)`
/// for Blu-ray (AACS 1.x), `None` if the Type record is absent.
pub fn mkb_is_uhd(mkb: &[u8]) -> Option<bool> {
    mkb_type(mkb).map(MkbType::is_uhd)
}

#[cfg(test)]
#[path = "mkb_tests.rs"]
mod tests;