// SPDX-License-Identifier: MIT-0

//! On-disk file format for fcry.
//!
//! Layout:
//! ```text
//!   magic             "fcry"           4 bytes
//!   version           u8               1
//!   alg_id            u8               1
//!   flags             u8               1
//!   reserved          u8               1   (must be 0)
//!   chunk_size        u32 LE           4   (plaintext bytes per chunk)
//!   kdf_id            u8               1
//!   kdf_params        variable             (depends on kdf_id)
//!   nonce_prefix      [u8; 19]         19  (STREAM nonce prefix)
//!   plaintext_length  u64 LE           8   (only if version >= 2 and flags & 0x01)
//!   key_commitment    [u8; 32]         32  (only if version >= 3 and flags & 0x02)
//!   --- end of header ---
//!   chunk[0..N]                            each chunk_size + 16 bytes,
//!                                          last may be shorter
//! ```
//!
//! The full encoded header is fed as AAD to every chunk, so any tampering
//! with chunk_size, nonce_prefix, kdf params, plaintext_length, etc. causes
//! authentication failure on every chunk.
//!
//! Versions:
//! * v1 — no length committed, no flag bits used.
//! * v2 — adds `FLAG_LENGTH_COMMITTED` (bit 0); when set, the total plaintext
//!   length is appended after `nonce_prefix`. This enables random-access
//!   decryption without scanning predecessors.
//! * v3 — adds `FLAG_KEY_COMMITTED` (bit 1) and an authenticated key
//!   commitment for fast wrong-key detection before chunk processing.

use std::io::Read;

use crate::{error::FcryError, policy};

const MAGIC: [u8; 4] = *b"fcry";
pub const VERSION_CURRENT: u8 = 3;
const VERSION_MIN: u8 = 1;

pub const NONCE_PREFIX_LEN: usize = 19;
pub const TAG_LEN: usize = 16;

/// Set in `flags` when the header carries an authenticated `plaintext_length`
/// field. Required for random-access decryption.
pub const FLAG_LENGTH_COMMITTED: u8 = 0x01;
pub const FLAG_KEY_COMMITTED: u8 = 0x02;

/// Mask of all flag bits this build understands. Unknown bits → reject.
const FLAG_KNOWN_MASK: u8 = FLAG_LENGTH_COMMITTED | FLAG_KEY_COMMITTED;
pub const KEY_COMMITMENT_LEN: usize = 32;

#[repr(u8)]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum AlgId {
    XChaCha20Poly1305 = 1,
}

impl AlgId {
    fn from_u8(v: u8) -> Result<Self, FcryError> {
        match v {
            1 => Ok(Self::XChaCha20Poly1305),
            _ => Err(FcryError::Format(format!("unknown alg id: {v}"))),
        }
    }
}

pub const ARGON2_SALT_LEN: usize = 16;

/// Key-derivation parameters stored in the header.
///
/// `Raw` means the key was supplied directly (no KDF). `Argon2id` carries
/// the salt and cost parameters needed to redo derivation on decrypt.
#[derive(Clone, Debug)]
pub enum KdfParams {
    Raw,
    Argon2id {
        salt: [u8; ARGON2_SALT_LEN],
        m_cost: u32,
        t_cost: u32,
        p_cost: u32,
    },
}

impl KdfParams {
    fn id(&self) -> u8 {
        match self {
            Self::Raw => 0,
            Self::Argon2id { .. } => 1,
        }
    }

    fn write_into(&self, out: &mut Vec<u8>) {
        match self {
            Self::Raw => {}
            Self::Argon2id {
                salt,
                m_cost,
                t_cost,
                p_cost,
            } => {
                out.extend_from_slice(salt);
                out.extend_from_slice(&m_cost.to_le_bytes());
                out.extend_from_slice(&t_cost.to_le_bytes());
                out.extend_from_slice(&p_cost.to_le_bytes());
            }
        }
    }

    fn read_from(id: u8, r: &mut impl Read) -> Result<Self, FcryError> {
        match id {
            0 => Ok(Self::Raw),
            1 => {
                let mut salt = [0u8; ARGON2_SALT_LEN];
                r.read_exact(&mut salt)?;
                let mut buf = [0u8; 4];
                r.read_exact(&mut buf)?;
                let m_cost = u32::from_le_bytes(buf);
                r.read_exact(&mut buf)?;
                let t_cost = u32::from_le_bytes(buf);
                r.read_exact(&mut buf)?;
                let p_cost = u32::from_le_bytes(buf);
                Ok(Self::Argon2id {
                    salt,
                    m_cost,
                    t_cost,
                    p_cost,
                })
            }
            _ => Err(FcryError::Format(format!("unknown kdf id: {id}"))),
        }
    }
}

#[derive(Clone, Debug)]
pub struct Header {
    /// On-disk format version. Set to `VERSION_CURRENT` for new encrypts;
    /// preserved as-read for decrypt so the AAD recomputed on decode matches
    /// the bytes that were authenticated when the file was written.
    pub version: u8,
    pub alg: AlgId,
    pub flags: u8,
    pub chunk_size: u32,
    pub kdf: KdfParams,
    pub nonce_prefix: [u8; NONCE_PREFIX_LEN],
    /// Total plaintext byte count. `Some` iff `flags & FLAG_LENGTH_COMMITTED`.
    pub plaintext_length: Option<u64>,
    /// v3 key commitment. `Some` iff `flags & FLAG_KEY_COMMITTED`.
    pub key_commitment: Option<[u8; KEY_COMMITMENT_LEN]>,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct HeaderReadOptions {
    pub max_argon_memory_mib: u32,
}

impl Default for HeaderReadOptions {
    fn default() -> Self {
        Self {
            max_argon_memory_mib: policy::default_argon_decrypt_cap_mib(),
        }
    }
}

impl Header {
    fn encode_without_commitment(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(104);
        out.extend_from_slice(&MAGIC);
        out.push(self.version);
        out.push(self.alg as u8);
        out.push(self.flags);
        out.push(0); // reserved
        out.extend_from_slice(&self.chunk_size.to_le_bytes());
        out.push(self.kdf.id());
        self.kdf.write_into(&mut out);
        out.extend_from_slice(&self.nonce_prefix);
        if (self.flags & FLAG_LENGTH_COMMITTED) != 0 {
            let len = self
                .plaintext_length
                .expect("FLAG_LENGTH_COMMITTED set but plaintext_length is None");
            out.extend_from_slice(&len.to_le_bytes());
        }
        out
    }

    pub fn encode(&self) -> Vec<u8> {
        let mut out = self.encode_without_commitment();
        if (self.flags & FLAG_KEY_COMMITTED) != 0 {
            let commitment = self
                .key_commitment
                .expect("FLAG_KEY_COMMITTED set but key_commitment is None");
            out.extend_from_slice(&commitment);
        }
        out
    }

    pub fn commitment_input_encoding(&self) -> Vec<u8> {
        self.encode_without_commitment()
    }

    pub fn read(r: &mut impl Read) -> Result<Self, FcryError> {
        Self::read_with_options(r, HeaderReadOptions::default())
    }

    pub fn read_with_options(
        r: &mut impl Read,
        options: HeaderReadOptions,
    ) -> Result<Self, FcryError> {
        let mut magic = [0u8; 4];
        r.read_exact(&mut magic)?;
        if magic != MAGIC {
            return Err(FcryError::Format("not an fcry file (bad magic)".into()));
        }

        let mut fixed = [0u8; 4];
        r.read_exact(&mut fixed)?;
        let [version, alg_id, flags, reserved] = fixed;
        if !(VERSION_MIN..=VERSION_CURRENT).contains(&version) {
            return Err(FcryError::Format(format!("unsupported version: {version}")));
        }
        if reserved != 0 {
            return Err(FcryError::Format("reserved byte must be zero".into()));
        }
        if (flags & !FLAG_KNOWN_MASK) != 0 {
            return Err(FcryError::Format(format!(
                "unknown flag bits: 0x{flags:02x}"
            )));
        }
        if version < 2 && flags != 0 {
            return Err(FcryError::Format("v1 header must have flags == 0".into()));
        }
        if version < 3 && (flags & FLAG_KEY_COMMITTED) != 0 {
            return Err(FcryError::Format(
                "key commitment flag requires v3 header".into(),
            ));
        }
        if version >= 3 && (flags & FLAG_KEY_COMMITTED) == 0 {
            return Err(FcryError::Format("v3 header must commit the key".into()));
        }
        let alg = AlgId::from_u8(alg_id)?;

        let mut chunk_size_bytes = [0u8; 4];
        r.read_exact(&mut chunk_size_bytes)?;
        let chunk_size = u32::from_le_bytes(chunk_size_bytes);
        policy::validate_chunk_size(chunk_size)?;

        let mut kdf_id = [0u8; 1];
        r.read_exact(&mut kdf_id)?;
        let kdf = KdfParams::read_from(kdf_id[0], r)?;
        policy::validate_header_kdf(&kdf, options.max_argon_memory_mib)?;

        let mut nonce_prefix = [0u8; NONCE_PREFIX_LEN];
        r.read_exact(&mut nonce_prefix)?;

        let plaintext_length = if (flags & FLAG_LENGTH_COMMITTED) != 0 {
            let mut b = [0u8; 8];
            r.read_exact(&mut b)?;
            Some(u64::from_le_bytes(b))
        } else {
            None
        };

        let key_commitment = if (flags & FLAG_KEY_COMMITTED) != 0 {
            let mut b = [0u8; KEY_COMMITMENT_LEN];
            r.read_exact(&mut b)?;
            Some(b)
        } else {
            None
        };

        Ok(Self {
            version,
            alg,
            flags,
            chunk_size,
            kdf,
            nonce_prefix,
            plaintext_length,
            key_commitment,
        })
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;

    #[test]
    fn roundtrip() {
        let h = Header {
            version: VERSION_CURRENT,
            alg: AlgId::XChaCha20Poly1305,
            flags: FLAG_KEY_COMMITTED,
            chunk_size: 1024 * 1024,
            kdf: KdfParams::Raw,
            nonce_prefix: [7u8; NONCE_PREFIX_LEN],
            plaintext_length: None,
            key_commitment: Some([1u8; KEY_COMMITMENT_LEN]),
        };
        let bytes = h.encode();
        let mut cur = Cursor::new(&bytes);
        let parsed = Header::read(&mut cur).unwrap();
        assert_eq!(parsed.version, h.version);
        assert_eq!(parsed.alg, h.alg);
        assert_eq!(parsed.flags, h.flags);
        assert_eq!(parsed.chunk_size, h.chunk_size);
        assert_eq!(parsed.nonce_prefix, h.nonce_prefix);
        assert_eq!(parsed.plaintext_length, None);
        assert_eq!(parsed.key_commitment, h.key_commitment);
        assert_eq!(cur.position() as usize, bytes.len());
    }

    #[test]
    fn roundtrip_length_committed() {
        let h = Header {
            version: VERSION_CURRENT,
            alg: AlgId::XChaCha20Poly1305,
            flags: FLAG_LENGTH_COMMITTED | FLAG_KEY_COMMITTED,
            chunk_size: 65536,
            kdf: KdfParams::Raw,
            nonce_prefix: [9u8; NONCE_PREFIX_LEN],
            plaintext_length: Some(123_456_789),
            key_commitment: Some([2u8; KEY_COMMITMENT_LEN]),
        };
        let bytes = h.encode();
        let mut cur = Cursor::new(&bytes);
        let parsed = Header::read(&mut cur).unwrap();
        assert_eq!(parsed.flags, FLAG_LENGTH_COMMITTED | FLAG_KEY_COMMITTED);
        assert_eq!(parsed.plaintext_length, Some(123_456_789));
        assert_eq!(parsed.key_commitment, h.key_commitment);
        assert_eq!(cur.position() as usize, bytes.len());
    }

    #[test]
    fn v3_encoding_layout_stable() {
        let h = Header {
            version: VERSION_CURRENT,
            alg: AlgId::XChaCha20Poly1305,
            flags: FLAG_LENGTH_COMMITTED | FLAG_KEY_COMMITTED,
            chunk_size: 0x0102_0304,
            kdf: KdfParams::Raw,
            nonce_prefix: [0x55u8; NONCE_PREFIX_LEN],
            plaintext_length: Some(0x0807_0605_0403_0201),
            key_commitment: Some([0xaau8; KEY_COMMITMENT_LEN]),
        };
        let commitment_input = h.commitment_input_encoding();
        assert_eq!(commitment_input.len(), 40);
        assert_eq!(&commitment_input[..4], b"fcry");
        assert_eq!(commitment_input[4], 3);
        assert_eq!(
            &commitment_input[32..40],
            &0x0807_0605_0403_0201u64.to_le_bytes()
        );

        let aad = h.encode();
        assert_eq!(aad.len(), 72);
        assert_eq!(&aad[..40], &commitment_input);
        assert_eq!(&aad[40..], &[0xaau8; KEY_COMMITMENT_LEN]);
    }

    #[test]
    fn rejects_bad_magic() {
        let mut bytes = Header {
            version: VERSION_CURRENT,
            alg: AlgId::XChaCha20Poly1305,
            flags: 0,
            chunk_size: 4096,
            kdf: KdfParams::Raw,
            nonce_prefix: [0u8; NONCE_PREFIX_LEN],
            plaintext_length: None,
            key_commitment: Some([3u8; KEY_COMMITMENT_LEN]),
        }
        .encode();
        bytes[0] ^= 1;
        assert!(matches!(
            Header::read(&mut Cursor::new(&bytes)),
            Err(FcryError::Format(_))
        ));
    }

    #[test]
    fn rejects_unknown_flag_bits() {
        let mut bytes = Header {
            version: VERSION_CURRENT,
            alg: AlgId::XChaCha20Poly1305,
            flags: 0,
            chunk_size: 4096,
            kdf: KdfParams::Raw,
            nonce_prefix: [0u8; NONCE_PREFIX_LEN],
            plaintext_length: None,
            key_commitment: Some([4u8; KEY_COMMITMENT_LEN]),
        }
        .encode();
        // flags byte is at offset 6 (4 magic + version + alg)
        bytes[6] = 0x80;
        assert!(matches!(
            Header::read(&mut Cursor::new(&bytes)),
            Err(FcryError::Format(_))
        ));
    }

    #[test]
    fn reads_v1_header() {
        // hand-crafted v1 header (raw kdf, no length field)
        let mut bytes = Vec::new();
        bytes.extend_from_slice(b"fcry");
        bytes.push(1); // version
        bytes.push(1); // alg
        bytes.push(0); // flags
        bytes.push(0); // reserved
        bytes.extend_from_slice(&1024u32.to_le_bytes());
        bytes.push(0); // kdf id raw
        bytes.extend_from_slice(&[3u8; NONCE_PREFIX_LEN]);
        let parsed = Header::read(&mut Cursor::new(&bytes)).unwrap();
        assert_eq!(parsed.version, 1);
        assert_eq!(parsed.flags, 0);
        assert_eq!(parsed.chunk_size, 1024);
        assert_eq!(parsed.plaintext_length, None);
        assert_eq!(parsed.key_commitment, None);
        // Re-encoding must reproduce the original v1 bytes exactly so the
        // recomputed AAD matches what the file was authenticated with.
        assert_eq!(parsed.encode(), bytes);
    }
}