// file: kb_lib/src/ws_hybrid_watch.rs

//! Hybrid WebSocket watch-target collection service.

/// Kind of hybrid WebSocket watch target.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum WsWatchTargetKind {
    /// Program-level watch target for `programSubscribe`.
    Program,
    /// Account-level watch target for `accountSubscribe`.
    Account,
}

/// One hybrid WebSocket watch target.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WsWatchTarget {
    /// Target kind.
    pub kind: crate::WsWatchTargetKind,
    /// Base-58 address to watch.
    pub address: std::string::String,
    /// Logical source label.
    pub logical_source: std::string::String,
}

/// Snapshot of hybrid WebSocket watch targets.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WsHybridWatchSnapshot {
    /// Program targets suitable for `programSubscribe`.
    pub program_targets: std::vec::Vec<crate::WsWatchTarget>,
    /// Account targets suitable for `accountSubscribe`.
    pub account_targets: std::vec::Vec<crate::WsWatchTarget>,
}

/// Hybrid WebSocket watch-target collection service.
#[derive(Debug, Clone)]
pub struct WsHybridWatchService {
    database: std::sync::Arc<crate::Database>,
}

impl WsHybridWatchService {
    /// Creates a new hybrid watch-target collection service.
    pub fn new(database: std::sync::Arc<crate::Database>) -> Self {
        return Self { database };
    }

    /// Collects the deduplicated set of DEX program targets.
    pub async fn collect_program_targets(
        &self,
    ) -> Result<std::vec::Vec<crate::WsWatchTarget>, crate::Error> {
        let dexes_result = crate::query_dexs_list(self.database.as_ref()).await;
        let dexes = match dexes_result {
            Ok(dexes) => dexes,
            Err(error) => return Err(error),
        };
        let mut seen = std::collections::HashSet::<std::string::String>::new();
        let mut targets = std::vec::Vec::new();
        for dex in dexes {
            if !dex.is_enabled {
                continue;
            }
            if let Some(program_id) = dex.program_id.clone() {
                if !program_id.is_empty() && !seen.contains(&program_id) {
                    seen.insert(program_id.clone());
                    targets.push(crate::WsWatchTarget {
                        kind: crate::WsWatchTargetKind::Program,
                        address: program_id,
                        logical_source: format!("dex:{}:program_id", dex.code),
                    });
                }
            }
            if let Some(router_program_id) = dex.router_program_id.clone() {
                if !router_program_id.is_empty() && !seen.contains(&router_program_id) {
                    seen.insert(router_program_id.clone());
                    targets.push(crate::WsWatchTarget {
                        kind: crate::WsWatchTargetKind::Program,
                        address: router_program_id,
                        logical_source: format!("dex:{}:router_program_id", dex.code),
                    });
                }
            }
        }
        targets.sort_by(|left, right| return left.address.cmp(&right.address));
        return Ok(targets);
    }

    /// Collects the deduplicated set of active pool-account targets.
    pub async fn collect_account_targets(
        &self,
    ) -> Result<std::vec::Vec<crate::WsWatchTarget>, crate::Error> {
        let pools_result = crate::query_pools_list(self.database.as_ref()).await;
        let pools = match pools_result {
            Ok(pools) => pools,
            Err(error) => return Err(error),
        };
        let mut seen = std::collections::HashSet::<std::string::String>::new();
        let mut targets = std::vec::Vec::new();
        for pool in pools {
            if pool.status != crate::PoolStatus::Active {
                continue;
            }
            if pool.address.is_empty() {
                continue;
            }
            if seen.contains(&pool.address) {
                continue;
            }
            seen.insert(pool.address.clone());
            targets.push(crate::WsWatchTarget {
                kind: crate::WsWatchTargetKind::Account,
                address: pool.address.clone(),
                logical_source: format!("pool:{}", pool.dex_id),
            });
        }
        targets.sort_by(|left, right| return left.address.cmp(&right.address));
        return Ok(targets);
    }

    /// Collects both program and account targets in one snapshot.
    pub async fn collect_snapshot(&self) -> Result<crate::WsHybridWatchSnapshot, crate::Error> {
        let program_targets_result = self.collect_program_targets().await;
        let program_targets = match program_targets_result {
            Ok(program_targets) => program_targets,
            Err(error) => return Err(error),
        };
        let account_targets_result = self.collect_account_targets().await;
        let account_targets = match account_targets_result {
            Ok(account_targets) => account_targets,
            Err(error) => return Err(error),
        };
        return Ok(crate::WsHybridWatchSnapshot { program_targets, account_targets });
    }
}

#[cfg(test)]
mod tests {
    async fn make_database() -> std::sync::Arc<crate::Database> {
        let tempdir_result = tempfile::tempdir();
        let tempdir = match tempdir_result {
            Ok(tempdir) => tempdir,
            Err(error) => panic!("tempdir must succeed: {}", error),
        };
        let database_path = tempdir.path().join("ws_hybrid_watch.sqlite3");
        let config = crate::DatabaseConfig {
            enabled: true,
            backend: crate::DatabaseBackend::Sqlite,
            sqlite: crate::SqliteDatabaseConfig {
                path: database_path.to_string_lossy().to_string(),
                create_if_missing: true,
                busy_timeout_ms: 5000,
                max_connections: 1,
                auto_initialize_schema: true,
                use_wal: true,
            },
        };
        let database_result = crate::Database::connect_and_initialize(&config).await;
        let database = match database_result {
            Ok(database) => database,
            Err(error) => panic!("database init must succeed: {}", error),
        };
        return std::sync::Arc::new(database);
    }
    async fn seed_dexes_and_pools(database: std::sync::Arc<crate::Database>) {
        let dex_a = crate::DexDto::new(
            "raydium_amm_v4".to_string(),
            "Raydium AmmV4".to_string(),
            Some("RaydiumProgram111".to_string()),
            None,
            true,
        );
        let dex_a_id_result = crate::query_dexs_upsert(database.as_ref(), &dex_a).await;
        let dex_a_id = match dex_a_id_result {
            Ok(dex_a_id) => dex_a_id,
            Err(error) => panic!("dex A upsert must succeed: {}", error),
        };
        let dex_b = crate::DexDto::new(
            "meteora_dbc".to_string(),
            "Meteora DBC".to_string(),
            Some("MeteoraProgram111".to_string()),
            Some("SharedRouter111".to_string()),
            true,
        );
        let dex_b_id_result = crate::query_dexs_upsert(database.as_ref(), &dex_b).await;
        let dex_b_id = match dex_b_id_result {
            Ok(dex_b_id) => dex_b_id,
            Err(error) => panic!("dex B upsert must succeed: {}", error),
        };
        let dex_c = crate::DexDto::new(
            "pump_swap".to_string(),
            "PumpSwap".to_string(),
            Some("PumpSwapProgram111".to_string()),
            Some("SharedRouter111".to_string()),
            true,
        );
        let dex_c_id_result = crate::query_dexs_upsert(database.as_ref(), &dex_c).await;
        let dex_c_id = match dex_c_id_result {
            Ok(dex_c_id) => dex_c_id,
            Err(error) => panic!("dex C upsert must succeed: {}", error),
        };
        let pool_a = crate::PoolDto::new(
            dex_a_id,
            "PoolAddress111".to_string(),
            crate::PoolKind::Amm,
            crate::PoolStatus::Active,
        );
        let pool_a_result = crate::query_pools_upsert(database.as_ref(), &pool_a).await;
        if let Err(error) = pool_a_result {
            panic!("pool A upsert must succeed: {}", error);
        }
        let pool_b = crate::PoolDto::new(
            dex_b_id,
            "PoolAddress222".to_string(),
            crate::PoolKind::Amm,
            crate::PoolStatus::Active,
        );
        let pool_b_result = crate::query_pools_upsert(database.as_ref(), &pool_b).await;
        if let Err(error) = pool_b_result {
            panic!("pool B upsert must succeed: {}", error);
        }
        let pool_c = crate::PoolDto::new(
            dex_c_id,
            "PoolAddress333".to_string(),
            crate::PoolKind::Amm,
            crate::PoolStatus::Inactive,
        );
        let pool_c_result = crate::query_pools_upsert(database.as_ref(), &pool_c).await;
        if let Err(error) = pool_c_result {
            panic!("pool C upsert must succeed: {}", error);
        }
    }

    #[tokio::test]
    async fn collect_program_targets_deduplicates_program_and_router_ids() {
        let database = make_database().await;
        seed_dexes_and_pools(database.clone()).await;
        let service = crate::WsHybridWatchService::new(database);
        let targets_result = service.collect_program_targets().await;
        let targets = match targets_result {
            Ok(targets) => targets,
            Err(error) => panic!("collect program targets must succeed: {}", error),
        };
        assert_eq!(targets.len(), 4);
        let mut addresses = std::vec::Vec::new();
        for target in &targets {
            assert_eq!(target.kind, crate::WsWatchTargetKind::Program);
            addresses.push(target.address.clone());
        }
        addresses.sort();
        assert_eq!(
            addresses,
            vec![
                "MeteoraProgram111".to_string(),
                "PumpSwapProgram111".to_string(),
                "RaydiumProgram111".to_string(),
                "SharedRouter111".to_string(),
            ]
        );
    }

    #[tokio::test]
    async fn collect_account_targets_only_keeps_active_pools() {
        let database = make_database().await;
        seed_dexes_and_pools(database.clone()).await;
        let service = crate::WsHybridWatchService::new(database);
        let targets_result = service.collect_account_targets().await;
        let targets = match targets_result {
            Ok(targets) => targets,
            Err(error) => panic!("collect account targets must succeed: {}", error),
        };
        assert_eq!(targets.len(), 2);
        let mut addresses = std::vec::Vec::new();
        for target in &targets {
            assert_eq!(target.kind, crate::WsWatchTargetKind::Account);
            addresses.push(target.address.clone());
        }
        addresses.sort();
        assert_eq!(addresses, vec!["PoolAddress111".to_string(), "PoolAddress222".to_string(),]);
    }
}
