#![forbid(unsafe_code)]

use camino::Utf8Path;
use ciborium::{de, ser};
use core::str::FromStr;
use quantum_sign::crypto as qs_crypto;
use quantum_sign::crypto::DigestAlg;
use quantum_sign::format::{QSig, Transparency};
use quantum_sign::policy::{Policy, RequiredSignatures};
use serde::{Deserialize, Serialize};
use serde_bytes::{ByteBuf, Bytes};
use std::{
    collections::BTreeMap,
    fs::{self, File, OpenOptions},
    io::{self, ErrorKind, Write},
    time::{SystemTime, UNIX_EPOCH},
};

const VERSION: u8 = 1;
const MAX_FRAGMENT_LEN: usize = 8192;
const MAX_SIGNERS: usize = 16;

#[inline]
fn err_invalid(msg: &str) -> io::Error {
    io::Error::new(ErrorKind::InvalidData, msg)
}

#[inline]
fn err_input(msg: &str) -> io::Error {
    io::Error::new(ErrorKind::InvalidInput, msg)
}

#[inline]
fn err_perm(msg: &str) -> io::Error {
    io::Error::new(ErrorKind::PermissionDenied, msg)
}

#[inline]
fn err_other(msg: &str) -> io::Error {
    io::Error::other(msg.to_string())
}

fn intent_default_digest_alg() -> String {
    "sha512".into()
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Intent {
    pub v: u8,
    pub alg: String,
    pub digest: ByteBuf,
    #[serde(default = "intent_default_digest_alg")]
    pub digest_alg: String,
    pub policy_hash: ByteBuf,
    pub m: u8,
    pub n: u8,
    pub allowed_kids: Vec<String>,
    pub created_unix: i64,
    #[serde(default)]
    pub display_name: Option<String>,
    #[serde(default)]
    pub owner: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CosignFragment {
    pub v: u8,
    pub kid: String,
    pub alg: String,
    pub digest: ByteBuf,
    pub sig: ByteBuf,
    pub created_unix: i64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Journal {
    pub v: u8,
    pub intent: Intent,
    pub rev: u64,
    pub fragments: BTreeMap<String, ByteBuf>,
}

fn now_unix() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64
}

pub(crate) fn write_atomic_sync(path: &Utf8Path, bytes: &[u8]) -> io::Result<()> {
    let mut opts = OpenOptions::new();
    opts.write(true).create_new(true);
    #[cfg(unix)]
    {
        use libc::O_NOFOLLOW;
        use std::os::unix::fs::OpenOptionsExt;
        opts.mode(0o600).custom_flags(O_NOFOLLOW);
    }
    let mut attempt = 0usize;
    let mut tmp;
    let mut file;
    loop {
        let suffix = if attempt == 0 {
            "tmp".to_string()
        } else {
            format!("tmp{attempt}")
        };
        tmp = path.with_extension(suffix);
        match opts.open(&tmp) {
            Ok(f) => {
                file = f;
                break;
            }
            Err(err) if err.kind() == io::ErrorKind::AlreadyExists => {
                attempt = attempt.saturating_add(1);
                continue;
            }
            Err(err) => return Err(err),
        }
    }
    file.write_all(bytes)?;
    file.sync_all()?;
    fs::rename(&tmp, path)?;
    if let Some(parent) = path.parent() {
        let dir = File::open(parent)?;
        dir.sync_all()?;
    }
    Ok(())
}

fn load<T: for<'de> serde::Deserialize<'de>>(path: &Utf8Path) -> io::Result<T> {
    let bytes = fs::read(path)?;
    de::from_reader(bytes.as_slice()).map_err(|e| err_invalid(&format!("{e}")))
}

fn save<T: serde::Serialize>(path: &Utf8Path, value: &T) -> io::Result<()> {
    let mut buf = Vec::new();
    ser::into_writer(value, &mut buf).map_err(|e| err_other(&format!("{e}")))?;
    write_atomic_sync(path, &buf)
}

fn to_cbor<T: serde::Serialize>(value: &T) -> io::Result<Vec<u8>> {
    let mut buf = Vec::new();
    ser::into_writer(value, &mut buf).map_err(|e| err_other(&format!("{e}")))?;
    Ok(buf)
}

fn encode_claims(
    policy_hash: &[u8],
    kids: &[String],
    m: u8,
    n: u8,
    digest_alg: &str,
    display_name: Option<&str>,
    owner: Option<&str>,
) -> io::Result<Vec<u8>> {
    let hash = Bytes::new(policy_hash);
    let kid_refs: Vec<&str> = kids.iter().map(|s| s.as_str()).collect();
    let tuple = (
        &hash,
        &kid_refs,
        m as u64,
        n as u64,
        digest_alg,
        display_name.unwrap_or(""),
        owner.unwrap_or(""),
    );
    let mut buf = Vec::new();
    ser::into_writer(&tuple, &mut buf).map_err(|e| err_other(&format!("{e}")))?;
    Ok(buf)
}

pub(crate) fn validate_required_signatures(req: &RequiredSignatures) -> io::Result<()> {
    if req.m == 0 || req.n == 0 || req.m > req.n {
        return Err(err_input("invalid required_signatures in policy"));
    }
    Ok(())
}

/// Load an intent from disk.
pub fn load_intent(path: &Utf8Path) -> io::Result<Intent> {
    load(path)
}

/// Persist a cosign fragment to disk.
pub fn save_fragment(path: &Utf8Path, fragment: &CosignFragment) -> io::Result<()> {
    save(path, fragment)
}

/// Create an intent and journal for quorum signing.
pub fn quorum_init(
    intent_path: &Utf8Path,
    policy: &Policy,
    alg: &str,
    digest: &[u8],
    digest_alg: &str,
    allowed_kids: &[String],
) -> io::Result<()> {
    let (m, n) = match policy.required_signatures {
        Some(req) => {
            validate_required_signatures(&req)?;
            (req.m, req.n)
        }
        None => (1, 1),
    };

    let digest_alg_enum =
        DigestAlg::from_str(digest_alg).map_err(|_| err_input("unsupported digest_alg"))?;
    if digest.len() != digest_alg_enum.output_len() {
        return Err(err_input("digest length does not match digest_alg"));
    }

    if policy.require_level5 && !policy.allow_lower_levels {
        if !qs_crypto::is_level5_sig_alg(alg) {
            return Err(err_input("algorithm not permitted under Level-5 policy"));
        }
        if !matches!(digest_alg_enum, DigestAlg::Sha512 | DigestAlg::Shake256_64) {
            return Err(err_input(
                "digest_alg must be sha512 or shake256-64 under Level-5 policy",
            ));
        }
    }

    let mut allowed: Vec<String> = allowed_kids.to_vec();
    allowed.sort();
    allowed.dedup();

    if !allowed.is_empty() {
        if allowed.len() < m as usize {
            return Err(err_input("not enough allowed kids to satisfy quorum"));
        }
        if allowed.len() > n as usize {
            return Err(err_input("allowed kid list exceeds policy quorum"));
        }
    }

    let intent = Intent {
        v: VERSION,
        alg: alg.to_string(),
        digest: ByteBuf::from(digest.to_vec()),
        digest_alg: digest_alg.to_string(),
        policy_hash: ByteBuf::from(policy.canonical_hash().to_vec()),
        m,
        n,
        allowed_kids: allowed,
        created_unix: now_unix(),
    };

    let journal = Journal {
        v: VERSION,
        intent: intent.clone(),
        rev: 0,
        fragments: BTreeMap::new(),
    };

    let part_path = intent_path.with_extension("qsig.part");
    save(intent_path, &intent)?;
    save(&part_path, &journal)
}

/// Append a cosign fragment to the journal (merge-retry to avoid concurrent stomps).
pub fn quorum_add(part_path: &Utf8Path, fragment_path: &Utf8Path) -> io::Result<()> {
    let fragment: CosignFragment = load(fragment_path)?;

    loop {
        let mut journal: Journal = load(part_path)?;
        if journal.v != VERSION {
            return Err(err_invalid("journal version mismatch"));
        }
        if fragment.v != VERSION {
            return Err(err_invalid("fragment version mismatch"));
        }
        if fragment.alg != journal.intent.alg {
            return Err(err_invalid("fragment algorithm mismatch"));
        }
        if fragment.digest.as_ref() != journal.intent.digest.as_ref() {
            return Err(err_invalid("fragment digest mismatch"));
        }
        if !journal.intent.allowed_kids.is_empty()
            && !journal.intent.allowed_kids.contains(&fragment.kid)
        {
            return Err(err_perm("kid not permitted by intent"));
        }

        let existed = journal.fragments.contains_key(&fragment.kid);
        if !existed && journal.fragments.len() >= journal.intent.n as usize {
            return Err(err_input("quorum already satisfied"));
        }

        journal
            .fragments
            .insert(fragment.kid.clone(), fragment.sig.clone());
        journal.rev = journal.rev.saturating_add(1);

        let bytes = to_cbor(&journal)?;
        write_atomic_sync(part_path, &bytes)?;

        let journal_after: Journal = load(part_path)?;
        if journal_after.fragments.contains_key(&fragment.kid) {
            break;
        }
    }

    Ok(())
}

/// Seal the journal into a final `.qsig`, verifying each fragment with the supplied callback.
pub fn quorum_seal(
    part_path: &Utf8Path,
    out_sig: &Utf8Path,
    verify_sig: impl Fn(&str, &str, &[u8], DigestAlg, &[u8], &[u8]) -> bool,
) -> io::Result<()> {
    let journal: Journal = load(part_path)?;
    if journal.fragments.len() < journal.intent.m as usize {
        return Err(err_other("quorum not satisfied"));
    }

    let digest_alg = DigestAlg::from_str(&journal.intent.digest_alg)
        .map_err(|_| err_invalid("unsupported digest_alg"))?;
    if journal.intent.digest.len() != digest_alg.output_len() {
        return Err(err_invalid("digest length does not match digest_alg"));
    }
    if !qs_crypto::is_level5_sig_alg(&journal.intent.alg) {
        return Err(err_invalid("intent algorithm is not Level-5"));
    }

    if journal.intent.n as usize > MAX_SIGNERS {
        return Err(err_invalid(
            "intent.n exceeds maximum supported signers (16)",
        ));
    }

    let mut entries: Vec<_> = journal.fragments.iter().collect();
    entries.sort_unstable_by(|(ka, _), (kb, _)| ka.cmp(kb));

    if entries.len() > journal.intent.n as usize {
        return Err(err_invalid("more fragments than n"));
    }

    for (kid, sig) in &entries {
        if sig.len() > MAX_FRAGMENT_LEN {
            return Err(err_invalid(&format!("signature too large for {kid}")));
        }
        if journal.intent.alg == "mldsa-87" && sig.len() != qs_crypto::mldsa87::SIGNATURE_LEN {
            return Err(err_invalid(&format!(
                "signature length mismatch for mldsa-87 (expected {} bytes)",
                qs_crypto::mldsa87::SIGNATURE_LEN
            )));
        }
        if !verify_sig(
            kid,
            &journal.intent.alg,
            journal.intent.digest.as_ref(),
            digest_alg,
            sig.as_ref(),
            journal.intent.policy_hash.as_ref(),
        ) {
            return Err(err_invalid(&format!(
                "signature verification failed for {kid}"
            )));
        }
    }

    let mut sigs_by_kid = BTreeMap::new();
    for (kid, sig) in &entries {
        sigs_by_kid.insert((*kid).clone(), sig.as_ref().to_vec());
    }

    let kids_sorted: Vec<String> = entries.iter().map(|(kid, _)| (*kid).clone()).collect();
    if kids_sorted
        .windows(2)
        .any(|w| w[0].as_str() >= w[1].as_str())
    {
        return Err(err_invalid("kids list must be strictly sorted and unique"));
    }
    if kids_sorted.len() < journal.intent.m as usize {
        return Err(err_invalid("kids length shorter than intent.m"));
    }
    if kids_sorted.len() > journal.intent.n as usize {
        return Err(err_invalid("kids length exceeds intent.n"));
    }
    let claims_bytes = encode_claims(
        journal.intent.policy_hash.as_ref(),
        &kids_sorted,
        journal.intent.m,
        journal.intent.n,
        &journal.intent.digest_alg,
        journal.intent.display_name.as_deref(),
        journal.intent.owner.as_deref(),
    )?;

    let mut ordered_sigs = Vec::with_capacity(kids_sorted.len());
    for kid in &kids_sorted {
        if let Some(sig) = sigs_by_kid.get(kid) {
            ordered_sigs.push(sig.clone());
        }
    }

    if ordered_sigs.is_empty() {
        return Err(err_invalid("no signatures available after verification"));
    }

    let mut sig_iter = ordered_sigs.into_iter();
    let primary_sig = sig_iter.next().expect("ordered_sigs non-empty");
    let co = sig_iter.map(ByteBuf::from).collect();

    let qsig = QSig {
        version: VERSION,
        alg: journal.intent.alg.clone(),
        digest: journal.intent.digest.as_ref().to_vec(),
        sig: primary_sig,
        co_sig: co,
        tst: None,
        tlog: None::<Transparency>,
        signer: qs_format::Signer {
            kid: format!("quorum:{}", kids_sorted[0]),
            claims: claims_bytes,
        },
        time_unix: now_unix(),
        domain_sep: "quantum-sign-v1".into(),
        meta: Some(journal.intent.policy_hash.clone()),
    };

    let mut encoded = Vec::new();
    ser::into_writer(&qsig, &mut encoded).map_err(|e| err_other(&format!("{e}")))?;

    let roundtrip: QSig = de::from_reader(encoded.as_slice())
        .map_err(|e| err_invalid(&format!("qsig decode: {e}")))?;
    let mut reencoded = Vec::new();
    ser::into_writer(&roundtrip, &mut reencoded).map_err(|e| err_other(&format!("{e}")))?;
    if reencoded != encoded {
        return Err(err_invalid("non-canonical .qsig (re-encode mismatch)"));
    }

    write_atomic_sync(out_sig, &encoded)
}

/// Construct a fragment from an `Intent` and raw signature bytes.
pub fn make_fragment(intent: &Intent, kid: &str, signature: Vec<u8>) -> CosignFragment {
    CosignFragment {
        v: VERSION,
        kid: kid.to_string(),
        alg: intent.alg.clone(),
        digest: intent.digest.clone(),
        sig: ByteBuf::from(signature),
        created_unix: now_unix(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use camino::Utf8PathBuf;
    use qs_crypto::{self, keypair_mldsa87, sign_mldsa87, DigestAlg, HmacSha512Drbg};
    use qs_policy::{load_policy_str, Format};
    use sha2::{Digest, Sha512};
    use std::collections::HashMap;
    use std::thread;
    use tempfile::tempdir;

    fn temp_path(name: &str) -> Utf8PathBuf {
        let base = std::env::temp_dir().join(format!("qs_quorum_test_{}_{}", name, now_unix()));
        Utf8PathBuf::from_path_buf(base).expect("utf8 path")
    }

    fn entropy(seed: u8) -> [u8; 48] {
        let mut out = [0u8; 48];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add(i as u8);
        }
        out
    }

    fn nonce(seed: u8) -> [u8; 16] {
        let mut out = [0u8; 16];
        for (i, byte) in out.iter_mut().enumerate() {
            *byte = seed.wrapping_add((i * 5) as u8);
        }
        out
    }

    #[test]
    fn quorum_round_trip() {
        let dir = temp_path("roundtrip");
        fs::create_dir_all(dir.as_std_path()).unwrap();

        let artifact = dir.join("artifact.bin");
        fs::write(artifact.as_std_path(), b"super secret").unwrap();

        let mut hasher = Sha512::new();
        hasher.update(b"super secret");
        let digest = hasher.finalize();
        let mut digest_arr = [0u8; 64];
        digest_arr.copy_from_slice(&digest);

        let policy_json = r#"{
            "default_alg": "mldsa-87",
            "allow_algs": ["mldsa-87"],
            "required_signatures": {"m": 2, "n": 3},
            "offline_ok": false,
            "require_fips_only": true,
            "digest_alg": "sha512"
        }"#;
        let policy = load_policy_str(policy_json, Some(Format::Json)).unwrap();

        // Generate keypairs for three participants so we can derive deterministic kids.
        let mut drbg_a = HmacSha512Drbg::new(&entropy(1), &nonce(2), Some(b"alice")).unwrap();
        let kp_a = keypair_mldsa87(&mut drbg_a).unwrap();
        let mut drbg_b = HmacSha512Drbg::new(&entropy(3), &nonce(4), Some(b"bob")).unwrap();
        let kp_b = keypair_mldsa87(&mut drbg_b).unwrap();
        let mut drbg_c = HmacSha512Drbg::new(&entropy(11), &nonce(12), Some(b"carol")).unwrap();
        let kp_c = keypair_mldsa87(&mut drbg_c).unwrap();

        let kid_a = qs_crypto::kid_from_public_key(&kp_a.public).unwrap();
        let kid_b = qs_crypto::kid_from_public_key(&kp_b.public).unwrap();
        let kid_c = qs_crypto::kid_from_public_key(&kp_c.public).unwrap();

        let intent_path = dir.join("intent.qsi");
        let allowed = vec![kid_a.clone(), kid_b.clone(), kid_c.clone()];
        quorum_init(
            &intent_path,
            &policy,
            &policy.default_alg,
            &digest_arr,
            &policy.digest_alg,
            &allowed,
        )
        .unwrap();
        let part_path = intent_path.with_extension("qsig.part");
        assert!(part_path.as_std_path().exists());

        // Produce fragments for two participants (quorum m=2)
        let intent = load_intent(&intent_path).unwrap();
        let mut sign_drbg_a = HmacSha512Drbg::new(&entropy(5), &nonce(6), Some(b"sign-a")).unwrap();
        let digest_alg_enum = DigestAlg::from_str(&intent.digest_alg).unwrap();
        let sig_a = sign_mldsa87(
            &mut sign_drbg_a,
            &kp_a.secret,
            intent.digest.as_ref(),
            digest_alg_enum,
            Some(intent.policy_hash.as_ref()),
        )
        .unwrap();
        let frag_a = make_fragment(&intent, &kid_a, sig_a);
        let frag_a_path = dir.join("alice.csf");
        save_fragment(&frag_a_path, &frag_a).unwrap();
        quorum_add(&part_path, &frag_a_path).unwrap();

        let mut sign_drbg_b = HmacSha512Drbg::new(&entropy(7), &nonce(8), Some(b"sign-b")).unwrap();
        let sig_b = sign_mldsa87(
            &mut sign_drbg_b,
            &kp_b.secret,
            intent.digest.as_ref(),
            digest_alg_enum,
            Some(intent.policy_hash.as_ref()),
        )
        .unwrap();
        let frag_b = make_fragment(&intent, &kid_b, sig_b);
        let frag_b_path = dir.join("bob.csf");
        save_fragment(&frag_b_path, &frag_b).unwrap();
        quorum_add(&part_path, &frag_b_path).unwrap();

        // Build trust store
        let mut trust = HashMap::new();
        trust.insert(
            kid_a.clone(),
            qs_crypto::public_key_to_spki(&kp_a.public).unwrap(),
        );
        trust.insert(
            kid_b.clone(),
            qs_crypto::public_key_to_spki(&kp_b.public).unwrap(),
        );
        trust.insert(
            kid_c.clone(),
            qs_crypto::public_key_to_spki(&kp_c.public).unwrap(),
        );
        let verify = move |kid: &str,
                           alg: &str,
                           digest: &[u8],
                           digest_alg: DigestAlg,
                           sig: &[u8],
                           policy: &[u8]|
              -> bool {
            if alg != "mldsa-87" {
                return false;
            }
            let spki = match trust.get(kid) {
                Some(k) => k,
                None => return false,
            };
            if qs_crypto::kid_from_spki_der(spki) != kid {
                return false;
            }
            qs_crypto::verify_mldsa87_spki(spki, digest, digest_alg, sig, Some(policy)).is_ok()
        };

        let final_sig = dir.join("artifact.qsig");
        quorum_seal(&part_path, &final_sig, verify).unwrap();
        assert!(final_sig.as_std_path().exists());
    }

    #[test]
    fn merge_retry_handles_concurrent_adds() {
        let dir = tempdir().unwrap();
        let part_std = dir.path().join("merge.qsig.part");
        let part_utf8 = Utf8PathBuf::from_path_buf(part_std.clone()).unwrap();

        let intent = Intent {
            v: VERSION,
            alg: "mldsa-87".into(),
            digest: ByteBuf::from(vec![0xAA; 64]),
            digest_alg: "sha512".into(),
            policy_hash: ByteBuf::from(vec![0xBB; 32]),
            m: 1,
            n: 3,
            allowed_kids: vec!["kid-a".into(), "kid-b".into(), "kid-c".into()],
            created_unix: 0,
        };
        let journal = Journal {
            v: VERSION,
            intent,
            rev: 0,
            fragments: BTreeMap::new(),
        };
        let bytes = to_cbor(&journal).unwrap();
        std::fs::write(&part_std, bytes).unwrap();

        let digest_len = DigestAlg::Sha512.output_len();
        let frag_bytes = |kid: &str| {
            to_cbor(&CosignFragment {
                v: VERSION,
                kid: kid.into(),
                alg: "mldsa-87".into(),
                digest: ByteBuf::from(vec![0xAA; digest_len]),
                sig: ByteBuf::from(vec![0x11; 2424]),
                created_unix: 0,
            })
            .unwrap()
        };

        let frag_a_std = dir.path().join("a.csf");
        std::fs::write(&frag_a_std, frag_bytes("kid-a")).unwrap();
        let frag_b_std = dir.path().join("b.csf");
        std::fs::write(&frag_b_std, frag_bytes("kid-b")).unwrap();

        let frag_a = Utf8PathBuf::from_path_buf(frag_a_std.clone()).unwrap();
        let part_for_a = part_utf8.clone();
        let t1 = thread::spawn(move || {
            quorum_add(&part_for_a, &frag_a).unwrap();
        });

        let frag_b = Utf8PathBuf::from_path_buf(frag_b_std.clone()).unwrap();
        let part_for_b = part_utf8.clone();
        let t2 = thread::spawn(move || {
            quorum_add(&part_for_b, &frag_b).unwrap();
        });

        t1.join().unwrap();
        t2.join().unwrap();

        let updated: Journal = load(&part_utf8).unwrap();
        assert!(updated.fragments.contains_key("kid-a"));
        assert!(updated.fragments.contains_key("kid-b"));
    }
}
