//! Mempool service.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use tokio::sync::{broadcast, RwLock};
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};

use crate::config::NodeConfig;
use crate::services::ConsensusService;

/// Transaction in mempool.
#[derive(Clone, Debug)]
pub struct MempoolTx {
    /// Transaction hash.
    pub hash: [u8; 32],
    /// Raw transaction bytes.
    pub data: Vec<u8>,
    /// Transaction mass (for prioritization).
    pub mass: u64,
    /// Fee in sompi.
    pub fee: u64,
    /// Fee per mass unit.
    pub fee_rate: f64,
    /// Entry timestamp.
    pub timestamp: u64,
    /// Dependencies (inputs from other mempool txs).
    pub dependencies: Vec<[u8; 32]>,
    /// Is high priority.
    pub high_priority: bool,
}

/// Maximum transaction age before expiration (24 hours in seconds).
const MAX_TX_AGE_SECS: u64 = 24 * 60 * 60;

/// Cleanup interval (5 minutes).
const CLEANUP_INTERVAL_SECS: u64 = 5 * 60;

/// Mempool service manages unconfirmed transactions.
pub struct MempoolService {
    /// Consensus reference.
    consensus: Arc<ConsensusService>,

    /// Transactions in mempool.
    txs: RwLock<HashMap<[u8; 32], MempoolTx>>,

    /// Maximum mempool size in bytes.
    max_size: usize,

    /// Current size in bytes.
    current_size: RwLock<usize>,

    /// Is running.
    running: RwLock<bool>,

    /// Shutdown receiver.
    shutdown_rx: broadcast::Receiver<()>,

    /// Transaction added channel.
    tx_added: broadcast::Sender<[u8; 32]>,

    /// Transaction removed channel.
    tx_removed: broadcast::Sender<[u8; 32]>,

    /// Cleanup task handle.
    cleanup_handle: RwLock<Option<JoinHandle<()>>>,
}

impl MempoolService {
    /// Creates a new mempool service.
    pub fn new(
        consensus: Arc<ConsensusService>,
        _config: &NodeConfig,
        shutdown_rx: broadcast::Receiver<()>,
    ) -> anyhow::Result<Self> {
        let (tx_added, _) = broadcast::channel(1000);
        let (tx_removed, _) = broadcast::channel(1000);

        Ok(MempoolService {
            consensus,
            txs: RwLock::new(HashMap::new()),
            max_size: 100 * 1024 * 1024, // 100MB
            current_size: RwLock::new(0),
            running: RwLock::new(false),
            shutdown_rx,
            tx_added,
            tx_removed,
            cleanup_handle: RwLock::new(None),
        })
    }

    /// Starts the mempool service.
    pub async fn start(&self) -> anyhow::Result<()> {
        info!("Starting mempool service");
        *self.running.write().await = true;
        Ok(())
    }

    /// Spawns the cleanup task. Must be called after wrapping in Arc.
    pub fn spawn_cleanup_task(self: &Arc<Self>) {
        let mempool = Arc::clone(self);
        let mut block_rx = self.consensus.subscribe_blocks();

        let handle = tokio::spawn(async move {
            let mut cleanup_interval =
                tokio::time::interval(Duration::from_secs(CLEANUP_INTERVAL_SECS));
            cleanup_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            info!("Mempool cleanup task started");

            loop {
                // Check if we should stop
                if !*mempool.running.read().await {
                    info!("Mempool cleanup task shutting down");
                    break;
                }

                tokio::select! {
                    // Handle new accepted blocks
                    result = block_rx.recv() => {
                        match result {
                            Ok(block_hash) => {
                                debug!(hash = hex::encode(&block_hash[..8]), "Block accepted, cleaning mempool");

                                // Get transactions from the block and remove them from mempool
                                // Note: In a full implementation, we'd query the block's transactions
                                // For now, we rely on external calls to clear_accepted()
                                mempool.remove_transaction(&block_hash).await;
                            }
                            Err(broadcast::error::RecvError::Lagged(n)) => {
                                warn!(missed = n, "Mempool cleanup lagged behind block notifications");
                            }
                            Err(broadcast::error::RecvError::Closed) => {
                                info!("Block channel closed, stopping cleanup task");
                                break;
                            }
                        }
                    }

                    // Periodic cleanup of expired transactions
                    _ = cleanup_interval.tick() => {
                        let expired = mempool.expire_old_transactions().await;
                        if expired > 0 {
                            info!(count = expired, "Expired old transactions from mempool");
                        }
                    }
                }
            }
        });

        // Store handle for cleanup on stop
        let mempool_clone = Arc::clone(self);
        tokio::spawn(async move {
            *mempool_clone.cleanup_handle.write().await = Some(handle);
        });
    }

    /// Stops the mempool service.
    pub async fn stop(&self) -> anyhow::Result<()> {
        info!("Stopping mempool service");
        *self.running.write().await = false;

        // Abort cleanup task if running
        if let Some(handle) = self.cleanup_handle.write().await.take() {
            handle.abort();
        }

        Ok(())
    }

    /// Expires old transactions from the mempool.
    /// Returns the number of expired transactions.
    async fn expire_old_transactions(&self) -> usize {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();

        let mut txs = self.txs.write().await;
        let mut current_size = self.current_size.write().await;

        let mut expired = Vec::new();

        // Find expired transactions
        for (hash, tx) in txs.iter() {
            let age_secs = now.saturating_sub(tx.timestamp / 1000); // timestamp is in millis
            if age_secs > MAX_TX_AGE_SECS && !tx.high_priority {
                expired.push((*hash, tx.data.len()));
            }
        }

        // Remove expired transactions
        let count = expired.len();
        for (hash, size) in expired {
            txs.remove(&hash);
            *current_size = current_size.saturating_sub(size);
            let _ = self.tx_removed.send(hash);
        }

        count
    }

    /// Adds a transaction to the mempool.
    pub async fn add_transaction(&self, tx: MempoolTx) -> anyhow::Result<()> {
        let hash = tx.hash;

        // Check if already exists
        {
            let txs = self.txs.read().await;
            if txs.contains_key(&hash) {
                return Ok(());
            }
        }

        // Check mempool size
        let tx_size = tx.data.len();
        {
            let current = *self.current_size.read().await;
            if current + tx_size > self.max_size {
                // Evict low-priority transactions
                self.evict(tx_size).await?;
            }
        }

        // Add transaction
        {
            let mut txs = self.txs.write().await;
            let mut current_size = self.current_size.write().await;

            txs.insert(hash, tx);
            *current_size += tx_size;
        }

        debug!(
            hash = hex::encode(&hash[..8]),
            "Transaction added to mempool"
        );
        let _ = self.tx_added.send(hash);

        Ok(())
    }

    /// Removes a transaction from the mempool.
    pub async fn remove_transaction(&self, hash: &[u8; 32]) -> Option<MempoolTx> {
        let mut txs = self.txs.write().await;
        if let Some(tx) = txs.remove(hash) {
            let mut current_size = self.current_size.write().await;
            *current_size = current_size.saturating_sub(tx.data.len());

            debug!(
                hash = hex::encode(&hash[..8]),
                "Transaction removed from mempool"
            );
            let _ = self.tx_removed.send(*hash);

            Some(tx)
        } else {
            None
        }
    }

    /// Gets a transaction from the mempool.
    pub async fn get_transaction(&self, hash: &[u8; 32]) -> Option<MempoolTx> {
        self.txs.read().await.get(hash).cloned()
    }

    /// Checks if transaction is in mempool.
    pub async fn contains(&self, hash: &[u8; 32]) -> bool {
        self.txs.read().await.contains_key(hash)
    }

    /// Returns the number of transactions.
    pub async fn count(&self) -> usize {
        self.txs.read().await.len()
    }

    /// Returns the current size in bytes.
    pub async fn size(&self) -> usize {
        *self.current_size.read().await
    }

    /// Gets all transaction hashes.
    pub async fn all_hashes(&self) -> Vec<[u8; 32]> {
        self.txs.read().await.keys().copied().collect()
    }

    /// Gets transactions for block template.
    pub async fn select_for_block(&self, max_mass: u64) -> Vec<MempoolTx> {
        let txs = self.txs.read().await;

        // Sort by fee rate descending
        let mut sorted: Vec<_> = txs.values().cloned().collect();
        sorted.sort_by(|a, b| {
            b.fee_rate
                .partial_cmp(&a.fee_rate)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Select transactions up to max mass
        let mut selected = Vec::new();
        let mut total_mass = 0u64;

        for tx in sorted {
            if total_mass + tx.mass <= max_mass {
                total_mass += tx.mass;
                selected.push(tx);
            }
        }

        selected
    }

    /// Evicts transactions to make room.
    async fn evict(&self, needed: usize) -> anyhow::Result<()> {
        let mut txs = self.txs.write().await;
        let mut current_size = self.current_size.write().await;

        // Sort by fee rate ascending (evict lowest first)
        let mut sorted: Vec<_> = txs.values().cloned().collect();
        sorted.sort_by(|a, b| {
            a.fee_rate
                .partial_cmp(&b.fee_rate)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        let mut freed = 0usize;
        let mut to_remove = Vec::new();

        for tx in sorted {
            if freed >= needed {
                break;
            }
            if !tx.high_priority {
                freed += tx.data.len();
                to_remove.push(tx.hash);
            }
        }

        for hash in to_remove {
            if let Some(tx) = txs.remove(&hash) {
                *current_size = current_size.saturating_sub(tx.data.len());
                let _ = self.tx_removed.send(hash);
            }
        }

        Ok(())
    }

    /// Clears transactions included in a block.
    pub async fn clear_accepted(&self, tx_hashes: &[[u8; 32]]) {
        for hash in tx_hashes {
            self.remove_transaction(hash).await;
        }
    }

    /// Subscribes to new transactions.
    pub fn subscribe_added(&self) -> broadcast::Receiver<[u8; 32]> {
        self.tx_added.subscribe()
    }

    /// Subscribes to removed transactions.
    pub fn subscribe_removed(&self) -> broadcast::Receiver<[u8; 32]> {
        self.tx_removed.subscribe()
    }

    /// Gets mempool statistics.
    pub async fn stats(&self) -> MempoolStats {
        let txs = self.txs.read().await;

        let total_fees: u64 = txs.values().map(|t| t.fee).sum();
        let total_mass: u64 = txs.values().map(|t| t.mass).sum();

        MempoolStats {
            tx_count: txs.len(),
            size_bytes: *self.current_size.read().await,
            total_fees,
            total_mass,
        }
    }
}

/// Mempool statistics.
#[derive(Clone, Debug)]
pub struct MempoolStats {
    pub tx_count: usize,
    pub size_bytes: usize,
    pub total_fees: u64,
    pub total_mass: u64,
}
