use std::collections::HashMap;

use crate::error::AxmlError;
use crate::proto;

pub(crate) const STRING_POOL_TYPE: u16 = 0x0001;
pub(crate) const UTF8_FLAG: u32 = 0x100;
pub(crate) const STRING_POOL_HEADER_SIZE: usize = 28;

/// Complete string pool stored as raw bytes for exact round-trip preservation,
/// plus decoded strings cached for fast access.
pub struct StringPool {
    /// Raw bytes of the entire string pool chunk (for exact repacking when unmodified)
    pub(crate) raw: Option<Vec<u8>>,
    /// If true, raw is invalid and pool must be recomputed on pack()
    pub(crate) dirty: bool,
    /// UTF-8 flag
    pub(crate) is_utf8: bool,
    /// Original flags field (preserved in raw mode)
    pub(crate) flags: u32,
    /// Decoded strings (populated lazily from raw, or directly in built mode)
    pub(crate) strings: Vec<Vec<u8>>, // raw data bytes per string (no length prefix, no null)
    /// O(1) reverse index: encoded bytes → index in `strings`.
    /// Kept in sync with `strings` by `add()`.
    pub(crate) index: HashMap<Vec<u8>, u32>,
    /// String offsets within the string data area (from raw parsing)
    pub(crate) string_offsets: Vec<u32>,
    /// Byte offset of string data within the raw chunk (preserved from parsing)
    #[allow(dead_code)]
    pub(crate) string_data_start: usize,
    /// Style offsets (preserved for round-trip)
    pub(crate) style_offsets: Vec<u32>,
}

impl Default for StringPool {
    fn default() -> Self {
        StringPool {
            raw: None,
            dirty: true,
            is_utf8: false,
            flags: 0,
            strings: Vec::new(),
            index: HashMap::new(),
            string_offsets: Vec::new(),
            string_data_start: 0,
            style_offsets: Vec::new(),
        }
    }
}

impl StringPool {
    /// Create a new empty string pool for building from XML (UTF-16 by default)
    pub fn new(is_utf8: bool) -> Self {
        let flags = if is_utf8 { UTF8_FLAG } else { 0 };
        StringPool {
            raw: None,
            dirty: true,
            is_utf8,
            flags,
            strings: Vec::new(),
            index: HashMap::new(),
            string_offsets: Vec::new(),
            string_data_start: 0,
            style_offsets: Vec::new(),
        }
    }

    /// Parse string pool from binary data. Returns (pool, bytes_consumed).
    pub fn parse(data: &[u8]) -> Result<(Self, usize), AxmlError> {
        if data.len() < STRING_POOL_HEADER_SIZE {
            return Err(AxmlError::UnexpectedEof);
        }

        let chunk_type = u16::from_le_bytes([data[0], data[1]]);
        if chunk_type != STRING_POOL_TYPE {
            return Err(AxmlError::InvalidChunkType(chunk_type));
        }
        let header_size = u16::from_le_bytes([data[2], data[3]]) as usize;
        let chunk_size = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;

        if chunk_size > data.len() {
            return Err(AxmlError::SliceError(format!(
                "String pool chunk_size {} > data.len() {}",
                chunk_size,
                data.len()
            )));
        }

        let string_count = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
        let style_count = u32::from_le_bytes([data[12], data[13], data[14], data[15]]) as usize;
        let flags = u32::from_le_bytes([data[16], data[17], data[18], data[19]]);
        let string_offset = u32::from_le_bytes([data[20], data[21], data[22], data[23]]) as usize;
        let style_offset = u32::from_le_bytes([data[24], data[25], data[26], data[27]]) as usize;
        let is_utf8 = (flags & UTF8_FLAG) != 0;

        // Read string offsets
        let offsets_start = header_size;
        // Cap capacity to prevent OOM from a malformed string_count field.
        let safe_str_cap = string_count.min((chunk_size.saturating_sub(offsets_start)) / 4);
        let mut string_offsets = Vec::with_capacity(safe_str_cap);
        for i in 0..string_count {
            let pos = offsets_start + i * 4;
            if pos + 4 > chunk_size {
                return Err(AxmlError::UnexpectedEof);
            }
            let off = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
            string_offsets.push(off);
        }

        // Read style offsets
        let style_offsets_start = offsets_start + string_count * 4;
        // Cap capacity: shared space with string offsets, each entry 4 bytes.
        let safe_sty_cap = style_count.min((chunk_size.saturating_sub(style_offsets_start)) / 4);
        let mut style_offsets = Vec::with_capacity(safe_sty_cap);
        for i in 0..style_count {
            let pos = style_offsets_start + i * 4;
            if pos + 4 > chunk_size {
                return Err(AxmlError::UnexpectedEof);
            }
            let off = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
            style_offsets.push(off);
        }

        // String data area bounds
        let string_data_end = if style_count > 0 && style_offset > 0 {
            style_offset.min(chunk_size)
        } else {
            chunk_size
        };
        let string_data_start = string_offset;

        // Decode each string (store raw data bytes without prefix/null).
        // Cap capacity: each string needs at least 2 bytes (1-byte prefix + 1-byte null).
        let safe_strings_cap = string_count.min(chunk_size / 2);
        let mut strings: Vec<Vec<u8>> = Vec::with_capacity(safe_strings_cap);

        if string_count > 0 {
            let string_data_area =
                if string_data_start <= string_data_end && string_data_end <= data.len() {
                    &data[string_data_start..string_data_end]
                } else {
                    return Err(AxmlError::SliceError(format!(
                        "Invalid string data range: {}..{}",
                        string_data_start, string_data_end
                    )));
                };

            for &raw_off in string_offsets.iter().take(string_count) {
                let off = raw_off as usize;
                if off >= string_data_area.len() {
                    // Clamp gracefully
                    strings.push(Vec::new());
                    continue;
                }
                let entry_data = &string_data_area[off..];
                let raw_data = decode_string_data(entry_data, is_utf8).unwrap_or_default();
                strings.push(raw_data);
            }
        }

        let raw_chunk = data[..chunk_size].to_vec();

        // Build reverse index from the parsed strings.
        let index = strings
            .iter()
            .enumerate()
            .map(|(i, data)| (data.clone(), i as u32))
            .collect();

        Ok((
            StringPool {
                raw: Some(raw_chunk),
                dirty: false,
                is_utf8,
                flags,
                strings,
                index,
                string_offsets,
                string_data_start,
                style_offsets,
            },
            chunk_size,
        ))
    }

    /// Decode string at index to a Rust String
    pub fn decode(&self, idx: u32) -> Result<String, AxmlError> {
        let idx = idx as usize;
        if idx == 0xffff_ffff {
            return Ok(String::new());
        }
        if idx >= self.strings.len() {
            return Err(AxmlError::InvalidStringIndex(idx, self.strings.len()));
        }
        let data = &self.strings[idx];
        if self.is_utf8 {
            // Zero-copy UTF-8 validation; only allocate on success.
            match std::str::from_utf8(data) {
                Ok(s) => Ok(s.to_owned()),
                Err(_) => Ok(String::from_utf8_lossy(data).into_owned()),
            }
        } else {
            if !data.len().is_multiple_of(2) {
                return Ok(String::from_utf8_lossy(data).into_owned());
            }
            let u16s: Vec<u16> = data
                .chunks_exact(2)
                .map(|b| u16::from_le_bytes([b[0], b[1]]))
                .collect();
            Ok(String::from_utf16_lossy(&u16s))
        }
    }

    /// Find index of a string, or None if not present.
    /// O(1) lookup via the reverse index.
    pub fn find(&self, s: &str) -> Option<u32> {
        let encoded = encode_string_data(s, self.is_utf8);
        self.index.get(&encoded).copied()
    }

    /// Add a string to the pool (or return existing index).
    /// O(1) amortised.  Marks the pool as dirty (raw bytes are stale).
    pub fn add(&mut self, s: &str) -> u32 {
        let encoded = encode_string_data(s, self.is_utf8);
        if let Some(&idx) = self.index.get(&encoded) {
            return idx;
        }
        let idx = self.strings.len() as u32;
        self.index.insert(encoded.clone(), idx);
        self.strings.push(encoded);
        self.dirty = true;
        idx
    }

    /// Force the pool to be recomputed on the next `pack()` call.
    /// Call this (via `Axml::compute()`) before `pack()` to produce a valid binary.
    pub fn mark_dirty(&mut self) {
        self.dirty = true;
        self.raw = None;
    }

    /// Apply new string content without updating the chunk-size header.
    ///
    /// Updates per-entry length prefixes and the string-offset table in the raw
    /// bytes, but leaves the chunk total-size field (`raw[4..8]`) at its original
    /// value.  `pack()` will return this intentionally corrupted binary until
    /// `mark_dirty()` is called (via `compute()`).
    ///
    /// If no raw bytes are available the pool is marked dirty instead (safe fallback).
    pub fn apply_no_compute(&mut self, new_strings: Vec<Vec<u8>>) {
        // Rebuild reverse index
        self.index.clear();
        for (i, s) in new_strings.iter().enumerate() {
            self.index.insert(s.clone(), i as u32);
        }
        self.strings = new_strings;

        if let Some(raw) = &self.raw {
            let string_count = self.strings.len();
            let _style_count = self.style_offsets.len();

            // Layout: 28-byte header | string offsets (4*n) | style offsets (4*m) | string data
            let header_end = STRING_POOL_HEADER_SIZE; // 28

            // Build new string entries (with updated per-entry length prefix)
            let entries: Vec<Vec<u8>> = self
                .strings
                .iter()
                .map(|data| pack_string_entry(data, self.is_utf8))
                .collect();

            // Compute updated per-string offsets
            let mut new_offsets = Vec::with_capacity(string_count);
            let mut cur = 0u32;
            for e in &entries {
                new_offsets.push(cur);
                cur += e.len() as u32;
            }
            self.string_offsets = new_offsets.clone();

            // Reconstruct: keep original header (size field intentionally stale),
            // write new string offsets, preserve style offsets, write new string data.
            let mut new_raw = raw[..header_end].to_vec();
            for off in &new_offsets {
                new_raw.extend_from_slice(&off.to_le_bytes());
            }
            for off in &self.style_offsets {
                new_raw.extend_from_slice(&off.to_le_bytes());
            }
            for e in entries {
                new_raw.extend_from_slice(&e);
            }

            self.raw = Some(new_raw);
            self.dirty = false; // pack() returns the corrupted raw
        } else {
            // No original raw bytes, fall back to full recompute on next pack()
            self.dirty = true;
        }
    }

    /// Serialize the string pool to bytes.
    /// Uses raw bytes if unmodified, otherwise recomputes.
    pub fn pack(&self) -> Vec<u8> {
        if !self.dirty {
            if let Some(raw) = &self.raw {
                return raw.clone();
            }
        }
        self.pack_from_strings()
    }

    fn pack_from_strings(&self) -> Vec<u8> {
        let string_count = self.strings.len() as u32;
        let style_count = self.style_offsets.len() as u32;

        // Extract style span data verbatim from original raw bytes (if any).
        // Style spans reference strings by index and remain valid when new strings
        // are appended because existing indices do not change.
        let style_data: &[u8] = if style_count > 0 {
            if let Some(raw) = &self.raw {
                let orig_styles_start =
                    u32::from_le_bytes([raw[24], raw[25], raw[26], raw[27]]) as usize;
                let orig_chunk_size = u32::from_le_bytes([raw[4], raw[5], raw[6], raw[7]]) as usize;
                if orig_styles_start > 0
                    && orig_styles_start < orig_chunk_size
                    && orig_chunk_size <= raw.len()
                {
                    &raw[orig_styles_start..orig_chunk_size]
                } else {
                    &[]
                }
            } else {
                &[]
            }
        } else {
            &[]
        };

        // Build string entries (with size prefix and null terminator)
        let entries: Vec<Vec<u8>> = self
            .strings
            .iter()
            .map(|data| pack_string_entry(data, self.is_utf8))
            .collect();

        // Compute total string data size (unpadded)
        let total_strings: usize = entries.iter().map(|e| e.len()).sum();

        // Pad string data to 4-byte alignment
        let padding_needed = total_strings.wrapping_neg() % 4;

        // stringsStart = header(28) + string_offsets(4*n) + style_offsets(4*m)
        let string_offset_val = STRING_POOL_HEADER_SIZE as u32 + 4 * string_count + 4 * style_count;
        // stylesStart = stringsStart + string_data + padding
        let styles_start_val = string_offset_val + total_strings as u32 + padding_needed as u32;
        let chunk_size = styles_start_val + style_data.len() as u32;

        let mut out = Vec::with_capacity(chunk_size as usize);

        // 8-byte chunk header
        out.extend_from_slice(&STRING_POOL_TYPE.to_le_bytes());
        out.extend_from_slice(&(STRING_POOL_HEADER_SIZE as u16).to_le_bytes());
        out.extend_from_slice(&chunk_size.to_le_bytes());

        // String pool specific fields (20 bytes)
        out.extend_from_slice(&string_count.to_le_bytes());
        out.extend_from_slice(&style_count.to_le_bytes());
        out.extend_from_slice(&self.flags.to_le_bytes());
        out.extend_from_slice(&string_offset_val.to_le_bytes());
        // stylesStart is only non-zero when there are style spans
        let styles_start_field = if style_count > 0 {
            styles_start_val
        } else {
            0u32
        };
        out.extend_from_slice(&styles_start_field.to_le_bytes());

        // String offsets
        let mut cur_off = 0u32;
        for entry in &entries {
            out.extend_from_slice(&cur_off.to_le_bytes());
            cur_off += entry.len() as u32;
        }

        // Style offsets (preserved from original; relative to stylesStart, still valid)
        for &off in &self.style_offsets {
            out.extend_from_slice(&off.to_le_bytes());
        }

        // String data
        for entry in &entries {
            out.extend_from_slice(entry);
        }

        // Alignment padding
        out.extend(std::iter::repeat_n(0u8, padding_needed));

        // Style span data (copied verbatim from original raw bytes)
        out.extend_from_slice(style_data);

        out
    }
}

// ─────────────────────────── StringBlocks wrapper ────────────────────────────

/// High-level wrapper that mirrors the Python `StringBlocks` class.
/// Holds both the efficient binary representation (`inner: StringPool`) and
/// the prost-generated protobuf message (`proto: proto::StringBlocks`).
/// The `proto` field is populated by `Axml::update_proto()` and should be
/// treated as a view that may lag behind `inner` if strings are added manually.
pub struct StringBlocks {
    pub(crate) proto: proto::StringBlocks,
    pub(crate) inner: StringPool,
}

impl StringBlocks {
    /// Create an empty `StringBlocks` (used when constructing a new `Axml`).
    pub fn new() -> Self {
        StringBlocks {
            proto: proto::StringBlocks::default(),
            inner: StringPool::default(),
        }
    }

    /// Wrap an existing `StringPool` together with its pre-computed proto message.
    pub fn from_pool_and_proto(inner: StringPool, proto: proto::StringBlocks) -> Self {
        StringBlocks { proto, inner }
    }

    /// Return a reference to the prost-generated protobuf message.
    #[allow(dead_code)]
    pub fn get_proto(&self) -> &proto::StringBlocks {
        &self.proto
    }
}

impl Default for StringBlocks {
    fn default() -> Self {
        Self::new()
    }
}

/// Decode just the string data bytes (without the size prefix, without null) from a buffer entry.
fn decode_string_data(data: &[u8], is_utf8: bool) -> Option<Vec<u8>> {
    if is_utf8 {
        let (byte_count, offset) = read_utf8_lengths(data)?;
        let end = offset + byte_count;
        if end > data.len() {
            // Graceful: take what's there
            Some(data[offset..data.len().min(end)].to_vec())
        } else {
            Some(data[offset..end].to_vec())
        }
    } else {
        let (char_count, offset) = read_utf16_length(data)?;
        let byte_count = char_count * 2;
        let end = offset + byte_count;
        if end > data.len() {
            Some(data[offset..data.len().min(end)].to_vec())
        } else {
            Some(data[offset..end].to_vec())
        }
    }
}

/// Read UTF-8 length prefix: returns (byte_count, data_start_offset).
/// UTF-8 format: [char_count_bytes][byte_count_bytes]
fn read_utf8_lengths(data: &[u8]) -> Option<(usize, usize)> {
    let mut pos = 0;

    // First field: char count (may be 1 or 2 bytes)
    if pos >= data.len() {
        return None;
    }
    let first = data[pos];
    pos += 1;
    let _char_count = if (first & 0x80) != 0 {
        if pos >= data.len() {
            return None;
        }
        let second = data[pos];
        pos += 1;
        (((first & 0x7f) as usize) << 8) | (second as usize)
    } else {
        first as usize
    };

    // Second field: byte count (may be 1 or 2 bytes)
    if pos >= data.len() {
        return None;
    }
    let first = data[pos];
    pos += 1;
    let byte_count = if (first & 0x80) != 0 {
        if pos >= data.len() {
            return None;
        }
        let second = data[pos];
        pos += 1;
        (((first & 0x7f) as usize) << 8) | (second as usize)
    } else {
        first as usize
    };

    Some((byte_count, pos))
}

/// Read UTF-16 length prefix: returns (char_count, data_start_offset).
fn read_utf16_length(data: &[u8]) -> Option<(usize, usize)> {
    if data.len() < 2 {
        return None;
    }
    let first = u16::from_le_bytes([data[0], data[1]]);
    let mut pos = 2;
    let char_count = if (first & 0x8000) != 0 {
        if data.len() < 4 {
            return None;
        }
        let second = u16::from_le_bytes([data[2], data[3]]);
        pos = 4;
        (((first & 0x7fff) as usize) << 16) | (second as usize)
    } else {
        first as usize
    };
    Some((char_count, pos))
}

/// Encode string data bytes (without prefix/null) for storage in strings vec.
fn encode_string_data(s: &str, is_utf8: bool) -> Vec<u8> {
    if is_utf8 {
        s.as_bytes().to_vec()
    } else {
        s.encode_utf16().flat_map(|c| c.to_le_bytes()).collect()
    }
}

/// Build a complete string entry (size_prefix + data + null).
fn pack_string_entry(data: &[u8], is_utf8: bool) -> Vec<u8> {
    let mut out = Vec::new();
    if is_utf8 {
        // Char count (attempt UTF-8 char count; fall back to byte count)
        let char_count = match std::str::from_utf8(data) {
            Ok(s) => s.chars().count(),
            Err(_) => data.len(),
        };
        let byte_count = data.len();

        // Encode char count
        encode_var_byte(char_count, &mut out);
        // Encode byte count
        encode_var_byte(byte_count, &mut out);

        out.extend_from_slice(data);
        out.push(0); // null terminator
    } else {
        let char_count = data.len() / 2;

        // Encode char count as u16(s)
        encode_var_u16(char_count, &mut out);

        out.extend_from_slice(data);
        out.extend_from_slice(&0u16.to_le_bytes()); // null terminator
    }
    out
}

fn encode_var_byte(val: usize, out: &mut Vec<u8>) {
    if val <= 0x7f {
        out.push(val as u8);
    } else {
        out.push(((val >> 8) as u8) | 0x80);
        out.push((val & 0xff) as u8);
    }
}

fn encode_var_u16(val: usize, out: &mut Vec<u8>) {
    if val <= 0x7fff {
        out.extend_from_slice(&(val as u16).to_le_bytes());
    } else {
        let high = ((val >> 16) as u16) | 0x8000;
        let low = (val & 0xffff) as u16;
        out.extend_from_slice(&high.to_le_bytes());
        out.extend_from_slice(&low.to_le_bytes());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── StringPool::add / find ────────────────────────────────────────────────

    #[test]
    fn test_add_returns_stable_index() {
        let mut pool = StringPool::new(false);
        let i0 = pool.add("hello");
        let i1 = pool.add("world");
        let i2 = pool.add("hello"); // duplicate, must return same index
        assert_eq!(i0, 0);
        assert_eq!(i1, 1);
        assert_eq!(i2, i0, "duplicate add should return the existing index");
    }

    #[test]
    fn test_find_existing() {
        let mut pool = StringPool::new(false);
        pool.add("foo");
        assert_eq!(pool.find("foo"), Some(0));
    }

    #[test]
    fn test_find_missing() {
        let pool = StringPool::new(false);
        assert_eq!(pool.find("bar"), None);
    }

    // ── StringPool::decode ────────────────────────────────────────────────────

    #[test]
    fn test_decode_utf16() {
        let mut pool = StringPool::new(false);
        pool.add("hello");
        assert_eq!(pool.decode(0).unwrap(), "hello");
    }

    #[test]
    fn test_decode_utf8() {
        let mut pool = StringPool::new(true);
        pool.add("héllo");
        assert_eq!(pool.decode(0).unwrap(), "héllo");
    }

    #[test]
    fn test_decode_sentinel_returns_empty() {
        let pool = StringPool::new(false);
        assert_eq!(pool.decode(0xffff_ffff).unwrap(), "");
    }

    #[test]
    fn test_decode_out_of_range() {
        let pool = StringPool::new(false);
        assert!(pool.decode(5).is_err());
    }

    // ── StringPool::pack / parse round-trip ───────────────────────────────────

    #[test]
    fn test_pack_parse_utf16_roundtrip() {
        let mut pool = StringPool::new(false);
        pool.add("hello");
        pool.add("world");
        pool.add("android");

        let packed = pool.pack();
        let (parsed, consumed) = StringPool::parse(&packed).expect("parse should succeed");

        assert_eq!(consumed, packed.len());
        assert_eq!(parsed.strings.len(), 3);
        assert_eq!(parsed.decode(0).unwrap(), "hello");
        assert_eq!(parsed.decode(1).unwrap(), "world");
        assert_eq!(parsed.decode(2).unwrap(), "android");
    }

    #[test]
    fn test_pack_parse_utf8_roundtrip() {
        let mut pool = StringPool::new(true);
        pool.add("café");
        pool.add("naïve");

        let packed = pool.pack();
        let (parsed, _) = StringPool::parse(&packed).expect("parse UTF-8 should succeed");

        assert_eq!(parsed.decode(0).unwrap(), "café");
        assert_eq!(parsed.decode(1).unwrap(), "naïve");
    }

    #[test]
    fn test_pack_empty_pool() {
        let pool = StringPool::new(false);
        let packed = pool.pack();
        let (parsed, _) = StringPool::parse(&packed).expect("parse empty pool should succeed");
        assert_eq!(parsed.strings.len(), 0);
    }

    // ── SEC-01: capacity cap ──────────────────────────────────────────────────

    #[test]
    fn test_parse_rejects_oversized_chunk() {
        // Claim chunk_size > data.len(), should return an error, not OOM.
        let mut data = vec![0u8; 28];
        data[0] = 0x01; // STRING_POOL_TYPE low byte
        data[1] = 0x00; // STRING_POOL_TYPE high byte
        data[2] = 28; // header_size
        data[3] = 0;
        // chunk_size = 0xFFFF_FFFF (way more than data.len())
        data[4] = 0xFF;
        data[5] = 0xFF;
        data[6] = 0xFF;
        data[7] = 0xFF;
        assert!(StringPool::parse(&data).is_err());
    }
}
