//! Storage Proofs - Proof of Spacetime
//!
//! Storage nodes must prove they're actually storing data.
//! Proofs are verified on-chain (L1) for reward distribution.

use crate::cid::ContentId;
use serde::{Deserialize, Serialize};

/// Node identifier (public key)
pub type NodeId = [u8; 32];

/// Storage proof challenge from L1
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Challenge {
    /// Block hash used as randomness source
    pub block_hash: [u8; 32],
    /// Challenge epoch number
    pub epoch: u64,
    /// CID being challenged
    pub cid: ContentId,
    /// Random chunk index to prove
    pub chunk_index: u32,
    /// Random byte range within chunk to return
    pub byte_offset: u64,
    pub byte_length: u64,
}

impl Challenge {
    /// Generate a challenge from block hash and CID
    pub fn generate(
        block_hash: [u8; 32],
        epoch: u64,
        cid: ContentId,
        total_chunks: u32,
        chunk_size: usize,
    ) -> Self {
        // Use block hash to deterministically select chunk and range
        let chunk_seed = u64::from_le_bytes(block_hash[0..8].try_into().unwrap());
        let chunk_index = (chunk_seed % total_chunks as u64) as u32;

        let offset_seed = u64::from_le_bytes(block_hash[8..16].try_into().unwrap());
        let byte_offset = offset_seed % (chunk_size as u64 / 2);
        let byte_length = 1024.min(chunk_size as u64 - byte_offset); // Max 1KB proof

        Self {
            block_hash,
            epoch,
            cid,
            chunk_index,
            byte_offset,
            byte_length,
        }
    }
}

/// Storage proof submitted by node
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageProof {
    /// Node submitting the proof
    pub node_id: NodeId,
    /// Challenge being answered
    pub challenge: Challenge,
    /// Merkle root of the full chunk
    pub chunk_merkle_root: [u8; 32],
    /// Proof data (requested byte range)
    #[serde(with = "serde_bytes")]
    pub proof_data: Vec<u8>,
    /// Merkle proof showing proof_data is part of chunk
    pub merkle_proof: Vec<[u8; 32]>,
    /// Timestamp of proof generation
    pub timestamp: u64,
    /// Signature over the proof (64 bytes, stored as Vec for serde compatibility)
    #[serde(with = "serde_bytes")]
    pub signature: Vec<u8>,
}

impl StorageProof {
    /// Verify the proof is valid
    pub fn verify(&self) -> bool {
        // 1. Verify proof data length matches challenge
        if self.proof_data.len() != self.challenge.byte_length as usize {
            return false;
        }

        // 2. Verify Merkle proof
        if !self.verify_merkle_proof() {
            return false;
        }

        // 3. Verify signature (TODO: implement with actual crypto)
        // self.verify_signature()

        true
    }

    fn verify_merkle_proof(&self) -> bool {
        if self.merkle_proof.is_empty() {
            return false;
        }

        // Compute hash of proof data
        let mut current = *blake3::hash(&self.proof_data).as_bytes();

        // Walk up the Merkle tree
        for sibling in &self.merkle_proof {
            let mut combined = [0u8; 64];
            combined[..32].copy_from_slice(&current);
            combined[32..].copy_from_slice(sibling);
            current = *blake3::hash(&combined).as_bytes();
        }

        // Should match chunk Merkle root
        current == self.chunk_merkle_root
    }
}

/// Build a Merkle tree from chunk data
pub fn build_merkle_tree(data: &[u8], leaf_size: usize) -> MerkleTree {
    let leaves: Vec<[u8; 32]> = data
        .chunks(leaf_size)
        .map(|chunk| *blake3::hash(chunk).as_bytes())
        .collect();

    MerkleTree::from_leaves(leaves)
}

/// Simple Merkle tree for chunk proofs
#[derive(Debug, Clone)]
pub struct MerkleTree {
    /// All nodes in the tree (leaves first, then internal nodes)
    nodes: Vec<[u8; 32]>,
    /// Number of leaves
    leaf_count: usize,
}

impl MerkleTree {
    /// Build tree from leaf hashes
    pub fn from_leaves(leaves: Vec<[u8; 32]>) -> Self {
        let leaf_count = leaves.len();
        let mut nodes = leaves;

        // Build tree bottom-up
        let mut level_start = 0;
        let mut level_size = leaf_count;

        while level_size > 1 {
            let next_level_size = level_size.div_ceil(2);

            for i in 0..next_level_size {
                let left_idx = level_start + i * 2;
                let right_idx = left_idx + 1;

                let left = nodes[left_idx];
                let right = if right_idx < level_start + level_size {
                    nodes[right_idx]
                } else {
                    left // Duplicate for odd number
                };

                let mut combined = [0u8; 64];
                combined[..32].copy_from_slice(&left);
                combined[32..].copy_from_slice(&right);
                nodes.push(*blake3::hash(&combined).as_bytes());
            }

            level_start += level_size;
            level_size = next_level_size;
        }

        Self { nodes, leaf_count }
    }

    /// Get the root hash
    pub fn root(&self) -> [u8; 32] {
        *self.nodes.last().unwrap_or(&[0u8; 32])
    }

    /// Generate proof for a leaf
    pub fn proof(&self, leaf_index: usize) -> Vec<[u8; 32]> {
        if leaf_index >= self.leaf_count {
            return Vec::new();
        }

        let mut proof = Vec::new();
        let mut idx = leaf_index;
        let mut level_start = 0;
        let mut level_size = self.leaf_count;

        while level_size > 1 {
            let sibling_idx = if idx.is_multiple_of(2) {
                idx + 1
            } else {
                idx - 1
            };

            if sibling_idx < level_size {
                proof.push(self.nodes[level_start + sibling_idx]);
            } else {
                proof.push(self.nodes[level_start + idx]);
            }

            idx /= 2;
            level_start += level_size;
            level_size = level_size.div_ceil(2);
        }

        proof
    }
}

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

    #[test]
    fn test_merkle_tree() {
        let data = b"AAAABBBBCCCCDDDD";
        let tree = build_merkle_tree(data, 4);

        assert!(!tree.nodes.is_empty());
        assert_eq!(tree.leaf_count, 4);
    }

    #[test]
    fn test_merkle_proof() {
        let leaves: Vec<[u8; 32]> = (0..4).map(|i| *blake3::hash(&[i]).as_bytes()).collect();

        let tree = MerkleTree::from_leaves(leaves.clone());
        let root = tree.root();

        // Verify proof for each leaf
        for (i, leaf) in leaves.iter().enumerate() {
            let proof = tree.proof(i);

            // Verify by walking up
            let mut current = *leaf;
            for (j, sibling) in proof.iter().enumerate() {
                let mut combined = [0u8; 64];
                if (i >> j).is_multiple_of(2) {
                    combined[..32].copy_from_slice(&current);
                    combined[32..].copy_from_slice(sibling);
                } else {
                    combined[..32].copy_from_slice(sibling);
                    combined[32..].copy_from_slice(&current);
                }
                current = *blake3::hash(&combined).as_bytes();
            }

            assert_eq!(current, root);
        }
    }

    #[test]
    fn test_challenge_generation() {
        let block_hash = [1u8; 32];
        let cid = ContentId::from_content(b"test");

        let challenge = Challenge::generate(block_hash, 1, cid, 100, 1024);

        assert!(challenge.chunk_index < 100);
        assert!(challenge.byte_length <= 1024);
    }
}