//! Block ordering and linearization for GHOSTDAG.
//!
//! Even though blocks form a DAG, transactions need a linear ordering for
//! execution. This module provides that ordering based on GHOSTDAG rules:
//!
//! 1. Blocks are ordered by their position in the selected chain
//! 2. Merge set blocks are ordered after the selected parent
//! 3. Within the merge set, blues come before reds
//! 4. Within blues/reds, ordering is by blue score then hash

use crate::{ghostdag::GhostdagManager, BlockId, BlueScore};
use std::cmp::Ordering;
use thiserror::Error;

/// An ordered block with its position and metadata.
#[derive(Clone, Debug)]
pub struct OrderedBlock {
    /// The block ID.
    pub id: BlockId,
    /// Position in the linear ordering.
    pub position: u64,
    /// The block's blue score.
    pub blue_score: BlueScore,
    /// Whether this block is blue (in the main chain).
    pub is_blue: bool,
    /// Whether this block is in a merge set (not on selected chain).
    pub is_merge_block: bool,
}

impl OrderedBlock {
    /// Creates a new ordered block.
    pub fn new(
        id: BlockId,
        position: u64,
        blue_score: BlueScore,
        is_blue: bool,
        is_merge_block: bool,
    ) -> Self {
        OrderedBlock {
            id,
            position,
            blue_score,
            is_blue,
            is_merge_block,
        }
    }
}

/// Block ordering manager.
pub struct BlockOrdering {
    /// Reference to the GHOSTDAG manager.
    ghostdag: std::sync::Arc<GhostdagManager>,
}

impl BlockOrdering {
    /// Creates a new block ordering manager.
    pub fn new(ghostdag: std::sync::Arc<GhostdagManager>) -> Self {
        BlockOrdering { ghostdag }
    }

    /// Gets the linear ordering of blocks from a given block back to genesis.
    ///
    /// This is the canonical ordering used for transaction execution.
    pub fn get_ordering(&self, from: &BlockId) -> Result<Vec<OrderedBlock>, OrderingError> {
        let mut ordering = Vec::new();
        let mut position = 0u64;

        // Get the selected chain first
        let selected_chain = self.ghostdag.get_selected_chain(from)?;

        // Process in reverse (genesis to tip)
        for chain_block in selected_chain.into_iter().rev() {
            let data = self.ghostdag.get_data(&chain_block)?;

            // First add the selected chain block
            ordering.push(OrderedBlock::new(
                chain_block,
                position,
                data.blue_score,
                true,
                false, // Not a merge block
            ));
            position += 1;

            // Then add merge set (blues first, then reds)
            let mut merge_blocks = Vec::new();

            // Add blues
            for blue in &data.merge_set_blues {
                let blue_score = self.ghostdag.get_blue_score(blue).unwrap_or(0);
                merge_blocks.push((*blue, blue_score, true));
            }

            // Add reds
            for red in &data.merge_set_reds {
                let blue_score = self.ghostdag.get_blue_score(red).unwrap_or(0);
                merge_blocks.push((*red, blue_score, false));
            }

            // Sort: blues before reds, then by blue score (desc), then by hash
            merge_blocks.sort_by(|a, b| {
                // Blues come first
                match (a.2, b.2) {
                    (true, false) => Ordering::Less,
                    (false, true) => Ordering::Greater,
                    _ => {
                        // Same color: sort by blue score descending
                        match b.1.cmp(&a.1) {
                            Ordering::Equal => {
                                // Same score: sort by hash for determinism
                                a.0.cmp(&b.0)
                            }
                            other => other,
                        }
                    }
                }
            });

            // Add sorted merge blocks
            for (block_id, blue_score, is_blue) in merge_blocks {
                ordering.push(OrderedBlock::new(
                    block_id, position, blue_score, is_blue, true, // Is a merge block
                ));
                position += 1;
            }
        }

        Ok(ordering)
    }

    /// Gets only the blocks that need to be processed between two points.
    ///
    /// Returns blocks that are in `to`'s past but not in `from`'s past.
    pub fn get_diff(
        &self,
        from: &BlockId,
        to: &BlockId,
    ) -> Result<Vec<OrderedBlock>, OrderingError> {
        let from_ordering = self.get_ordering(from)?;
        let to_ordering = self.get_ordering(to)?;

        // Get all blocks in from's past
        let from_blocks: std::collections::HashSet<_> =
            from_ordering.iter().map(|b| b.id).collect();

        // Return blocks in to's past that aren't in from's past
        let diff: Vec<_> = to_ordering
            .into_iter()
            .filter(|b| !from_blocks.contains(&b.id))
            .collect();

        Ok(diff)
    }

    /// Checks if block A comes before block B in the canonical ordering.
    pub fn is_before(&self, a: &BlockId, b: &BlockId) -> Result<bool, OrderingError> {
        let ordering = self.get_ordering(b)?;

        let pos_a = ordering.iter().position(|block| &block.id == a);
        let pos_b = ordering.iter().position(|block| &block.id == b);

        match (pos_a, pos_b) {
            (Some(pa), Some(pb)) => Ok(pa < pb),
            (Some(_), None) => Ok(true),  // A is in past of B
            (None, Some(_)) => Ok(false), // A is not in past of B
            (None, None) => Err(OrderingError::BlocksNotRelated(*a, *b)),
        }
    }

    /// Gets the merge set at a specific block in canonical order.
    pub fn get_merge_set_ordered(
        &self,
        block_id: &BlockId,
    ) -> Result<Vec<OrderedBlock>, OrderingError> {
        let data = self.ghostdag.get_data(block_id)?;
        let mut merge_blocks = Vec::new();
        let mut position = 0u64;

        // Add blues
        for blue in &data.merge_set_blues {
            let blue_score = self.ghostdag.get_blue_score(blue).unwrap_or(0);
            merge_blocks.push(OrderedBlock::new(*blue, position, blue_score, true, true));
            position += 1;
        }

        // Add reds
        for red in &data.merge_set_reds {
            let blue_score = self.ghostdag.get_blue_score(red).unwrap_or(0);
            merge_blocks.push(OrderedBlock::new(*red, position, blue_score, false, true));
            position += 1;
        }

        // Sort by blue score and hash
        merge_blocks.sort_by(|a, b| match (a.is_blue, b.is_blue) {
            (true, false) => Ordering::Less,
            (false, true) => Ordering::Greater,
            _ => match b.blue_score.cmp(&a.blue_score) {
                Ordering::Equal => a.id.cmp(&b.id),
                other => other,
            },
        });

        // Update positions after sorting
        for (i, block) in merge_blocks.iter_mut().enumerate() {
            block.position = i as u64;
        }

        Ok(merge_blocks)
    }
}

impl std::fmt::Debug for BlockOrdering {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BlockOrdering").finish()
    }
}

/// Errors that can occur during block ordering.
#[derive(Debug, Error)]
pub enum OrderingError {
    #[error("GHOSTDAG error: {0}")]
    GhostdagError(#[from] crate::ghostdag::GhostdagError),

    #[error("Blocks not related: {0} and {1}")]
    BlocksNotRelated(BlockId, BlockId),

    #[error("Block not found: {0}")]
    BlockNotFound(BlockId),
}

/// Iterator that yields blocks in canonical order.
pub struct OrderingIterator<'a> {
    ordering: &'a BlockOrdering,
    current_chain: Vec<BlockId>,
    chain_index: usize,
    merge_set_queue: Vec<OrderedBlock>,
}

impl<'a> OrderingIterator<'a> {
    /// Creates a new ordering iterator starting from the given block.
    pub fn new(ordering: &'a BlockOrdering, from: BlockId) -> Result<Self, OrderingError> {
        let selected_chain = ordering.ghostdag.get_selected_chain(&from)?;

        Ok(OrderingIterator {
            ordering,
            current_chain: selected_chain.into_iter().rev().collect(), // Reverse to go genesis -> tip
            chain_index: 0,
            merge_set_queue: Vec::new(),
        })
    }
}

impl<'a> Iterator for OrderingIterator<'a> {
    type Item = Result<OrderedBlock, OrderingError>;

    fn next(&mut self) -> Option<Self::Item> {
        // First drain the merge set queue
        if let Some(block) = self.merge_set_queue.pop() {
            return Some(Ok(block));
        }

        // Get next block from chain
        if self.chain_index >= self.current_chain.len() {
            return None;
        }

        let chain_block = self.current_chain[self.chain_index];
        self.chain_index += 1;

        // Get GHOSTDAG data for this block
        let data = match self.ordering.ghostdag.get_data(&chain_block) {
            Ok(d) => d,
            Err(e) => return Some(Err(e.into())),
        };

        // Queue up the merge set (in reverse order so we pop in correct order)
        match self.ordering.get_merge_set_ordered(&chain_block) {
            Ok(merge_set) => {
                self.merge_set_queue = merge_set.into_iter().rev().collect();
            }
            Err(e) => return Some(Err(e)),
        }

        // Return the chain block
        Some(Ok(OrderedBlock::new(
            chain_block,
            (self.chain_index - 1) as u64,
            data.blue_score,
            true,
            false,
        )))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{dag::BlockDag, reachability::ReachabilityStore};
    use synor_types::Hash256;

    fn make_block_id(n: u8) -> BlockId {
        let mut bytes = [0u8; 32];
        bytes[0] = n;
        Hash256::from_bytes(bytes)
    }

    fn setup_test() -> (
        std::sync::Arc<BlockDag>,
        std::sync::Arc<ReachabilityStore>,
        std::sync::Arc<GhostdagManager>,
        BlockOrdering,
    ) {
        let genesis = make_block_id(0);
        let dag = std::sync::Arc::new(BlockDag::new(genesis, 0));
        let reachability = std::sync::Arc::new(ReachabilityStore::new(genesis));
        let ghostdag = std::sync::Arc::new(GhostdagManager::with_k(
            dag.clone(),
            reachability.clone(),
            3,
        ));
        let ordering = BlockOrdering::new(ghostdag.clone());
        (dag, reachability, ghostdag, ordering)
    }

    #[test]
    fn test_simple_chain_ordering() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag, ordering) = setup_test();

        // Build a simple chain
        let block1 = make_block_id(1);
        let block2 = make_block_id(2);

        dag.insert_block(block1, vec![genesis], 100).unwrap();
        dag.insert_block(block2, vec![block1], 200).unwrap();

        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        reachability.add_block(block2, block1, &[block1]).unwrap();

        ghostdag.add_block(block1, &[genesis]).unwrap();
        ghostdag.add_block(block2, &[block1]).unwrap();

        let order = ordering.get_ordering(&block2).unwrap();

        // Should be: genesis, block1, block2
        assert_eq!(order.len(), 3);
        assert_eq!(order[0].id, genesis);
        assert_eq!(order[1].id, block1);
        assert_eq!(order[2].id, block2);
    }

    #[test]
    fn test_is_before() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag, ordering) = setup_test();

        let block1 = make_block_id(1);
        let block2 = make_block_id(2);

        dag.insert_block(block1, vec![genesis], 100).unwrap();
        dag.insert_block(block2, vec![block1], 200).unwrap();

        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        reachability.add_block(block2, block1, &[block1]).unwrap();

        ghostdag.add_block(block1, &[genesis]).unwrap();
        ghostdag.add_block(block2, &[block1]).unwrap();

        assert!(ordering.is_before(&genesis, &block2).unwrap());
        assert!(ordering.is_before(&block1, &block2).unwrap());
        assert!(!ordering.is_before(&block2, &genesis).unwrap());
    }

    #[test]
    fn test_ordering_positions() {
        let genesis = make_block_id(0);
        let (dag, reachability, ghostdag, ordering) = setup_test();

        let block1 = make_block_id(1);
        dag.insert_block(block1, vec![genesis], 100).unwrap();
        reachability.add_block(block1, genesis, &[genesis]).unwrap();
        ghostdag.add_block(block1, &[genesis]).unwrap();

        let order = ordering.get_ordering(&block1).unwrap();

        // Check positions are sequential
        for (i, block) in order.iter().enumerate() {
            assert_eq!(block.position, i as u64);
        }
    }
}
