use anyhow::Result;
use crc32fast::Hasher;
use fatfs::{Dir, FileSystem, FormatVolumeOptions, FsOptions, ReadWriteSeek};
use flate2::bufread::ZlibEncoder;
use flate2::read::ZlibDecoder;
use flate2::Compression;
use fscommon::BufStream;
use gpt::mbr::{PartRecord, ProtectiveMBR};
use std::fs::File;
use std::io::{BufReader, BufWriter, Cursor, Read, Seek, SeekFrom, Write};
use std::path::Path;

mod blkx;
mod koly;
mod xml;

pub use crate::blkx::*;
pub use crate::koly::*;
pub use crate::xml::*;

pub struct DmgReader<R: Read + Seek> {
    koly: KolyTrailer,
    xml: Plist,
    r: R,
}

impl DmgReader<BufReader<File>> {
    pub fn open(path: &Path) -> Result<Self> {
        let r = BufReader::new(File::open(path)?);
        Self::new(r)
    }
}

impl<R: Read + Seek> DmgReader<R> {
    pub fn new(mut r: R) -> Result<Self> {
        let koly = KolyTrailer::read_from(&mut r)?;
        r.seek(SeekFrom::Start(koly.plist_offset))?;
        let mut xml = Vec::with_capacity(koly.plist_length as usize);
        (&mut r).take(koly.plist_length).read_to_end(&mut xml)?;
        let xml: Plist = plist::from_reader_xml(&xml[..])?;
        Ok(Self { koly, xml, r })
    }

    pub fn koly(&self) -> &KolyTrailer {
        &self.koly
    }

    pub fn plist(&self) -> &Plist {
        &self.xml
    }

    pub fn sector(&mut self, chunk: &BlkxChunk) -> Result<impl Read + '_> {
        self.r.seek(SeekFrom::Start(chunk.compressed_offset))?;
        let compressed_chunk = (&mut self.r).take(chunk.compressed_length);
        match chunk.ty().expect("unknown chunk type") {
            ChunkType::Ignore | ChunkType::Zero | ChunkType::Comment => {
                Ok(Box::new(std::io::repeat(0).take(chunk.compressed_length)) as Box<dyn Read>)
            }
            ChunkType::Raw => Ok(Box::new(compressed_chunk)),
            ChunkType::Zlib => Ok(Box::new(ZlibDecoder::new(compressed_chunk))),
            ChunkType::Adc | ChunkType::Bzlib | ChunkType::Lzfse => unimplemented!(),
            ChunkType::Term => Ok(Box::new(std::io::empty())),
        }
    }

    pub fn data_checksum(&mut self) -> Result<u32> {
        self.r.seek(SeekFrom::Start(self.koly.data_fork_offset))?;
        let mut data_fork = Vec::with_capacity(self.koly.data_fork_length as usize);
        (&mut self.r)
            .take(self.koly.data_fork_length)
            .read_to_end(&mut data_fork)?;
        Ok(crc32fast::hash(&data_fork))
    }

    pub fn partition_table(&self, i: usize) -> Result<BlkxTable> {
        self.plist().partitions()[i].table()
    }

    pub fn partition_name(&self, i: usize) -> &str {
        &self.plist().partitions()[i].name
    }

    pub fn partition_data(&mut self, i: usize) -> Result<Vec<u8>> {
        let table = self.plist().partitions()[i].table()?;
        let mut partition = vec![];
        for chunk in &table.chunks {
            std::io::copy(&mut self.sector(chunk)?, &mut partition)?;
        }
        Ok(partition)
    }
}

pub struct DmgWriter<W: Write + Seek> {
    xml: Plist,
    w: W,
    data_hasher: Hasher,
    main_hasher: Hasher,
    sector_number: u64,
    compressed_offset: u64,
}

impl DmgWriter<BufWriter<File>> {
    pub fn create(path: &Path) -> Result<Self> {
        let w = BufWriter::new(File::create(path)?);
        Ok(Self::new(w))
    }
}

impl<W: Write + Seek> DmgWriter<W> {
    pub fn new(w: W) -> Self {
        Self {
            xml: Default::default(),
            w,
            data_hasher: Hasher::new(),
            main_hasher: Hasher::new(),
            sector_number: 0,
            compressed_offset: 0,
        }
    }

    pub fn create_fat32(mut self, fat32: &[u8]) -> Result<()> {
        anyhow::ensure!(fat32.len() % 512 == 0);
        let sector_count = fat32.len() as u64 / 512;
        let mut mbr = ProtectiveMBR::new();
        let mut partition = PartRecord::new_protective(Some(sector_count.try_into()?));
        partition.os_type = 11;
        mbr.set_partition(0, partition);
        let mbr = mbr.as_bytes()?;
        self.add_partition("Master Boot Record (MBR : 0)", &mbr)?;
        self.add_partition("FAT32 (FAT32 : 1)", fat32)?;
        self.finish()?;
        Ok(())
    }

    pub fn add_partition(&mut self, name: &str, bytes: &[u8]) -> Result<()> {
        anyhow::ensure!(bytes.len() % 512 == 0);
        let id = self.xml.partitions().len() as u32;
        let name = name.to_string();
        let mut table = BlkxTable::new(id, self.sector_number, crc32fast::hash(bytes));
        for chunk in bytes.chunks(2048 * 512) {
            let mut encoder = ZlibEncoder::new(chunk, Compression::best());
            let mut compressed = vec![];
            encoder.read_to_end(&mut compressed)?;
            let compressed_length = compressed.len() as u64;
            let sector_count = chunk.len() as u64 / 512;
            self.w.write_all(&compressed)?;
            self.data_hasher.update(&compressed);
            table.add_chunk(BlkxChunk::new(
                ChunkType::Zlib,
                self.sector_number,
                sector_count,
                self.compressed_offset,
                compressed_length,
            ));
            self.sector_number += sector_count;
            self.compressed_offset += compressed_length;
        }
        table.add_chunk(BlkxChunk::term(self.sector_number, self.compressed_offset));
        self.main_hasher.update(&table.checksum.data[..4]);
        self.xml
            .add_partition(Partition::new(id as i32 - 1, name, table));
        Ok(())
    }

    pub fn finish(mut self) -> Result<()> {
        let mut xml = vec![];
        plist::to_writer_xml(&mut xml, &self.xml)?;
        let pos = self.w.seek(SeekFrom::Current(0))?;
        let data_digest = self.data_hasher.finalize();
        let main_digest = self.main_hasher.finalize();
        let koly = KolyTrailer::new(
            pos,
            self.sector_number,
            pos,
            xml.len() as _,
            data_digest,
            main_digest,
        );
        self.w.write_all(&xml)?;
        koly.write_to(&mut self.w)?;
        Ok(())
    }
}

fn symlink(target: &str) -> Vec<u8> {
    let xsym = format!(
        "XSym\n{:04}\n{:x}\n{}\n",
        target.as_bytes().len(),
        md5::compute(target.as_bytes()),
        target,
    );
    let mut xsym = xsym.into_bytes();
    xsym.resize(1067, b' ');
    xsym
}

fn add_dir<T: ReadWriteSeek>(src: &Path, dest: &Dir<'_, T>) -> Result<()> {
    for entry in std::fs::read_dir(src)? {
        let entry = entry?;
        let file_name = entry.file_name();
        let file_name = file_name.to_str().unwrap();
        let source = src.join(&file_name);
        let file_type = entry.file_type()?;
        if file_type.is_dir() {
            let d = dest.create_dir(file_name)?;
            add_dir(&source, &d)?;
        } else if file_type.is_file() {
            let mut f = dest.create_file(file_name)?;
            std::io::copy(&mut File::open(source)?, &mut f)?;
        } else if file_type.is_symlink() {
            let target = std::fs::read_link(&source)?;
            let xsym = symlink(target.to_str().unwrap());
            let mut f = dest.create_file(file_name)?;
            std::io::copy(&mut &xsym[..], &mut f)?;
        }
    }
    Ok(())
}

pub fn create_dmg(dir: &Path, dmg: &Path, volume_label: &str, total_sectors: u32) -> Result<()> {
    let mut fat32 = vec![0; total_sectors as usize * 512];
    {
        let mut volume_label_bytes = [0; 11];
        let end = std::cmp::min(volume_label_bytes.len(), volume_label.len());
        volume_label_bytes[..end].copy_from_slice(&volume_label.as_bytes()[..end]);
        let volume_options = FormatVolumeOptions::new()
            .volume_label(volume_label_bytes)
            .bytes_per_sector(512)
            .total_sectors(total_sectors);
        let mut disk = BufStream::new(Cursor::new(&mut fat32));
        fatfs::format_volume(&mut disk, volume_options)?;
        let fs = FileSystem::new(disk, FsOptions::new())?;
        let file_name = dir.file_name().unwrap().to_str().unwrap();
        let dest = fs.root_dir().create_dir(file_name)?;
        add_dir(dir, &dest)?;
    }
    DmgWriter::create(dmg)?.create_fat32(&fat32)
}

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

    static DMG: &[u8] = include_bytes!("../assets/raqote-winit.dmg");

    fn print_dmg<R: Read + Seek>(dmg: &DmgReader<R>) -> Result<()> {
        println!("{:?}", dmg.koly());
        println!("{:?}", dmg.plist());
        for partition in dmg.plist().partitions() {
            let table = partition.table()?;
            println!("{:?}", table);
            println!("table checksum 0x{:x}", u32::from(table.checksum));
            for (i, chunk) in table.chunks.iter().enumerate() {
                println!("{} {:?}", i, chunk);
            }
        }
        Ok(())
    }

    #[test]
    fn read_koly_trailer() -> Result<()> {
        let koly = KolyTrailer::read_from(&mut Cursor::new(DMG))?;
        //println!("{:#?}", koly);
        let mut bytes = [0; 512];
        koly.write_to(&mut &mut bytes[..])?;
        let koly2 = KolyTrailer::read_from(&mut Cursor::new(&bytes))?;
        assert_eq!(koly, koly2);
        Ok(())
    }

    #[test]
    fn only_read_dmg() -> Result<()> {
        let mut dmg = DmgReader::new(Cursor::new(DMG))?;
        print_dmg(&dmg)?;
        assert_eq!(
            UdifChecksum::new(dmg.data_checksum()?),
            dmg.koly().data_fork_digest
        );
        let mut buffer = vec![];
        let mut dmg2 = DmgWriter::new(Cursor::new(&mut buffer));
        for i in 0..dmg.plist().partitions().len() {
            let data = dmg.partition_data(i)?;
            let name = dmg.partition_name(i);
            dmg2.add_partition(name, &data)?;
        }
        dmg2.finish()?;
        let mut dmg2 = DmgReader::new(Cursor::new(buffer))?;
        print_dmg(&dmg2)?;
        assert_eq!(
            UdifChecksum::new(dmg2.data_checksum()?),
            dmg2.koly().data_fork_digest
        );
        for i in 0..dmg.plist().partitions().len() {
            let table = dmg.partition_table(i)?;
            let data = dmg.partition_data(i)?;
            let expected = u32::from(table.checksum);
            let calculated = crc32fast::hash(&data);
            assert_eq!(expected, calculated);
        }
        assert_eq!(dmg.koly().main_digest, dmg2.koly().main_digest);
        println!("data crc32 0x{:x}", u32::from(dmg.koly().data_fork_digest));
        println!("main crc32 0x{:x}", u32::from(dmg.koly().main_digest));
        Ok(())
    }

    #[test]
    fn read_dmg_partition_mbr() -> Result<()> {
        let mut dmg = DmgReader::new(Cursor::new(DMG))?;
        let mbr = dmg.partition_data(0)?;
        println!("{:?}", mbr);
        let mbr = ProtectiveMBR::from_bytes(&mbr, LogicalBlockSize::Lb512)?;
        println!("{:?}", mbr);
        Ok(())
    }

    #[test]
    fn read_dmg_partition_fat32() -> Result<()> {
        let mut dmg = DmgReader::new(Cursor::new(DMG))?;
        let fat32 = dmg.partition_data(1)?;
        let fs = FileSystem::new(Cursor::new(fat32), FsOptions::new())?;
        println!("volume: {}", fs.volume_label());
        for entry in fs.root_dir().iter() {
            let entry = entry?;
            println!("{}", entry.file_name());
        }
        Ok(())
    }

    #[test]
    fn checksum() -> Result<()> {
        let mut dmg = DmgReader::new(Cursor::new(DMG))?;
        assert_eq!(
            UdifChecksum::new(dmg.data_checksum()?),
            dmg.koly().data_fork_digest
        );
        for i in 0..dmg.plist().partitions().len() {
            let table = dmg.partition_table(i)?;
            let data = dmg.partition_data(i)?;
            let expected = u32::from(table.checksum);
            let calculated = crc32fast::hash(&data);
            assert_eq!(expected, calculated);
        }
        Ok(())
    }
}