use crate::block_io::BlockIo;
use crate::cache::BlockCache;
use crate::error::FsError;
use crate::integrity::crc32c;
use lancer_core::fs::{BlockRef, DEDUP_SHARDS, SUPERBLOCK_SIZE, Superblock, SuperblockSlot};

fn superblock_to_bytes(sb: &Superblock) -> [u8; SUPERBLOCK_SIZE] {
    let mut buf = [0u8; SUPERBLOCK_SIZE];
    let src = sb as *const Superblock as *const u8;
    unsafe { core::ptr::copy_nonoverlapping(src, buf.as_mut_ptr(), SUPERBLOCK_SIZE) };
    buf
}

fn superblock_from_bytes(buf: &[u8; SUPERBLOCK_SIZE]) -> Superblock {
    unsafe { core::ptr::read_unaligned(buf.as_ptr() as *const Superblock) }
}

pub struct CommitRoots {
    pub tree_root: BlockRef,
    pub freemap_root: BlockRef,
    pub dedup_roots: [BlockRef; DEDUP_SHARDS],
    pub snapshot_root: BlockRef,
    pub scrub_cursor: u64,
    pub next_object_id: u64,
}

pub struct SuperblockPair {
    slots: [Superblock; 2],
}

impl SuperblockPair {
    pub fn load(a: Superblock, b: Superblock) -> Self {
        Self { slots: [a, b] }
    }

    pub fn active(&self) -> &Superblock {
        lancer_core::fs::select_superblock(&self.slots[0], &self.slots[1])
    }

    fn commit_target(&self) -> SuperblockSlot {
        lancer_core::fs::commit_target(&self.slots[0], &self.slots[1])
    }

    fn update_slot(&mut self, slot: SuperblockSlot, sb: Superblock) {
        match slot {
            SuperblockSlot::A => self.slots[0] = sb,
            SuperblockSlot::B => self.slots[1] = sb,
        }
    }
}

pub fn atomic_commit(
    bio: &mut BlockIo,
    cache: &mut BlockCache,
    pair: &mut SuperblockPair,
    roots: &CommitRoots,
) -> Result<(), FsError> {
    cache.cache_flush(bio)?;
    bio.flush()?;

    let current = pair.active();
    let next_seq = current.next_sequence().ok_or(FsError::CorruptSuperblock)?;
    let next_txn = current
        .next_transaction()
        .ok_or(FsError::CorruptSuperblock)?;

    let mut new_sb = Superblock::new(current.total_blocks, current.block_size);
    new_sb.sequence = next_seq;
    new_sb.transaction_id = next_txn;
    new_sb.tree_root = roots.tree_root;
    new_sb.freemap_root = roots.freemap_root;
    new_sb.dedup_roots = roots.dedup_roots;
    new_sb.snapshot_root = roots.snapshot_root;
    new_sb.scrub_cursor = roots.scrub_cursor;
    new_sb.next_object_id = roots.next_object_id;
    new_sb.checksum = 0;

    let mut buf = superblock_to_bytes(&new_sb);
    let checksum = crc32c(&buf[..SUPERBLOCK_SIZE - 4]);
    buf[SUPERBLOCK_SIZE - 4..].copy_from_slice(&checksum.to_le_bytes());

    let target = pair.commit_target();
    bio.write_blocks(target.block_number(), 1, &buf)?;
    bio.flush()?;

    let final_sb = superblock_from_bytes(&buf);
    pair.update_slot(target, final_sb);

    Ok(())
}

pub fn verify_superblock_crc(sb: &Superblock) -> bool {
    let buf = superblock_to_bytes(sb);
    let computed = crc32c(&buf[..SUPERBLOCK_SIZE - 4]);
    computed == sb.checksum
}

fn read_raw_superblock(bio: &mut BlockIo, slot: u64) -> Result<Superblock, FsError> {
    let mut buf = [0u8; SUPERBLOCK_SIZE];
    bio.read_blocks(slot, 1, &mut buf)?;
    Ok(superblock_from_bytes(&buf))
}

fn validate_one(sb: Superblock) -> Option<Superblock> {
    match sb.is_valid_magic() && verify_superblock_crc(&sb) {
        true => Some(sb),
        false => None,
    }
}

pub fn load_validated_superblock_pair(bio: &mut BlockIo) -> Result<SuperblockPair, FsError> {
    let sb_a = read_raw_superblock(bio, 0).ok().and_then(validate_one);
    let sb_b = read_raw_superblock(bio, 1).ok().and_then(validate_one);

    match (sb_a, sb_b) {
        (Some(a), Some(b)) => Ok(SuperblockPair::load(a, b)),
        (Some(a), None) => {
            let clone = superblock_from_bytes(&superblock_to_bytes(&a));
            Ok(SuperblockPair::load(a, clone))
        }
        (None, Some(b)) => {
            let clone = superblock_from_bytes(&superblock_to_bytes(&b));
            Ok(SuperblockPair::load(clone, b))
        }
        (None, None) => Err(FsError::CorruptSuperblock),
    }
}