rostrum 14.0.1

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use anyhow::Result;
use bitcoin_hashes::Hash;
use futures::StreamExt;
use sha2::{Digest, Sha256};
use tokio::sync::{mpsc, Mutex};

use crate::chaindef::ScriptHash;
use crate::indexes::scripthashindex::ScriptHashIndexRow;
use crate::indexes::DBRow;
use crate::mempool::Tracker;
use crate::query::Query;
use crate::rpc::daemon::server::ConnectionId;
use crate::rpc::daemon::ScriptHashUpdate;
use crate::store::{DBStore, Row};

/// Computes a statushash that is not true to the electrum spec; but very fast.
async fn faster_statushash(
    confirmed: &Arc<DBStore>,
    mempool: &Tracker,
    scripthash: ScriptHash,
    max_elements: Option<usize>,
) -> Result<Option<[u8; 32]>> {
    const DEFAULT_MAX_ELEMENTS: usize = 10_000;
    let max_elements = max_elements.unwrap_or(DEFAULT_MAX_ELEMENTS);

    // Mempool transactions
    let scan_filter = ScriptHashIndexRow::filter_include_all(scripthash.into_inner());
    let (query, stream) = mempool
        .index()
        .scan(ScriptHashIndexRow::CF, scan_filter.clone(), None)
        .await;

    let mut hasher = Sha256::new();
    let mut count = 0;

    stream
        .take(max_elements)
        .for_each(|row| {
            let row = ScriptHashIndexRow::from_row(&row);
            let oph_inner = row.outpointhash_inner();
            hasher.update(oph_inner);
            let height = row.get_height();
            let height_bytes = height.to_le_bytes();
            hasher.update(height_bytes);
            count += 1;
            futures::future::ready(())
        })
        .await;

    // When we stop consuming early (via take), the channel receiver is dropped,
    // causing the scan task to get "channel closed" error. This is expected and we should ignore it.
    if let Err(e) = query.await? {
        if !e
            .downcast_ref::<mpsc::error::SendTimeoutError<Row>>()
            .map(|err| matches!(err, mpsc::error::SendTimeoutError::Closed(_)))
            .unwrap_or(false)
        {
            return Err(e);
        }
    }

    // Using reverse scan to get most recent transactions first, capping at max_elements.
    let (query, stream) = confirmed
        .rscan(ScriptHashIndexRow::CF, scan_filter, None)
        .await;
    stream
        .take(max_elements.saturating_sub(count))
        .for_each(|row| {
            let row = ScriptHashIndexRow::from_row(&row);
            let oph_inner = row.outpointhash_inner();
            hasher.update(oph_inner);
            let height = row.get_height();
            let height_bytes = height.to_le_bytes();
            hasher.update(height_bytes);
            count += 1;
            futures::future::ready(())
        })
        .await;

    // When we stop consuming early (via take), the channel receiver is dropped,
    // causing the scan task to get "channel closed" error. This is expected and we should ignore it.
    if let Err(e) = query.await? {
        if !e
            .downcast_ref::<mpsc::error::SendTimeoutError<Row>>()
            .map(|err| matches!(err, mpsc::error::SendTimeoutError::Closed(_)))
            .unwrap_or(false)
        {
            return Err(e);
        }
    }

    if count != 0 {
        let combined_hash = hasher.finalize().into();
        return Ok(Some(combined_hash));
    }

    Ok(None)
}

/// Subscription metadata stored per connection
#[derive(Clone)]
pub struct SubscriptionMetadata {
    pub old_statushash: Option<[u8; 32]>,
    pub alias: Option<String>,
}

/// Manages scripthash subscriptions - maps scripthashes to connections that subscribe to them
pub struct ScripthashSubscriptions {
    /// Maps scripthash -> map of connection IDs to their subscription metadata
    index: Arc<Mutex<HashMap<ScriptHash, HashMap<ConnectionId, SubscriptionMetadata>>>>,
    /// Maps scripthash -> current statushash (cached, shared across all subscriptions)
    statushash_cache: Arc<Mutex<HashMap<ScriptHash, Option<[u8; 32]>>>>,
    query: Arc<Query>,
    max_elements: usize,
}

impl ScripthashSubscriptions {
    pub fn new(query: Arc<Query>, max_elements: usize) -> Self {
        ScripthashSubscriptions {
            index: Arc::new(Mutex::new(HashMap::new())),
            statushash_cache: Arc::new(Mutex::new(HashMap::new())),
            query,
            max_elements,
        }
    }

    /// Get or compute the current statushash for a scripthash
    /// Returns the statushash, computing it if not cached
    async fn get_or_compute_statushash(&self, scripthash: ScriptHash) -> Result<Option<[u8; 32]>> {
        // Check cache first
        {
            let cache = self.statushash_cache.lock().await;
            if let Some(cached) = cache.get(&scripthash) {
                return Ok(*cached);
            }
        }

        let statushash = faster_statushash(
            self.query.confirmed_index(),
            self.query.unconfirmed_index(),
            scripthash,
            Some(self.max_elements),
        )
        .await?;
        // Store in cache
        {
            let mut cache = self.statushash_cache.lock().await;
            cache.insert(scripthash, statushash);
        }

        Ok(statushash)
    }

    /// Register a subscription: connection `conn_id` subscribes to `scripthash`
    /// Returns the current statushash (computed if needed)
    pub async fn subscribe(
        &self,
        conn_id: ConnectionId,
        scripthash: ScriptHash,
        alias: Option<String>,
    ) -> Result<Option<[u8; 32]>> {
        // Get or compute statushash
        let statushash = self.get_or_compute_statushash(scripthash).await?;

        // Store subscription
        let mut index = self.index.lock().await;
        index.entry(scripthash).or_insert_with(HashMap::new).insert(
            conn_id,
            SubscriptionMetadata {
                old_statushash: statushash,
                alias,
            },
        );

        Ok(statushash)
    }

    /// Unregister a subscription: connection `conn_id` unsubscribes from `scripthash`
    pub async fn unsubscribe(&self, conn_id: ConnectionId, scripthash: ScriptHash) {
        let mut index = self.index.lock().await;
        if let Some(conns) = index.get_mut(&scripthash) {
            conns.remove(&conn_id);
            if conns.is_empty() {
                index.remove(&scripthash);
                // Clean up cache when no subscriptions remain
                let mut cache = self.statushash_cache.lock().await;
                cache.remove(&scripthash);
            }
        }
    }

    /// Unregister all subscriptions for a connection (when it disconnects)
    pub async fn remove_connection(&self, conn_id: ConnectionId) {
        let mut index = self.index.lock().await;
        let mut cache = self.statushash_cache.lock().await;
        let mut scripthashes_to_remove = Vec::new();

        for (scripthash, conns) in index.iter_mut() {
            if conns.remove(&conn_id).is_some() && conns.is_empty() {
                scripthashes_to_remove.push(*scripthash);
            }
        }

        // Remove empty subscriptions and clean up cache
        for scripthash in scripthashes_to_remove {
            index.remove(&scripthash);
            cache.remove(&scripthash);
        }
    }

    /// Get all subscriptions for the given scripthashes with their metadata
    pub async fn get_subscriptions(
        &self,
        scripthashes: &HashSet<ScriptHash>,
    ) -> HashMap<ScriptHash, HashMap<ConnectionId, SubscriptionMetadata>> {
        let index = self.index.lock().await;
        let mut result = HashMap::new();

        for scripthash in scripthashes {
            if let Some(conns) = index.get(scripthash) {
                result.insert(*scripthash, conns.clone());
            }
        }

        result
    }

    /// Update statushash for all subscriptions of a scripthash after notification
    pub async fn update_statushash(
        &self,
        scripthash: ScriptHash,
        new_statushash: Option<[u8; 32]>,
    ) {
        // Update cache
        {
            let mut cache = self.statushash_cache.lock().await;
            cache.insert(scripthash, new_statushash);
        }

        // Update all subscriptions for this scripthash
        let mut index = self.index.lock().await;
        if let Some(conns) = index.get_mut(&scripthash) {
            for metadata in conns.values_mut() {
                metadata.old_statushash = new_statushash;
            }
        }
    }

    /// Notify that scripthashes have changed, compute statushash updates, and return grouped by connection
    pub async fn notify_scripthashes_changed(
        &self,
        scripthashes: HashSet<ScriptHash>,
    ) -> Result<HashMap<ConnectionId, Vec<ScriptHashUpdate>>> {
        if scripthashes.is_empty() {
            return Ok(HashMap::new());
        }

        // Invalidate cache for affected scripthashes to ensure fresh computation
        {
            let mut cache = self.statushash_cache.lock().await;
            for scripthash in &scripthashes {
                cache.remove(scripthash);
            }
        }

        // Get all subscriptions for affected scripthashes
        let subscriptions = self.get_subscriptions(&scripthashes).await;

        if subscriptions.is_empty() {
            return Ok(HashMap::new());
        }

        // Compute statushash once per unique scripthash
        let mut statushash_updates: HashMap<ScriptHash, Option<[u8; 32]>> = HashMap::new();

        for scripthash in &scripthashes {
            if !subscriptions.contains_key(scripthash) {
                continue;
            }

            // Only compute if we haven't already
            if statushash_updates.contains_key(scripthash) {
                continue;
            }

            let new_statushash = faster_statushash(
                self.query.confirmed_index(),
                self.query.unconfirmed_index(),
                *scripthash,
                Some(self.max_elements),
            )
            .await?;

            statushash_updates.insert(*scripthash, new_statushash);
        }

        // Group updates by connection
        let mut connection_updates: HashMap<ConnectionId, Vec<ScriptHashUpdate>> = HashMap::new();

        for (scripthash, conn_subscriptions) in subscriptions {
            let new_statushash = match statushash_updates.get(&scripthash) {
                Some(s) => *s,
                None => continue, // Skip if computation failed
            };

            // Update cache and all subscriptions
            self.update_statushash(scripthash, new_statushash).await;

            for (conn_id, metadata) in conn_subscriptions {
                // Only notify if statushash actually changed
                if new_statushash == metadata.old_statushash {
                    continue;
                }

                connection_updates
                    .entry(conn_id)
                    .or_default()
                    .push(ScriptHashUpdate {
                        scripthash,
                        old_statushash: metadata.old_statushash,
                        new_statushash,
                        alias: metadata.alias.clone(),
                    });
            }
        }

        Ok(connection_updates)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::chaindef::OutPointHash;
    use crate::indexes::scripthashindex::OutputFlags;
    use crate::mempool::Tracker;
    use crate::metrics::Metrics;
    use crate::store::{DBContents, DBStore};
    use crate::writebatch::WriteBatch;
    use std::env;
    use std::sync::Arc;

    #[tokio::test]
    async fn test_faster_statushash_max_elements() {
        // Create temporary directories for stores
        let temp_dir = env::temp_dir();
        let test_id = format!("rostrum_test_{}", std::process::id());
        let confirmed_path = temp_dir.join(&format!("{}_confirmed", test_id));
        let mempool_path = temp_dir.join(&format!("{}_mempool", test_id));

        // Clean up any existing test directories
        let _ = std::fs::remove_dir_all(&confirmed_path);
        let _ = std::fs::remove_dir_all(&mempool_path);

        // Create metrics (disabled)
        let metrics = Arc::new(
            Metrics::new("127.0.0.1:0".parse().unwrap(), false).expect("Failed to create metrics"),
        );

        // Create stores
        let confirmed = Arc::new(
            DBStore::open(DBContents::ConfirmedIndex, &confirmed_path, &metrics, true)
                .expect("Failed to create confirmed store"),
        );

        let mempool = Arc::new(Tracker::new(&mempool_path, &metrics));

        // Create a test scripthash
        let scripthash = ScriptHash::hash(&[42u8; 32]);

        const TOTAL_ELEMENTS: usize = 15_000;
        const TEST_MAX_ELEMENTS: usize = 5_000;

        let batch = WriteBatch::new();
        let rows: Vec<_> = (0..TOTAL_ELEMENTS)
            .map(|i| {
                let outpointhash = OutPointHash::hash(&i.to_le_bytes());
                let txid = [i as u8; 32];
                ScriptHashIndexRow::new_funding(
                    &scripthash,
                    &outpointhash,
                    OutputFlags::FundingNone,
                    txid,
                    0,
                    i as u32,
                )
                .to_row()
            })
            .collect();
        batch.insert(ScriptHashIndexRow::CF, rows);

        // Write to confirmed store
        confirmed.write_batch(&batch);
        confirmed.flush().expect("Failed to flush");
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        // Test with a smaller max_elements than total entries
        let result = faster_statushash(&confirmed, &mempool, scripthash, Some(TEST_MAX_ELEMENTS))
            .await
            .expect("faster_statushash should succeed");
        assert!(result.is_some(), "Should return a hash when entries exist");

        // Test with max_elements larger than total entries
        let result_all = faster_statushash(
            &confirmed,
            &mempool,
            scripthash,
            Some(TOTAL_ELEMENTS + 1000),
        )
        .await
        .expect("faster_statushash should succeed with larger max_elements");
        assert!(
            result_all.is_some(),
            "Should return a hash when processing all entries"
        );

        // Adding a new entry at a higher height should always change the hash,
        // even when total entries exceed max_elements (this is the bug regression test).
        let old_hash = result.unwrap();
        let new_entry_batch = WriteBatch::new();
        let new_outpointhash = OutPointHash::hash(&TOTAL_ELEMENTS.to_le_bytes());
        let new_row = ScriptHashIndexRow::new_funding(
            &scripthash,
            &new_outpointhash,
            OutputFlags::FundingNone,
            [0xFFu8; 32],
            0,
            TOTAL_ELEMENTS as u32, // higher height than all existing entries
        )
        .to_row();
        new_entry_batch.insert(ScriptHashIndexRow::CF, vec![new_row]);
        confirmed.write_batch(&new_entry_batch);
        confirmed.flush().expect("Failed to flush");
        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;

        let new_hash = faster_statushash(&confirmed, &mempool, scripthash, Some(TEST_MAX_ELEMENTS))
            .await
            .expect("faster_statushash should succeed after adding entry")
            .expect("Should return a hash");

        assert_ne!(
            old_hash, new_hash,
            "Statushash must change when a new entry is added at a higher height"
        );

        // Clean up
        let _ = std::fs::remove_dir_all(&confirmed_path);
        let _ = std::fs::remove_dir_all(&mempool_path);
    }
}