//! Compute derived size, count, and offset fields for every proto message type.
#![allow(dead_code)]
//!
//! All functions take a `&mut proto::X` and update it in place.  Fields that
//! carry user data (string content, attribute values, line numbers, …) are
//! never touched.
//!
//! The `recursive` parameter, where present, controls whether child messages
//! are also updated before the parent's derived fields are computed.

use prost::Message as _;

use crate::arsc::Arsc;
use crate::axml::Axml;
use crate::error::AxmlError;
use crate::proto;

// ─────────────────────────── StringBlock ─────────────────────────────────────

/// Compute `size` and `redundant_size` of a single StringBlock entry.
///
/// - UTF-8 pool: `size` = Unicode char count; `redundant_size` = byte count.
/// - UTF-16 pool: `size` = char count (UTF-16 code units); `redundant_size` is cleared.
pub fn compute_string_block(sb: &mut proto::StringBlock, is_utf8: bool) {
    let data = &sb.data;
    if is_utf8 {
        let char_count = std::str::from_utf8(data)
            .map(|s| s.chars().count())
            .unwrap_or(data.len()) as u32;
        sb.size = char_count;
        sb.redundant_size = Some(data.len() as u32);
        sb.padding = vec![0u8];
    } else {
        sb.size = data.len() as u32 / 2;
        sb.redundant_size = None;
        sb.padding = vec![0u8, 0u8];
    }
}

// ─────────────────────────── StringBlocks ────────────────────────────────────

/// Byte length of a single string entry as packed by the string pool encoder.
fn string_entry_len(data: &[u8], is_utf8: bool) -> usize {
    if is_utf8 {
        let char_count = std::str::from_utf8(data)
            .map(|s| s.chars().count())
            .unwrap_or(data.len());
        var_byte_len(char_count) + var_byte_len(data.len()) + data.len() + 1
    } else {
        var_u16_len(data.len() / 2) + data.len() + 2
    }
}

fn var_byte_len(n: usize) -> usize {
    if n <= 0x7f {
        1
    } else {
        2
    }
}

fn var_u16_len(n: usize) -> usize {
    if n <= 0x7fff {
        2
    } else {
        4
    }
}

/// Compute all derived fields of a [`proto::StringBlocks`] (string pool) message.
///
/// Updates: `hnd.hnd.size`, `hnd.len_stringblocks`, `hnd.len_styleblocks`,
/// `hnd.stringoffset`, `hnd.styleoffset`, `stringoffsets`.
///
/// If `recursive` is true, also calls [`compute_string_block`] on every entry.
pub fn compute_string_blocks(sbs: &mut proto::StringBlocks, recursive: bool) {
    let is_utf8 = sbs
        .hnd
        .as_ref()
        .map(|h| (h.flag & 256) != 0)
        .unwrap_or(false);
    let flag = sbs.hnd.as_ref().map(|h| h.flag).unwrap_or(0);

    if recursive {
        for blk in &mut sbs.stringblocks {
            compute_string_block(blk, is_utf8);
        }
    }

    let string_count = sbs.stringblocks.len() as u32;
    let style_count = sbs.styleblocks.len() as u32;

    // Compute per-string byte offsets (from start of string data region).
    let mut offsets: Vec<u32> = Vec::with_capacity(string_count as usize);
    let mut cur = 0u32;
    for blk in &sbs.stringblocks {
        offsets.push(cur);
        cur += string_entry_len(&blk.data, is_utf8) as u32;
    }
    let total_strings = cur;
    let padding_needed = total_strings.wrapping_neg() % 4;

    // Header (28) + offset tables + string data + alignment padding.
    let strings_start = 28 + 4 * string_count + 4 * style_count;
    let chunk_size = strings_start + total_strings + padding_needed;

    let hnd = sbs.hnd.get_or_insert_with(Default::default);
    let h = hnd.hnd.get_or_insert_with(Default::default);
    h.r#type = proto::ResType::ResStringPoolType as i32;
    h.header_size = 28;
    h.size = chunk_size;
    hnd.len_stringblocks = string_count;
    hnd.len_styleblocks = style_count;
    hnd.flag = flag;
    hnd.stringoffset = strings_start;
    hnd.styleoffset = 0;

    sbs.stringoffsets = offsets;
    sbs.pad_stringblocks = vec![0u8; padding_needed as usize];
}

// ─────────────────────────── ResourceMap ─────────────────────────────────────

/// Compute `header.size` for a [`proto::ResourceMap`].
pub fn compute_resource_map(rm: &mut proto::ResourceMap) {
    let h = rm.header.get_or_insert_with(Default::default);
    h.r#type = proto::ResType::ResXmlResourceMapType as i32;
    h.header_size = 8;
    h.size = 8 + 4 * rm.res.len() as u32;
}

// ─────────────────────────── XMLElement ──────────────────────────────────────

/// Compute `header.size` (and `start_elt.len_attributes`) for an XML element chunk.
pub fn compute_xml_element(elt: &mut proto::XmlElement) {
    let h = elt.header.get_or_insert_with(Default::default);
    h.header_size = 16;

    if elt.start_ns.is_some() {
        // 8 chunk header + 8 generic (line+comment) + 8 (prefix+uri)
        h.r#type = proto::ResType::ResXmlFirstChunkType as i32;
        h.size = 24;
    } else if elt.end_ns.is_some() {
        h.r#type = proto::ResType::ResXmlEndNamespaceType as i32;
        h.size = 24;
    } else if elt.end_elt.is_some() {
        h.r#type = proto::ResType::ResXmlEndElementType as i32;
        h.size = 24;
    } else if elt.cdata.is_some() {
        // 8 chunk header + 8 generic + 4 name + 8 Res_value
        h.r#type = proto::ResType::ResXmlCdataType as i32;
        h.size = 28;
    } else if let Some(se) = elt.start_elt.as_mut() {
        h.r#type = proto::ResType::ResXmlStartElementType as i32;
        let attr_count = se.attributes.len() as u32;
        // at_size is bytes per attribute (default 20 = 5 × u32)
        let per_attr = if se.at_size > 0 { se.at_size } else { 20 };
        // 8 chunk header + 8 generic + 8 (ns+name) + 4 (at_start+at_size)
        //   + 4 (attr_value_index byte + len_attributes) + 4 (style+class)
        // = 8 + 28 body = 36 … actually: body starts at 8, body = 28 bytes before attrs
        h.size = 8 + 28 + per_attr * attr_count;
        se.len_attributes = attr_count;
        if se.at_start == 0 {
            se.at_start = 0x14;
        }
        if se.at_size == 0 {
            se.at_size = 0x14;
        }
    }
}

// ─────────────────────────── ARSCResTableConfig ───────────────────────────────

/// Compute `size` of an [`proto::ArscResTableConfig`] from its present optional fields.
///
/// Base size is 16 (size field + 3 required fields); each present optional field adds 4.
pub fn compute_arsc_res_table_config(cfg: &mut proto::ArscResTableConfig) {
    let mut size = 16u32; // size(4) + imsi(4) + locale(4) + screenType(4)
    if cfg.input.is_some() {
        size += 4;
    }
    if cfg.screen_size.is_some() {
        size += 4;
    }
    if cfg.version.is_some() {
        size += 4;
    }
    if cfg.screen_config.is_some() {
        size += 4;
    }
    if cfg.screen_size_dp.is_some() {
        size += 4;
    }
    if cfg.locale_script.is_some() {
        size += 4;
    }
    if cfg.locale_variant.is_some() {
        size += 4;
    }
    if cfg.screen_config2.is_some() {
        size += 4;
    }
    size += cfg.padding.len() as u32;
    cfg.size = size;
}

// ─────────────────────────── ARSCResType ─────────────────────────────────────

/// Compute chunk-size and `entry_count` for a TypeSpec [`proto::ArscResType`].
fn compute_arsc_res_type_spec(rt: &mut proto::ArscResType) {
    let spec = match rt.typespec.as_ref() {
        Some(s) => s,
        None => return,
    };
    let entry_count = spec.entries.len() as u32;
    if let Some(s) = rt.typespec.as_mut() {
        s.entry_count = entry_count;
    }
    // Layout: 8 chunk header + 8 body (id:1+res0:1+res1:2+entryCount:4) + 4*N entries.
    let chunk_size = 8 + 8 + 4 * entry_count;
    let h = rt.hnd.get_or_insert_with(Default::default);
    h.r#type = proto::ResType::ResTableTypeSpecType as i32;
    h.header_size = 8;
    h.size = chunk_size;
}

/// Compute chunk-size, `entry_count`, and `entry_start` for a TypeType [`proto::ArscResType`].
fn compute_arsc_res_type_type(rt: &mut proto::ArscResType, recursive: bool) {
    let tt_ref = match rt.typetype.as_ref() {
        Some(t) => t,
        None => return,
    };
    let entry_count = tt_ref.tables.len() as u32;
    let config_size = tt_ref
        .config
        .as_ref()
        .map(|c| c.size)
        .filter(|&s| s > 0)
        .unwrap_or(32);
    let present_count = tt_ref.tables.iter().filter(|e| e.present).count() as u32;

    if recursive {
        if let Some(cfg) = rt.typetype.as_mut().and_then(|t| t.config.as_mut()) {
            compute_arsc_res_table_config(cfg);
        }
    }

    // entries_start is measured from the start of the chunk.
    // Chunk layout: 8 (chunk hdr) + 4 (id+flags+reserved) + 4 (entryCount) + 4 (entryStart)
    //             + config_size + 4*entry_count (offsets)
    let entries_start = 8 + 4 + 4 + 4 + config_size + 4 * entry_count;
    let entry_data_size = 16 * present_count; // 8 ResTable_entry + 8 Res_value each
    let chunk_size = entries_start + entry_data_size;

    let tt = rt.typetype.as_mut().unwrap();
    tt.entry_count = entry_count;
    tt.entry_start = entries_start;

    let h = rt.hnd.get_or_insert_with(Default::default);
    h.r#type = proto::ResType::ResTableTypeType as i32;
    h.header_size = 8;
    h.size = chunk_size;
}

/// Compute derived fields for an [`proto::ArscResType`] (TypeSpec or TypeType).
///
/// Dispatches to the appropriate sub-function based on which `oneof` variant is set.
/// When `recursive` is true, also computes the TypeType config if present.
pub fn compute_arsc_res_type(rt: &mut proto::ArscResType, recursive: bool) {
    if rt.typespec.is_some() {
        compute_arsc_res_type_spec(rt);
    } else if rt.typetype.is_some() {
        compute_arsc_res_type_type(rt, recursive);
    }
}

// ─────────────────────────── AXMLResTablePackage ─────────────────────────────

/// Compute all derived fields for a resource-table package.
///
/// Updates `hnd.size`, `type_strings`, `last_public_type`, `key_strings`,
/// `last_public_key`.  When `recursive` is true, also computes both child
/// string pools and every ResType chunk.
pub fn compute_arsc_res_table_package(pkg: &mut proto::AxmlResTablePackage, recursive: bool) {
    if recursive {
        if let Some(sbs) = pkg.type_sp_string.as_mut() {
            compute_string_blocks(sbs, true);
        }
        if let Some(sbs) = pkg.key_sp_string.as_mut() {
            compute_string_blocks(sbs, true);
        }
        for rt in &mut pkg.restypes {
            compute_arsc_res_type(rt, true);
        }
    }

    const PKG_HEADER: u32 = 288;

    let type_sp_size = pkg
        .type_sp_string
        .as_ref()
        .and_then(|s| s.hnd.as_ref())
        .and_then(|h| h.hnd.as_ref())
        .map(|h| h.size)
        .unwrap_or(0);
    let key_sp_size = pkg
        .key_sp_string
        .as_ref()
        .and_then(|s| s.hnd.as_ref())
        .and_then(|h| h.hnd.as_ref())
        .map(|h| h.size)
        .unwrap_or(0);
    let chunks_size: u32 = pkg
        .restypes
        .iter()
        .filter_map(|rt| rt.hnd.as_ref())
        .map(|h| h.size)
        .sum();

    let last_public_type = pkg
        .type_sp_string
        .as_ref()
        .map(|s| s.stringblocks.len() as u32)
        .unwrap_or(0);
    let last_public_key = pkg
        .key_sp_string
        .as_ref()
        .map(|s| s.stringblocks.len() as u32)
        .unwrap_or(0);

    pkg.type_strings = PKG_HEADER;
    pkg.last_public_type = last_public_type;
    pkg.key_strings = PKG_HEADER + type_sp_size;
    pkg.last_public_key = last_public_key;

    let h = pkg.hnd.get_or_insert_with(Default::default);
    h.r#type = proto::ResType::ResTablePackageType as i32;
    h.header_size = 8;
    h.size = PKG_HEADER + type_sp_size + key_sp_size + chunks_size;
}

// ─────────────────────────── AXML (top-level) ────────────────────────────────

/// Compute all derived size/count fields in an [`proto::Axml`] message.
///
/// Converts through the internal [`Axml`] representation to guarantee byte-exact sizes.
/// If `recursive` is true, all sub-components (string pool, elements) are recomputed.
pub fn compute_axml_proto(axml: &mut proto::Axml, _recursive: bool) -> Result<(), AxmlError> {
    let mut a = Axml::from_proto(axml.clone())?;
    a.compute();
    a.update_proto();
    *axml = a.to_proto();
    Ok(())
}

// ─────────────────────────── ARSC (top-level) ────────────────────────────────

/// Compute all derived size/count fields in an [`proto::Arsc`] message.
///
/// Converts through the internal [`Arsc`] representation to guarantee byte-exact sizes.
/// `recursive` is forwarded to [`Arsc::compute`] to control sub-package recomputation.
pub fn compute_arsc_proto(arsc: &mut proto::Arsc, recursive: bool) -> Result<(), AxmlError> {
    let mut a = Arsc::from_proto(arsc.clone())?;
    a.compute(recursive);
    a.update_proto();
    *arsc = a.to_proto();
    Ok(())
}

// ─────────────────────────── bytes-in / bytes-out ────────────────────────────
// These are the functions exposed to Python.  Each decodes a proto message from
// raw bytes, applies the corresponding compute function, and re-encodes.

/// Compute a single StringBlock proto message; return updated bytes.
pub fn compute_string_block_bytes(data: &[u8], is_utf8: bool) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::StringBlock::decode(data)?;
    compute_string_block(&mut msg, is_utf8);
    Ok(msg.encode_to_vec())
}

/// Compute a StringBlocks proto message; return updated bytes.
pub fn compute_string_blocks_bytes(data: &[u8], recursive: bool) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::StringBlocks::decode(data)?;
    compute_string_blocks(&mut msg, recursive);
    Ok(msg.encode_to_vec())
}

/// Compute a ResourceMap proto message; return updated bytes.
pub fn compute_resource_map_bytes(data: &[u8]) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::ResourceMap::decode(data)?;
    compute_resource_map(&mut msg);
    Ok(msg.encode_to_vec())
}

/// Compute an XMLElement proto message; return updated bytes.
pub fn compute_xml_element_bytes(data: &[u8]) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::XmlElement::decode(data)?;
    compute_xml_element(&mut msg);
    Ok(msg.encode_to_vec())
}

/// Compute an ARSCResTableConfig proto message; return updated bytes.
pub fn compute_arsc_res_table_config_bytes(data: &[u8]) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::ArscResTableConfig::decode(data)?;
    compute_arsc_res_table_config(&mut msg);
    Ok(msg.encode_to_vec())
}

/// Compute an ARSCResType proto message; return updated bytes.
pub fn compute_arsc_res_type_bytes(data: &[u8], recursive: bool) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::ArscResType::decode(data)?;
    compute_arsc_res_type(&mut msg, recursive);
    Ok(msg.encode_to_vec())
}

/// Compute an AXMLResTablePackage proto message; return updated bytes.
pub fn compute_arsc_res_table_package_bytes(
    data: &[u8],
    recursive: bool,
) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::AxmlResTablePackage::decode(data)?;
    compute_arsc_res_table_package(&mut msg, recursive);
    Ok(msg.encode_to_vec())
}

/// Compute an AXML proto message; return updated bytes.
pub fn compute_axml_bytes(data: &[u8], recursive: bool) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::Axml::decode(data)?;
    compute_axml_proto(&mut msg, recursive)?;
    Ok(msg.encode_to_vec())
}

/// Compute an ARSC proto message; return updated bytes.
pub fn compute_arsc_bytes(data: &[u8], recursive: bool) -> Result<Vec<u8>, AxmlError> {
    let mut msg = proto::Arsc::decode(data)?;
    compute_arsc_proto(&mut msg, recursive)?;
    Ok(msg.encode_to_vec())
}
