use crate::header::*;
use aes::cipher::{BlockDecrypt, KeyInit};
use aes::Aes128;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{self, Cursor, Read, Seek, SeekFrom};

/// Parse a XEX2 header from raw file data.
pub fn parse_xex2_header(data: &[u8]) -> io::Result<Xex2Header> {
    let mut cursor = Cursor::new(data);

    let magic = cursor.read_u32::<BigEndian>()?;
    if magic != XEX2_MAGIC {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!("Invalid XEX2 magic: {:#010x} (expected {:#010x})", magic, XEX2_MAGIC),
        ));
    }

    let module_flags = cursor.read_u32::<BigEndian>()?;
    let header_size = cursor.read_u32::<BigEndian>()?;
    let _reserved = cursor.read_u32::<BigEndian>()?;
    let security_offset = cursor.read_u32::<BigEndian>()?;
    let header_count = cursor.read_u32::<BigEndian>()?;

    let mut optional_headers = Vec::new();
    for _ in 0..header_count {
        let key = cursor.read_u32::<BigEndian>()?;
        let value = cursor.read_u32::<BigEndian>()?;
        optional_headers.push(Xex2OptionalHeader { key, value });
    }

    // Parse security info
    let security_info = if (security_offset as usize) < data.len() {
        cursor.seek(SeekFrom::Start(security_offset as u64))?;
        Some(parse_security_info(&mut cursor)?)
    } else {
        None
    };

    // Parse file format info
    let file_format_info = parse_file_format_info(data, &optional_headers);

    // Parse import libraries (addresses only; call resolve_imports after decompression)
    let import_libraries = parse_import_libraries(data, &optional_headers);

    // Parse execution info
    let execution_info = parse_execution_info(data, &optional_headers);

    // Parse original PE name
    let original_pe_name = parse_original_pe_name(data, &optional_headers);

    Ok(Xex2Header {
        magic,
        module_flags,
        header_size,
        security_offset,
        header_count,
        optional_headers,
        security_info,
        file_format_info,
        import_libraries,
        execution_info,
        original_pe_name,
    })
}

fn parse_security_info(cursor: &mut Cursor<&[u8]>) -> io::Result<Xex2SecurityInfo> {
    // xex2_security_info layout (from xex2_info.h):
    // 0x000: header_size (u32)
    // 0x004: image_size (u32)
    // 0x008: rsa_signature (0x100 bytes)
    // 0x108: unk_108 (u32)
    // 0x10C: image_flags (u32)
    // 0x110: load_address (u32)
    // 0x114: section_digest (0x14 bytes)
    // 0x128: import_table_count (u32)
    // 0x12C: import_table_digest (0x14 bytes)
    // 0x140: xgd2_media_id (0x10 bytes)
    // 0x150: aes_key (0x10 bytes)
    // 0x160: export_table (u32)
    // 0x164: header_digest (0x14 bytes)
    // 0x178: region (u32)
    // 0x17C: allowed_media_types (u32)
    // 0x180: page_descriptor_count (u32)
    // 0x184: page_descriptors[] (each is 0x18 bytes: u32 value + 0x14 digest)

    let _header_size = cursor.read_u32::<BigEndian>()?;     // 0x000
    let image_size = cursor.read_u32::<BigEndian>()?;       // 0x004

    // Skip RSA signature (0x100 bytes)
    let mut rsa_sig = [0u8; 0x100];
    cursor.read_exact(&mut rsa_sig)?;                       // 0x008

    let _unk_108 = cursor.read_u32::<BigEndian>()?;         // 0x108
    let image_flags = cursor.read_u32::<BigEndian>()?;      // 0x10C
    let load_address = cursor.read_u32::<BigEndian>()?;     // 0x110

    // Skip section_digest (0x14 bytes)
    let mut digest = [0u8; 0x14];
    cursor.read_exact(&mut digest)?;                        // 0x114

    let _import_table_count = cursor.read_u32::<BigEndian>()?; // 0x128

    // Skip import_table_digest (0x14 bytes)
    cursor.read_exact(&mut digest)?;                        // 0x12C

    // Skip xgd2_media_id (0x10 bytes)
    let mut media_id = [0u8; 0x10];
    cursor.read_exact(&mut media_id)?;                      // 0x140

    // Read aes_key (0x10 bytes)
    let mut aes_key = [0u8; 0x10];
    cursor.read_exact(&mut aes_key)?;                       // 0x150

    let export_table_address = cursor.read_u32::<BigEndian>()?; // 0x160

    // Skip header_digest (0x14 bytes)
    cursor.read_exact(&mut digest)?;                        // 0x164

    let _region = cursor.read_u32::<BigEndian>()?;          // 0x178
    let _allowed_media = cursor.read_u32::<BigEndian>()?;   // 0x17C

    let page_descriptor_count = cursor.read_u32::<BigEndian>()?; // 0x180

    let mut page_descriptors = Vec::new();
    for _ in 0..page_descriptor_count {
        let size_and_info = cursor.read_u32::<BigEndian>()?;
        // Skip data_digest (0x14 bytes per descriptor)
        cursor.read_exact(&mut digest)?;
        page_descriptors.push(Xex2PageDescriptor { size_and_info });
    }

    Ok(Xex2SecurityInfo {
        image_size,
        load_address,
        export_table_address,
        image_flags,
        aes_key,
        page_descriptors,
    })
}

/// Parse file format info from the optional header data.
fn parse_file_format_info(data: &[u8], headers: &[Xex2OptionalHeader]) -> Option<FileFormatInfo> {
    // The key format: low 8 bits indicate the data size category
    // 0xFF = data offset is a pointer to variable-size data in the header area
    let header = headers.iter().find(|h| h.key == header_keys::FILE_FORMAT_INFO)?;
    let offset = header.value as usize;
    if offset + 8 > data.len() {
        return None;
    }

    let mut cursor = Cursor::new(data);
    cursor.seek(SeekFrom::Start(offset as u64)).ok()?;

    let info_size = cursor.read_u32::<BigEndian>().ok()?;
    let encryption_type = cursor.read_u16::<BigEndian>().ok()?;
    let compression_type = cursor.read_u16::<BigEndian>().ok()?;

    let mut basic_blocks = Vec::new();
    let mut normal_window_size = 0u32;
    let mut normal_first_block_size = 0u32;
    let mut normal_first_block_hash = [0u8; 20];

    match compression_type {
        COMPRESSION_BASIC => {
            // Basic compression blocks: (data_size, zero_size) pairs
            // Number of blocks = (info_size - 8) / 8
            let block_count = if info_size > 8 { (info_size - 8) / 8 } else { 0 };
            for _ in 0..block_count {
                let data_size = cursor.read_u32::<BigEndian>().ok()?;
                let zero_size = cursor.read_u32::<BigEndian>().ok()?;
                basic_blocks.push(BasicCompressionBlock { data_size, zero_size });
            }
        }
        COMPRESSION_NORMAL => {
            normal_window_size = cursor.read_u32::<BigEndian>().ok()?;
            // Read first_block: block_size (4) + block_hash (20)
            normal_first_block_size = cursor.read_u32::<BigEndian>().ok()?;
            cursor.read_exact(&mut normal_first_block_hash).ok()?;
        }
        _ => {}
    }

    Some(FileFormatInfo {
        info_size,
        encryption_type,
        compression_type,
        basic_blocks,
        normal_window_size,
        normal_first_block_size,
        normal_first_block_hash,
    })
}

/// Parse import libraries from the optional header data.
/// At this stage, only record addresses are read; ordinals and record types
/// are resolved later by `resolve_imports` once the PE image is decompressed.
fn parse_import_libraries(data: &[u8], headers: &[Xex2OptionalHeader]) -> Vec<ImportLibrary> {
    let header = match headers.iter().find(|h| h.key == header_keys::IMPORT_LIBRARIES) {
        Some(h) => h,
        None => return Vec::new(),
    };

    let offset = header.value as usize;
    if offset + 12 > data.len() {
        return Vec::new();
    }

    fn be_u32(data: &[u8], off: usize) -> u32 {
        u32::from_be_bytes([data[off], data[off+1], data[off+2], data[off+3]])
    }
    fn be_u16(data: &[u8], off: usize) -> u16 {
        u16::from_be_bytes([data[off], data[off+1]])
    }

    let total_size = be_u32(data, offset) as usize;
    let string_table_size = be_u32(data, offset + 4) as usize;
    let string_count = be_u32(data, offset + 8) as usize;

    // Parse string table (null-terminated, 4-byte aligned)
    let string_data_start = offset + 12;
    let mut strings = Vec::new();
    let mut spos = 0usize;
    for _ in 0..string_count {
        let start = string_data_start + spos;
        let mut end = start;
        while end < data.len() && data[end] != 0 { end += 1; }
        let name = std::str::from_utf8(&data[start..end]).unwrap_or("???").to_string();
        spos += name.len() + 1;
        // 4-byte alignment
        if !spos.is_multiple_of(4) { spos += 4 - (spos % 4); }
        strings.push(name);
    }

    // Parse libraries
    let mut libs = Vec::new();
    let mut lib_off = offset + 12 + string_table_size;

    while lib_off + 0x28 <= data.len() && lib_off < offset + total_size {
        let lib_size = be_u32(data, lib_off) as usize;
        if lib_size == 0 { break; }

        let id = be_u32(data, lib_off + 0x18);
        let version_cur = be_u32(data, lib_off + 0x1C);
        let version_min = be_u32(data, lib_off + 0x20);
        let name_index = (be_u16(data, lib_off + 0x24) & 0xFF) as usize;
        let count = be_u16(data, lib_off + 0x26) as usize;

        let lib_name = strings.get(name_index).cloned().unwrap_or_else(|| format!("lib_{name_index}"));

        let mut imports = Vec::new();
        for i in 0..count {
            let record_addr = be_u32(data, lib_off + 0x28 + i * 4);
            imports.push(ImportEntry {
                ordinal: 0,
                record_type: 0xFF,
                address: record_addr,
            });
        }

        libs.push(ImportLibrary {
            name: lib_name,
            id,
            version_min,
            version_cur,
            imports,
        });
        lib_off += lib_size;
    }

    libs
}

/// Resolve import ordinals and record types from the decompressed PE image.
/// Must be called after `load_image` provides the PE data.
pub fn resolve_imports(header: &mut Xex2Header, pe_image: &[u8]) {
    let image_base = get_image_base(header).unwrap_or(0);

    for lib in &mut header.import_libraries {
        for imp in &mut lib.imports {
            let pe_off = imp.address.wrapping_sub(image_base) as usize;
            if pe_off + 4 <= pe_image.len() {
                // PE image values are big-endian (Xbox 360 native)
                let val = u32::from_be_bytes([
                    pe_image[pe_off], pe_image[pe_off+1],
                    pe_image[pe_off+2], pe_image[pe_off+3],
                ]);
                imp.record_type = ((val >> 24) & 0xFF) as u8;
                imp.ordinal = (val & 0xFFFF) as u16;
            }
        }
    }
}

/// Parse execution info from optional header data.
fn parse_execution_info(data: &[u8], headers: &[Xex2OptionalHeader]) -> Option<ExecutionInfo> {
    // EXECUTION_INFO key is 0x00040006 — the low byte 0x06 means the value
    // is an inline struct of 6 u32 words (24 bytes total).
    // Layout: media_id(4), version(4), base_version(4), title_id(4),
    //         platform(1), exec_type(1), disc_number(1), disc_count(1)
    let header = headers.iter().find(|h| h.key == header_keys::EXECUTION_INFO)?;
    let off = header.value as usize;
    if off + 20 > data.len() {
        return None;
    }

    let media_id = u32::from_be_bytes([data[off], data[off+1], data[off+2], data[off+3]]);
    let title_id = u32::from_be_bytes([data[off+12], data[off+13], data[off+14], data[off+15]]);
    let disc_number = data[off + 18];
    let disc_count = data[off + 19];

    Some(ExecutionInfo {
        media_id,
        title_id,
        disc_number,
        disc_count,
    })
}

/// Parse original PE name from optional header data.
fn parse_original_pe_name(data: &[u8], headers: &[Xex2OptionalHeader]) -> Option<String> {
    let header = headers.iter().find(|h| h.key == header_keys::ORIGINAL_PE_NAME)?;
    let off = header.value as usize;
    if off + 4 > data.len() {
        return None;
    }

    let size = u32::from_be_bytes([data[off], data[off+1], data[off+2], data[off+3]]) as usize;
    if off + size > data.len() || size <= 4 {
        return None;
    }

    let name_bytes = &data[off + 4..off + size];
    Some(String::from_utf8_lossy(name_bytes).trim_end_matches('\0').to_string())
}

/// Get an optional header value by key.
pub fn get_opt_header(header: &Xex2Header, key: u32) -> Option<u32> {
    header.optional_headers.iter()
        .find(|h| h.key == key)
        .map(|h| h.value)
}

/// Get the entry point address from the XEX2 header.
pub fn get_entry_point(header: &Xex2Header) -> Option<u32> {
    get_opt_header(header, header_keys::ENTRY_POINT)
}

/// Get the image base address.
pub fn get_image_base(header: &Xex2Header) -> Option<u32> {
    get_opt_header(header, header_keys::IMAGE_BASE_ADDRESS)
}

/// Get the default stack size.
pub fn get_stack_size(header: &Xex2Header) -> u32 {
    get_opt_header(header, header_keys::DEFAULT_STACK_SIZE).unwrap_or(0x10_0000) // Default 1MB
}

/// XEX `XEX_HEADER_SYSTEM_FLAGS` (key `0x00030000`) — the privilege bitmap
/// queried by `XexCheckExecutablePrivilege`. Low byte 0x00 means the inline
/// `value` field is the u32 itself (canary `xex_module.cc:103-108`). Returns
/// 0 when the header is absent (matches canary's `GetOptHeader` zero-init).
pub fn get_system_flags(header: &Xex2Header) -> u32 {
    get_opt_header(header, header_keys::SYSTEM_FLAGS).unwrap_or(0)
}

/// Load the XEX image data into a flat buffer (decompressing if needed).
/// Returns the decompressed image bytes ready to map into guest memory.
#[tracing::instrument(skip_all, fields(bytes = data.len()))]
pub fn load_image(data: &[u8], header: &Xex2Header) -> io::Result<Vec<u8>> {
    let started = std::time::Instant::now();
    let source = &data[header.header_size as usize..];
    let bytes_in = source.len();

    let output = match &header.file_format_info {
        Some(info) if info.compression_type == COMPRESSION_BASIC => {
            tracing::debug!(compression = "basic", "decompressing");
            load_basic_compressed(source, info)?
        }
        Some(info) if info.compression_type == COMPRESSION_NORMAL => {
            tracing::debug!(compression = "normal/LZX", "decompressing");
            load_normal_compressed(source, info, header)?
        }
        _ => source.to_vec(),
    };

    let elapsed_ms = started.elapsed().as_millis() as f64;
    metrics::histogram!("xex.load_image_ms").record(elapsed_ms);
    metrics::counter!("xex.bytes_in").increment(bytes_in as u64);
    metrics::counter!("xex.bytes_out").increment(output.len() as u64);
    let ratio = if bytes_in == 0 { 0.0 } else { output.len() as f64 / bytes_in as f64 };
    tracing::info!(bytes_in, bytes_out = output.len(), ratio, elapsed_ms, "image loaded");
    Ok(output)
}

/// Load basic compressed image data.
fn load_basic_compressed(source: &[u8], info: &FileFormatInfo) -> io::Result<Vec<u8>> {
    // Calculate total uncompressed size
    let total_size: u64 = info.basic_blocks.iter()
        .map(|b| b.data_size as u64 + b.zero_size as u64)
        .sum();

    let mut output = vec![0u8; total_size as usize];
    let mut src_offset = 0usize;
    let mut dst_offset = 0usize;

    for block in &info.basic_blocks {
        let data_size = block.data_size as usize;
        let zero_size = block.zero_size as usize;

        if src_offset + data_size > source.len() {
            return Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                format!("Basic compression block data extends past end of file (src_offset={:#x}, data_size={:#x}, source_len={:#x})",
                    src_offset, data_size, source.len()),
            ));
        }

        // Copy data block
        if dst_offset + data_size <= output.len() {
            output[dst_offset..dst_offset + data_size]
                .copy_from_slice(&source[src_offset..src_offset + data_size]);
        }
        src_offset += data_size;
        dst_offset += data_size;

        // Zero-filled gap (already zeroed from vec initialization)
        dst_offset += zero_size;
    }

    Ok(output)
}

/// Xbox 360 retail AES key for XEX2 session key decryption.
const XEX2_RETAIL_KEY: [u8; 16] = [
    0x20, 0xB1, 0x85, 0xA5, 0x9D, 0x28, 0xFD, 0xC3,
    0x40, 0x58, 0x3F, 0xBB, 0x08, 0x96, 0xBF, 0x91,
];

/// Xbox 360 devkit AES key (all zeros).
#[allow(dead_code)]
const XEX2_DEVKIT_KEY: [u8; 16] = [0u8; 16];

/// AES-128-CBC decryption with zero IV (matching Xbox 360 XEX decryption).
#[tracing::instrument(skip_all, fields(bytes = input.len()))]
fn aes_decrypt_cbc(key: &[u8; 16], input: &[u8]) -> Vec<u8> {
    let cipher = Aes128::new(key.into());
    let mut output = vec![0u8; input.len()];
    let mut iv = [0u8; 16];

    for (i, chunk) in input.chunks(16).enumerate() {
        if chunk.len() < 16 {
            // Partial block at end - copy as-is
            output[i * 16..i * 16 + chunk.len()].copy_from_slice(chunk);
            break;
        }
        let mut block = aes::Block::clone_from_slice(chunk);
        cipher.decrypt_block(&mut block);
        // XOR with IV (previous ciphertext block)
        for j in 0..16 {
            block[j] ^= iv[j];
        }
        iv.copy_from_slice(chunk);
        output[i * 16..(i + 1) * 16].copy_from_slice(&block);
    }

    output
}

/// Derive the session key by decrypting the XEX's aes_key field with the retail key.
/// Falls back to devkit key if retail produces invalid results.
fn derive_session_key(header: &Xex2Header) -> [u8; 16] {
    let sec = match &header.security_info {
        Some(s) => s,
        None => return [0u8; 16],
    };

    let decrypted = aes_decrypt_cbc(&XEX2_RETAIL_KEY, &sec.aes_key);
    let mut session_key = [0u8; 16];
    session_key.copy_from_slice(&decrypted[..16]);
    session_key
}

/// De-block compressed data: strip block headers and extract chunk payloads.
///
/// The first block's size comes from the file format header (first_block_size).
/// Each block in the data starts with a block_info struct for the NEXT block:
///   - block_size: u32 BE (size of the next block)
///   - block_hash: [u8; 20] (SHA1 of the next block)
/// Followed by chunks: { chunk_size: u16 BE, data: [u8; chunk_size] }, terminated by chunk_size=0
fn deblock(input: &[u8], first_block_size: u32) -> io::Result<Vec<u8>> {
    let mut output = Vec::new();
    let mut pos = 0usize;
    let mut cur_block_size = first_block_size as usize;

    while cur_block_size > 0 && pos < input.len() {
        let next_block_pos = pos + cur_block_size;

        // Read next block's info from start of current block data
        let next_block_size = if pos + 4 <= input.len() {
            u32::from_be_bytes([
                input[pos], input[pos + 1], input[pos + 2], input[pos + 3],
            ]) as usize
        } else {
            0
        };

        // Skip block_info header (4 bytes size + 20 bytes hash)
        let mut p = pos + 4 + 20;

        // Read chunks within this block
        loop {
            if p + 2 > input.len() {
                break;
            }
            let chunk_size = ((input[p] as usize) << 8) | (input[p + 1] as usize);
            p += 2;
            if chunk_size == 0 {
                break;
            }
            if p + chunk_size > input.len() {
                return Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    format!("De-block chunk extends past input (pos={:#x}, chunk_size={:#x}, input_len={:#x})",
                        p, chunk_size, input.len()),
                ));
            }
            output.extend_from_slice(&input[p..p + chunk_size]);
            p += chunk_size;
        }

        if next_block_pos <= pos {
            break; // Prevent infinite loop
        }
        pos = next_block_pos;
        cur_block_size = next_block_size;
    }

    Ok(output)
}

/// Load normal (LZX) compressed image data.
/// Pipeline: decrypt → de-block → LZX decompress (pure Rust)
#[tracing::instrument(skip_all, fields(bytes_in = source.len()))]
fn load_normal_compressed(source: &[u8], info: &FileFormatInfo, header: &Xex2Header) -> io::Result<Vec<u8>> {
    let uncompressed_size = header.security_info.as_ref()
        .map(|s| s.image_size as usize)
        .unwrap_or(0);

    if uncompressed_size == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "Cannot decompress: image_size is 0",
        ));
    }

    // Step 1: Decrypt if needed
    let decrypted;
    let input = if info.encryption_type == ENCRYPTION_NORMAL {
        let session_key = derive_session_key(header);
        decrypted = aes_decrypt_cbc(&session_key, source);
        &decrypted
    } else {
        source
    };

    // Step 2: De-block (strip block headers, extract chunk payloads)
    let deblocked = deblock(input, info.normal_first_block_size)?;

    if deblocked.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "De-blocking produced no data",
        ));
    }

    // Step 3: LZX decompress using pure Rust decoder
    let window_bits = match info.normal_window_size {
        s if s == 0 => 15, // default
        s => (s as f64).log2() as u32,
    };

    let mut decoder = crate::lzx::LzxDecoder::new(window_bits);
    let output = decoder.decompress(&deblocked, uncompressed_size)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, format!("LZX decompression failed: {e}")))?;

    tracing::info!("LZX decompressed: {} -> {} bytes", deblocked.len(), uncompressed_size);

    Ok(output)
}