pub mod helper;

use anyhow::{Result, anyhow};
use helper::ReadExt;
use log::{debug, info};
use std::{
    env::args,
    fs::File,
    io::{BufReader, Cursor, Read, Seek, SeekFrom, Write, copy},
};

fn main() -> anyhow::Result<()> {
    env_logger::init_from_env("LOG");
    let mut file = BufReader::new(File::open(args().nth(1).unwrap())?);

    let signature = file.read_cstr()?;
    let file_version = file.read_u32_be()?;
    let player_version = file.read_cstr()?;
    let unity_version = file.read_cstr()?;
    let size = file.read_u64_be()?;
    let blockindex_comp_size = file.read_u32_be()?;
    let blockindex_decomp_size = file.read_u32_be()?;
    let flags = file.read_u32_be()?;

    let meta_comp_scheme = CompressionScheme::from_flag_num((flags & 0x3f) as u8)
        .ok_or(anyhow!("unknown block compression"))?;
    let blockindex_eof = flags & 0x80 != 0;
    let blockindex_has_directory = flags & 0x40 != 0;
    let blockindex_need_padding = flags & 0x200 != 0;

    info!("signature={signature:?}");
    info!("file_version={file_version:?}");
    info!("player_version={player_version:?}");
    info!("unity_version={unity_version:?}");
    debug!("size={size:?}");
    debug!("meta_comp_size={blockindex_comp_size:?}");
    debug!("meta_decomp_size={blockindex_decomp_size:?}");
    debug!("flags={flags:?}");
    debug!("meta_comp_scheme={meta_comp_scheme:?}");
    debug!("blockindex_eof={blockindex_eof:?}");
    debug!("blockindex_has_directory={blockindex_has_directory:?}");
    debug!("blockindex_need_padding={blockindex_need_padding:?}");

    let mut blockindex = {
        let restore_position = if blockindex_eof {
            let pos = file.stream_position()?;
            file.seek(SeekFrom::End(-(blockindex_comp_size as i64)))?;
            Some(pos)
        } else {
            None
        };

        let mut blockindex = vec![0u8; blockindex_comp_size as usize];
        file.read_exact(&mut blockindex)?;

        if let Some(pos) = restore_position {
            file.seek(SeekFrom::Start(pos))?;
        }
        let blockindex =
            meta_comp_scheme.decompress(blockindex, blockindex_decomp_size as usize)?;
        Cursor::new(blockindex)
    };
    {
        // align stream
        let off = file.stream_position()? % 16;
        if off != 0 {
            file.seek_relative(16 - off as i64)?;
        }
    }

    blockindex.read_u128_be()?;

    let num_blocks = blockindex.read_u32_be()?;
    debug!("num_blocks={num_blocks:?}");
    let mut blocks = Vec::new();
    for _ in 0..num_blocks {
        let block_decomp_size = blockindex.read_u32_be()?;
        let block_comp_size = blockindex.read_u32_be()?;
        let block_flags = blockindex.read_u16_be()?;
        let block_comp_scheme = CompressionScheme::from_flag_num((block_flags & 0x3f) as u8)
            .ok_or(anyhow!("unknown block compression"))?;
        blocks.push((block_comp_size, block_decomp_size, block_comp_scheme))
    }

    let num_nodes = blockindex.read_u32_be()?;
    debug!("num_nodes={num_nodes:?}");
    let mut nodes = Vec::new();
    for _ in 0..num_nodes {
        let offset = blockindex.read_u64_be()?;
        let size = blockindex.read_u64_be()?;
        let status = blockindex.read_u32_be()?;
        let name = blockindex.read_cstr()?;
        info!("found node {name:?} (offset={offset}, size={size}, status={status})");
        nodes.push((offset, size, status, name))
    }

    let mut storage = Vec::new();
    for (comp_size, decomp_size, comp_scheme) in blocks {
        let mut comp_buf = vec![0u8; comp_size as usize];
        file.read_exact(&mut comp_buf)?;
        let decomp_buf = comp_scheme.decompress(comp_buf, decomp_size as usize)?;
        assert_eq!(decomp_size, decomp_buf.len() as u32);
        storage.extend_from_slice(&decomp_buf);
    }
    let mut storage = Cursor::new(storage);

    for (offset, size, _status, name) in nodes {
        storage.seek(SeekFrom::Start(offset))?;
        info!("extracting {name:?}");
        let mut file = File::create(format!("/tmp/{}", name))?;
        copy(&mut (&mut storage).take(size), &mut file)?;
    }

    Ok(())
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum CompressionScheme {
    None,
    LZMA,
    LZ4,
    LZ4HC,
    LZHAM,
}
impl CompressionScheme {
    pub fn from_flag_num(n: u8) -> Option<CompressionScheme> {
        Some(match n & 0x3f {
            0 => CompressionScheme::None,
            1 => CompressionScheme::LZMA,
            2 => CompressionScheme::LZ4,
            3 => CompressionScheme::LZ4HC,
            4 => CompressionScheme::LZHAM,
            _ => return None,
        })
    }
    pub fn decompress(&self, block: Vec<u8>, decomp_size: usize) -> Result<Vec<u8>> {
        match self {
            CompressionScheme::None => Ok(block),
            CompressionScheme::LZMA => {
                let mut r = lzma::Reader::from(Cursor::new(block))?;
                let mut buf = Vec::new();
                r.read_to_end(&mut buf)?;
                Ok(buf)
            }
            // CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
            //     Ok(lz4_flex::block::decompress(&block, decomp_size).context("lz4 decomp")?)
            // }
            CompressionScheme::LZ4HC | CompressionScheme::LZ4 => {
                File::create("/tmp/a")?.write_all(&block)?;
                Ok(lz4::block::decompress(&block, Some(decomp_size as i32))?)
                // Ok(lz4::block::decompress(&block, None)?)
            }
            CompressionScheme::LZHAM => todo!(),
        }
    }
}