//! Conversions between the internal [`Axml`] representation and the
//! prost-generated protobuf types in [`crate::proto`].

use prost::Message;

use crate::axml::Axml;
use crate::error::AxmlError;
use crate::proto;
use crate::resource_map::ResourceMap;
use crate::string_pool::{StringBlocks, StringPool};
use crate::xml_element::{Attribute, XmlElement};

// ─────────────────────────── Public API ──────────────────────────────────────

impl Axml {
    /// Serialize this AXML file as a binary protobuf message (proto encoding,
    /// not the Android binary AXML format).
    pub fn to_proto_bytes(&self) -> Vec<u8> {
        axml_to_proto(self).encode_to_vec()
    }

    /// Deserialize an AXML file from binary protobuf bytes produced by
    /// [`to_proto_bytes`].
    pub fn from_proto_bytes(data: &[u8]) -> Result<Self, AxmlError> {
        let p = proto::Axml::decode(data)?;
        proto_to_axml(p)
    }

    /// Convert to the prost-generated [`proto::Axml`] message for in-process
    /// manipulation.
    pub fn to_proto(&self) -> proto::Axml {
        axml_to_proto(self)
    }

    /// Reconstruct an [`Axml`] from a prost-generated [`proto::Axml`] message.
    pub fn from_proto(p: proto::Axml) -> Result<Self, AxmlError> {
        proto_to_axml(p)
    }

    /// Recompute `self.proto` and `self.stringblocks.proto` from the current
    /// in-memory state.  Called automatically at the end of `from_axml()` and
    /// `from_xml()`.  Call manually if you modify `stringblocks.inner` or
    /// `elements` directly.
    pub fn update_proto(&mut self) {
        let sb_proto: proto::StringBlocks = string_pool_to_proto(&self.stringblocks.inner);
        self.stringblocks.proto = sb_proto.clone();
        self.proto = proto::Axml {
            header_xml: Some(proto::AxmlHeader {
                r#type: proto::ResType::ResXmlType as i32,
                header_size: self.file_header_size as u32,
                size: self.total_size,
            }),
            stringblocks: Some(sb_proto),
            resourcemap: self.resource_map.as_ref().map(resource_map_to_proto),
            resourcexml: Some(elements_to_proto(&self.elements)),
        };
    }
}

// ───────────────────────── Axml → proto ──────────────────────────────────────

pub(crate) fn axml_to_proto(axml: &Axml) -> proto::Axml {
    proto::Axml {
        header_xml: Some(proto::AxmlHeader {
            r#type: proto::ResType::ResXmlType as i32,
            header_size: axml.file_header_size as u32,
            size: axml.total_size,
        }),
        stringblocks: Some(string_pool_to_proto(&axml.stringblocks.inner)),
        resourcemap: axml.resource_map.as_ref().map(resource_map_to_proto),
        resourcexml: Some(elements_to_proto(&axml.elements)),
    }
}

pub(crate) fn string_pool_to_proto(sp: &StringPool) -> proto::StringBlocks {
    // Read all layout fields directly from the struct to avoid a full re-serialization.
    let string_count = sp.strings.len() as u32;
    let style_count = sp.style_offsets.len() as u32;

    // When the pool is clean, read the exact byte offsets from the cached raw buffer;
    // otherwise compute them mathematically (pack_from_strings never writes styles, so
    // strings_start = header(28) + 4*string_count).
    let (pool_size, strings_start, styles_start) =
        if let Some(raw) = sp.raw.as_ref().filter(|_| !sp.dirty) {
            if raw.len() >= 28 {
                (
                    raw.len() as u32,
                    u32::from_le_bytes([raw[20], raw[21], raw[22], raw[23]]),
                    u32::from_le_bytes([raw[24], raw[25], raw[26], raw[27]]),
                )
            } else {
                (0u32, 28 + 4 * string_count, 0u32)
            }
        } else {
            // pool_size is informational; 0 is safe when dirty.
            (0u32, 28 + 4 * string_count + 4 * style_count, 0u32)
        };

    let hnd = proto::AxmlHeaderStringpool {
        hnd: Some(proto::AxmlHeader {
            r#type: proto::ResType::ResStringPoolType as i32,
            header_size: 28,
            size: pool_size,
        }),
        len_stringblocks: string_count,
        len_styleblocks: style_count,
        flag: sp.flags,
        stringoffset: strings_start,
        styleoffset: styles_start,
    };

    // Reuse cached string offsets when the pool is clean; they are informational in
    // the proto view and recomputed from scratch by pack_from_strings() on dirty pools.
    let stringoffsets = if !sp.dirty && !sp.string_offsets.is_empty() {
        sp.string_offsets.clone()
    } else {
        Vec::new()
    };

    let styleoffsets = sp.style_offsets.clone();

    // Build StringBlock entries from the raw string bytes stored in the pool
    let stringblocks: Vec<proto::StringBlock> = sp
        .strings
        .iter()
        .map(|raw| {
            let (size, redundant_size, padding) = if sp.is_utf8 {
                let char_count = String::from_utf8_lossy(raw).chars().count() as u32;
                let byte_count = raw.len() as u32;
                (char_count, Some(byte_count), vec![0u8])
            } else {
                (raw.len() as u32 / 2, None, vec![0u8, 0u8])
            };
            proto::StringBlock {
                size,
                redundant_size,
                data: raw.clone(),
                padding,
            }
        })
        .collect();

    proto::StringBlocks {
        hnd: Some(hnd),
        stringoffsets,
        styleoffsets,
        stringblocks,
        pad_stringblocks: vec![],
        styleblocks: vec![],
        pad_styleblocks: vec![],
    }
}

fn resource_map_to_proto(rm: &ResourceMap) -> proto::ResourceMap {
    proto::ResourceMap {
        header: Some(proto::AxmlHeader {
            r#type: proto::ResType::ResXmlResourceMapType as i32,
            header_size: 8,
            size: (8 + rm.ids.len() * 4) as u32,
        }),
        res: rm.ids.clone(),
    }
}

fn elements_to_proto(elements: &[XmlElement]) -> proto::ResourceXml {
    proto::ResourceXml {
        elts: elements.iter().map(xml_element_to_proto).collect(),
    }
}

fn xml_element_to_proto(elt: &XmlElement) -> proto::XmlElement {
    let chunk_size = elt.pack().len() as u32;
    let chunk_type = elt.chunk_type() as i32;

    let _header = Some(proto::AxmlHeader {
        r#type: chunk_type,
        header_size: 16,
        size: chunk_size,
    });

    match elt {
        XmlElement::StartNamespace {
            header_size,
            chunk_size,
            line_number,
            comment,
            prefix,
            uri,
            ..
        } => proto::XmlElement {
            header: Some(proto::AxmlHeader {
                r#type: proto::ResType::ResXmlFirstChunkType as i32,
                header_size: *header_size as u32,
                size: *chunk_size,
            }),
            start_ns: Some(proto::ResXmlStartNamespace {
                generic: Some(res_xml_generic(*line_number, *comment)),
                prefix: *prefix,
                uri: *uri,
            }),
            ..Default::default()
        },

        XmlElement::EndNamespace {
            header_size,
            chunk_size,
            line_number,
            comment,
            prefix,
            uri,
            ..
        } => proto::XmlElement {
            header: Some(proto::AxmlHeader {
                r#type: proto::ResType::ResXmlEndNamespaceType as i32,
                header_size: *header_size as u32,
                size: *chunk_size,
            }),
            end_ns: Some(proto::ResXmlEndNamespace {
                generic: Some(res_xml_generic(*line_number, *comment)),
                prefix: *prefix,
                uri: *uri,
            }),
            ..Default::default()
        },

        XmlElement::StartElement {
            header_size,
            chunk_size,
            line_number,
            comment,
            namespace_uri,
            name,
            at_start,
            at_size,
            style_attribute,
            class_attribute,
            attributes,
            ..
        } => {
            let proto_attrs = attributes
                .iter()
                .map(|a| proto::Attribute {
                    namespace_uri: a.namespace_uri,
                    name: a.name,
                    value: a.value,
                    r#type: a.type_,
                    data: a.data,
                    padding: a.padding.clone(),
                })
                .collect();

            proto::XmlElement {
                header: Some(proto::AxmlHeader {
                    r#type: proto::ResType::ResXmlStartElementType as i32,
                    header_size: *header_size as u32,
                    size: *chunk_size,
                }),
                start_elt: Some(proto::ResXmlStartElement {
                    generic: Some(res_xml_generic(*line_number, *comment)),
                    namespace_uri: *namespace_uri,
                    name: *name,
                    at_start: *at_start as u32,
                    at_size: *at_size as u32,
                    attributevalue: 0,
                    len_attributes: attributes.len() as u32,
                    style_attribute: *style_attribute as i32,
                    class_attribute: *class_attribute as i32,
                    attributes: proto_attrs,
                }),
                ..Default::default()
            }
        }

        XmlElement::EndElement {
            header_size,
            chunk_size,
            line_number,
            comment,
            namespace_uri,
            name,
            ..
        } => proto::XmlElement {
            header: Some(proto::AxmlHeader {
                r#type: proto::ResType::ResXmlEndElementType as i32,
                header_size: *header_size as u32,
                size: *chunk_size,
            }),
            end_elt: Some(proto::ResXmlEndElement {
                generic: Some(res_xml_generic(*line_number, *comment)),
                namespace_uri: *namespace_uri,
                name: *name,
            }),
            ..Default::default()
        },

        XmlElement::CData {
            header_size,
            chunk_size,
            line_number,
            comment,
            name,
            res_size,
            res_res0,
            res_data_type,
            res_data,
            ..
        } => proto::XmlElement {
            header: Some(proto::AxmlHeader {
                r#type: proto::ResType::ResXmlCdataType as i32,
                header_size: *header_size as u32,
                size: *chunk_size,
            }),
            cdata: Some(proto::ResXmlcdata {
                generic: Some(res_xml_generic(*line_number, *comment)),
                name: *name,
                res: Some(proto::ArscResStringPoolRef {
                    size: *res_size as u32,
                    res0: *res_res0 as u32,
                    data_type: *res_data_type as u32,
                    data: *res_data,
                }),
            }),
            ..Default::default()
        },
    }
}

#[inline]
fn res_xml_generic(line_number: u32, comment: u32) -> proto::ResXml {
    proto::ResXml {
        line_number,
        comment,
    }
}

// ───────────────────────── proto → Axml ──────────────────────────────────────

fn proto_to_axml(p: proto::Axml) -> Result<Axml, AxmlError> {
    let (file_type, file_header_size, total_size) = p
        .header_xml
        .map(|h| (h.r#type as u16, h.header_size as u16, h.size))
        .unwrap_or((crate::axml::RES_XML_TYPE, 8, 0));

    let sb_proto = p.stringblocks.unwrap_or_default();
    let string_pool = proto_to_string_pool(sb_proto.clone());

    let resource_map = p.resourcemap.map(proto_to_resource_map);

    let elements = p
        .resourcexml
        .map(proto_to_elements)
        .transpose()?
        .unwrap_or_default();

    let mut axml = Axml {
        proto: proto::Axml::default(),
        stringblocks: StringBlocks::from_pool_and_proto(string_pool, sb_proto),
        resource_map,
        elements,
        file_type,
        file_header_size,
        total_size,
    };
    axml.update_proto();
    Ok(axml)
}

pub(crate) fn proto_to_string_pool(sb: proto::StringBlocks) -> StringPool {
    let (is_utf8, flags) = sb
        .hnd
        .as_ref()
        .map(|h| {
            let f = h.flag;
            ((f & crate::string_pool::UTF8_FLAG) != 0, f)
        })
        .unwrap_or((false, 0));

    let strings: Vec<Vec<u8>> = sb.stringblocks.into_iter().map(|s| s.data).collect();
    let string_offsets = sb.stringoffsets;
    let style_offsets = sb.styleoffsets;

    // Rebuild the reverse index so find() / add() stay O(1).
    let index = strings
        .iter()
        .enumerate()
        .map(|(i, data)| (data.clone(), i as u32))
        .collect();

    StringPool {
        raw: None,
        dirty: true,
        is_utf8,
        flags,
        strings,
        index,
        string_offsets,
        string_data_start: 0,
        style_offsets,
    }
}

fn proto_to_resource_map(rm: proto::ResourceMap) -> ResourceMap {
    let header_size = rm
        .header
        .as_ref()
        .map(|h| h.header_size as u16)
        .unwrap_or(8);
    let chunk_size = rm.header.as_ref().map(|h| h.size).unwrap_or(0);
    ResourceMap {
        ids: rm.res,
        header_size,
        chunk_size,
    }
}

fn proto_to_elements(rxml: proto::ResourceXml) -> Result<Vec<XmlElement>, AxmlError> {
    rxml.elts
        .into_iter()
        .filter_map(proto_to_xml_element)
        .collect::<Result<Vec<_>, _>>()
}

pub(crate) fn proto_to_xml_element(
    elt: proto::XmlElement,
) -> Option<Result<XmlElement, AxmlError>> {
    let header_size = elt
        .header
        .as_ref()
        .map(|h| h.header_size as u16)
        .unwrap_or(16);
    let chunk_size = elt.header.as_ref().map(|h| h.size).unwrap_or(0);

    if let Some(ns) = elt.start_ns {
        let (ln, cm) = line_comment(&ns.generic);
        return Some(Ok(XmlElement::StartNamespace {
            header_size,
            chunk_size,
            line_number: ln,
            comment: cm,
            prefix: ns.prefix,
            uri: ns.uri,
        }));
    }
    if let Some(ns) = elt.end_ns {
        let (ln, cm) = line_comment(&ns.generic);
        return Some(Ok(XmlElement::EndNamespace {
            header_size,
            chunk_size,
            line_number: ln,
            comment: cm,
            prefix: ns.prefix,
            uri: ns.uri,
        }));
    }
    if let Some(se) = elt.start_elt {
        let (ln, cm) = line_comment(&se.generic);
        let attrs = se
            .attributes
            .into_iter()
            .map(|a| Attribute {
                namespace_uri: a.namespace_uri,
                name: a.name,
                value: a.value,
                type_: a.r#type,
                data: a.data,
                padding: a.padding,
            })
            .collect();
        return Some(Ok(XmlElement::StartElement {
            header_size,
            chunk_size,
            line_number: ln,
            comment: cm,
            namespace_uri: se.namespace_uri,
            name: se.name,
            at_start: se.at_start as i16,
            at_size: se.at_size as i16,
            style_attribute: se.style_attribute as i16,
            class_attribute: se.class_attribute as i16,
            attributes: attrs,
        }));
    }
    if let Some(ee) = elt.end_elt {
        let (ln, cm) = line_comment(&ee.generic);
        return Some(Ok(XmlElement::EndElement {
            header_size,
            chunk_size,
            line_number: ln,
            comment: cm,
            namespace_uri: ee.namespace_uri,
            name: ee.name,
        }));
    }
    if let Some(cd) = elt.cdata {
        let (ln, cm) = line_comment(&cd.generic);
        let (res_size, res_res0, res_data_type, res_data) = cd
            .res
            .map(|r| (r.size as u16, r.res0 as u8, r.data_type as u8, r.data))
            .unwrap_or((0, 0, 0, 0));
        return Some(Ok(XmlElement::CData {
            header_size,
            chunk_size,
            line_number: ln,
            comment: cm,
            name: cd.name,
            res_size,
            res_res0,
            res_data_type,
            res_data,
        }));
    }
    None
}

#[inline]
fn line_comment(g: &Option<proto::ResXml>) -> (u32, u32) {
    g.as_ref()
        .map(|r| (r.line_number, r.comment))
        .unwrap_or((0, 0xffff_ffff))
}

// ─────────────────────────────── Arsc ────────────────────────────────────────

use crate::arsc::{
    Arsc, ArscEntry, ArscPackage, ArscResChunk, ArscTypeSpec, ArscTypeType, ArscValue,
};

impl Arsc {
    /// Serialize `self.proto` as binary protobuf bytes.
    pub fn to_proto_bytes(&self) -> Vec<u8> {
        self.proto.encode_to_vec()
    }

    /// Deserialize an [`Arsc`] from binary protobuf bytes produced by [`to_proto_bytes`].
    pub fn from_proto_bytes(data: &[u8]) -> Result<Self, AxmlError> {
        let p = proto::Arsc::decode(data)?;
        proto_to_arsc(p)
    }

    /// Reconstruct an [`Arsc`] from a prost-generated [`proto::Arsc`] message.
    pub fn from_proto(p: proto::Arsc) -> Result<Self, AxmlError> {
        proto_to_arsc(p)
    }

    /// Convert to the prost-generated [`proto::Arsc`] message for in-process manipulation.
    pub fn to_proto(&self) -> proto::Arsc {
        self.proto.clone()
    }

    /// Recompute `self.proto` and `self.stringblocks.proto` from the current
    /// in-memory state.  Called automatically at the end of `from_axml()`.
    /// Call manually if you modify `stringblocks.inner` or `packages` directly.
    pub fn update_proto(&mut self) {
        let sb_proto = string_pool_to_proto(&self.stringblocks.inner);
        self.stringblocks.proto = sb_proto.clone();

        let restablespackage: Vec<proto::AxmlResTablePackage> = self
            .packages
            .iter_mut()
            .map(|pkg| {
                let p = arsc_package_to_proto(pkg);
                pkg.proto = p.clone();
                p
            })
            .collect();

        self.proto = proto::Arsc {
            header_res: Some(proto::AxmlHeaderRestable {
                hnd: Some(proto::AxmlHeader {
                    r#type: proto::ResType::ResTableType as i32,
                    header_size: 12,
                    size: self.total_size,
                }),
                package_count: self.package_count,
            }),
            stringblocks: Some(sb_proto),
            restablespackage,
        };
    }
}

fn arsc_package_to_proto(pkg: &ArscPackage) -> proto::AxmlResTablePackage {
    let type_sp = string_pool_to_proto(&pkg.type_strings);
    let key_sp = string_pool_to_proto(&pkg.key_strings);

    let restypes: Vec<proto::ArscResType> = pkg.chunks.iter().map(arsc_chunk_to_proto).collect();

    proto::AxmlResTablePackage {
        hnd: Some(proto::AxmlHeader {
            r#type: crate::arsc::RES_TABLE_PACKAGE_TYPE as i32,
            header_size: 8,
            size: pkg.chunk_size,
        }),
        id: pkg.id,
        name: pkg.name.clone(),
        type_strings: 0,
        last_public_type: 0,
        key_strings: 0,
        last_public_key: 0,
        padding: vec![],
        type_sp_string: Some(type_sp),
        key_sp_string: Some(key_sp),
        restypes,
    }
}

fn arsc_chunk_to_proto(chunk: &ArscResChunk) -> proto::ArscResType {
    match chunk {
        ArscResChunk::Spec(spec) => proto::ArscResType {
            hnd: Some(proto::AxmlHeader {
                r#type: crate::arsc::RES_TABLE_TYPE_SPEC_TYPE as i32,
                header_size: 8,
                size: (8 + 8 + spec.entries.len() * 4) as u32,
            }),
            typespec: Some(proto::ArscResTypeSpec {
                id: spec.id as u32,
                res0: 0,
                res1: 0,
                entry_count: spec.entries.len() as u32,
                entries: spec.entries.clone(),
            }),
            typetype: None,
        },
        ArscResChunk::Type(tt) => {
            let tables: Vec<proto::ArscResTableEntry> =
                tt.tables.iter().map(arsc_entry_to_proto).collect();
            proto::ArscResType {
                hnd: Some(proto::AxmlHeader {
                    r#type: crate::arsc::RES_TABLE_TYPE_TYPE as i32,
                    header_size: 8,
                    size: 0,
                }),
                typespec: None,
                typetype: Some(proto::ArscResTypeType {
                    id: tt.id as u32,
                    flags: 0,
                    reserved: 0,
                    entry_count: tt.tables.len() as u32,
                    entry_start: 0,
                    config: None,
                    entries: vec![],
                    tables,
                }),
            }
        }
    }
}

fn arsc_entry_to_proto(entry_opt: &Option<ArscEntry>) -> proto::ArscResTableEntry {
    match entry_opt {
        None => proto::ArscResTableEntry {
            size: 0,
            flags: 0,
            index: 0,
            present: false,
            item: None,
            key: None,
        },
        Some(entry) => {
            let key = entry.value.as_ref().map(|v| proto::ArscResStringPoolRef {
                size: v.size as u32,
                res0: v.res0 as u32,
                data_type: v.data_type as u32,
                data: v.data,
            });
            proto::ArscResTableEntry {
                size: 8,
                flags: 0,
                index: entry.key_index,
                present: true,
                item: None,
                key,
            }
        }
    }
}

// ─────────────────────── proto → Arsc ────────────────────────────────────────

pub(crate) fn proto_to_arsc(p: proto::Arsc) -> Result<Arsc, AxmlError> {
    let (total_size, package_count) = p
        .header_res
        .as_ref()
        .map(|h| {
            let sz = h.hnd.as_ref().map(|hnd| hnd.size).unwrap_or(0);
            (sz, h.package_count)
        })
        .unwrap_or((0, 0));

    let sb_proto = p.stringblocks.unwrap_or_default();
    let string_pool = proto_to_string_pool(sb_proto.clone());
    let stringblocks = StringBlocks::from_pool_and_proto(string_pool, sb_proto);

    let packages: Vec<ArscPackage> = p
        .restablespackage
        .into_iter()
        .map(proto_to_arsc_package)
        .collect();

    let mut arsc = Arsc {
        proto: proto::Arsc::default(),
        stringblocks,
        packages,
        total_size,
        package_count,
    };
    arsc.update_proto();
    Ok(arsc)
}

pub(crate) fn proto_to_arsc_package(p: proto::AxmlResTablePackage) -> ArscPackage {
    let type_strings = p
        .type_sp_string
        .map(proto_to_string_pool)
        .unwrap_or_default();
    let key_strings = p
        .key_sp_string
        .map(proto_to_string_pool)
        .unwrap_or_default();

    let chunks: Vec<ArscResChunk> = p
        .restypes
        .into_iter()
        .filter_map(proto_to_arsc_chunk)
        .collect();

    let chunk_size = p.hnd.as_ref().map(|h| h.size).unwrap_or(0);

    ArscPackage {
        proto: proto::AxmlResTablePackage::default(),
        id: p.id,
        name: p.name,
        type_strings,
        key_strings,
        chunks,
        chunk_size,
        last_public_type: 0,
        last_public_key: 0,
    }
}

pub(crate) fn proto_to_arsc_chunk(p: proto::ArscResType) -> Option<ArscResChunk> {
    if let Some(spec) = p.typespec {
        return Some(ArscResChunk::Spec(ArscTypeSpec {
            id: spec.id as u8,
            res0: 0,
            res1: 0,
            entries: spec.entries,
        }));
    }
    if let Some(tt) = p.typetype {
        let tables: Vec<Option<ArscEntry>> =
            tt.tables.into_iter().map(proto_to_arsc_entry).collect();
        let config_raw = {
            let mut cfg = vec![0u8; 32];
            cfg[0] = 32;
            cfg
        };
        return Some(ArscResChunk::Type(ArscTypeType {
            id: tt.id as u8,
            flags: 0,
            reserved: 0,
            language: String::new(),
            region: String::new(),
            config_raw,
            tables,
        }));
    }
    None
}

fn proto_to_arsc_entry(p: proto::ArscResTableEntry) -> Option<ArscEntry> {
    if !p.present {
        return None;
    }
    let value = p.key.map(|k| ArscValue {
        size: k.size as u16,
        res0: k.res0 as u8,
        data_type: k.data_type as u8,
        data: k.data,
    });
    let mut raw_tail = Vec::new();
    let size = if let Some(ref v) = value {
        raw_tail.extend_from_slice(&v.size.to_le_bytes());
        raw_tail.push(v.res0);
        raw_tail.push(v.data_type);
        raw_tail.extend_from_slice(&v.data.to_le_bytes());
        16 // ResTable_entry (8) + Res_value (8)
    } else {
        8 // Just the ResTable_entry header
    };
    Some(ArscEntry {
        size: size as u16,
        flags: 0,
        key_index: p.index,
        value,
        raw_tail,
    })
}
