use std::collections::HashSet;
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::Result;
use indicatif::{ProgressBar, ProgressStyle};
use tokio::sync::{mpsc, watch, Barrier};

use super::client::StreamClient;
use super::config::{ConnectionParams, GeyserConfig};
use super::types::{AccountUpdate, EndpointData, SlotStatus, SlotUpdate};

pub struct ProviderMeta {
    pub plugin_name: String,
    pub version: String,
    pub avg_ping: Duration,
}

pub struct DualEndpointCollector {
    cfg: GeyserConfig,
    url1: String,
    url2: String,
    token1: Option<String>,
    token2: Option<String>,
    target_slots: usize,
    track_accounts: bool,
    owner_filter: Option<String>,
}

impl DualEndpointCollector {
    pub fn new(
        cfg: GeyserConfig,
        url1: String,
        url2: String,
        token1: Option<String>,
        token2: Option<String>,
        target_slots: usize,
        track_accounts: bool,
        owner_filter: Option<String>,
    ) -> Self {
        Self {
            cfg,
            url1,
            url2,
            token1,
            token2,
            target_slots,
            track_accounts,
            owner_filter,
        }
    }

    pub async fn collect(
        &self,
        verbose: bool,
    ) -> Result<(EndpointData, EndpointData, ProviderMeta, ProviderMeta, f64)> {
        let start = Instant::now();
        let extra_pct = self.cfg.benchmark.extra_slots_pct;
        let buffer_slots = (self.target_slots as f32 * (1.0 + extra_pct)) as usize;

        let params1 = ConnectionParams::from_config(
            self.url1.clone(),
            self.token1.clone(),
            &self.cfg,
        );
        let params2 = ConnectionParams::from_config(
            self.url2.clone(),
            self.token2.clone(),
            &self.cfg,
        );

        let client1 = StreamClient::new(params1.clone(), self.track_accounts, self.owner_filter.clone());
        let client2 = StreamClient::new(params2.clone(), self.track_accounts, self.owner_filter.clone());

        if verbose { println!("fetching provider versions..."); }

        let (ver1, ver2) = tokio::try_join!(
            client1.get_version(),
            client2.get_version()
        )?;

        if verbose {
            println!("  provider 1: {}", ver1);
            println!("  provider 2: {}", ver2);
        }

        if verbose { println!("measuring latency ({} samples)...", self.cfg.benchmark.ping_samples); }

        let (pings1, pings2) = tokio::try_join!(
            client1.measure_ping(self.cfg.benchmark.ping_samples),
            client2.measure_ping(self.cfg.benchmark.ping_samples)
        )?;

        let avg1 = pings1.iter().sum::<Duration>() / pings1.len() as u32;
        let avg2 = pings2.iter().sum::<Duration>() / pings2.len() as u32;

        if verbose {
            println!("  provider 1 ping: {:.2}ms", avg1.as_secs_f64() * 1000.0);
            println!("  provider 2 ping: {:.2}ms", avg2.as_secs_f64() * 1000.0);
        }

        let meta1 = ProviderMeta {
            plugin_name: "yellowstone".to_string(),
            version: ver1,
            avg_ping: avg1,
        };
        let meta2 = ProviderMeta {
            plugin_name: "yellowstone".to_string(),
            version: ver2,
            avg_ping: avg2,
        };

        let barrier = Arc::new(Barrier::new(2));

        let (stop_tx, stop_rx1) = watch::channel(false);
        let stop_rx2 = stop_tx.subscribe();

        let (slot_tx1, mut slot_rx1) = mpsc::channel::<SlotUpdate>(10000);
        let (slot_tx2, mut slot_rx2) = mpsc::channel::<SlotUpdate>(10000);

        let (acct_tx1, mut acct_rx1) = mpsc::channel::<AccountUpdate>(50000);
        let (acct_tx2, mut acct_rx2) = mpsc::channel::<AccountUpdate>(50000);

        let client1 = StreamClient::new(params1, self.track_accounts, self.owner_filter.clone());
        let client2 = StreamClient::new(params2, self.track_accounts, self.owner_filter.clone());

        let b1 = barrier.clone();
        let track1 = self.track_accounts;
        let handle1 = tokio::spawn(async move {
            b1.wait().await;
            let acct_tx = if track1 { Some(acct_tx1) } else { None };
            client1.stream_slots(slot_tx1, acct_tx, stop_rx1).await
        });

        let b2 = barrier.clone();
        let track2 = self.track_accounts;
        let handle2 = tokio::spawn(async move {
            b2.wait().await;
            let acct_tx = if track2 { Some(acct_tx2) } else { None };
            client2.stream_slots(slot_tx2, acct_tx, stop_rx2).await
        });

        let mut data1 = EndpointData::with_capacity(
            self.url1.clone(),
            self.target_slots,
            extra_pct,
        );
        let mut data2 = EndpointData::with_capacity(
            self.url2.clone(),
            self.target_slots,
            extra_pct,
        );

        let pb = if verbose {
            let pb = ProgressBar::new(buffer_slots as u64);
            if let Ok(style) = ProgressStyle::default_bar()
                .template("{spinner:.green} [{bar:40.cyan/blue}] {pos}/{len} slots ({eta})")
            {
                pb.set_style(style.progress_chars("#>-"));
            }
            Some(pb)
        } else {
            None
        };

        let mut finalized = HashSet::with_capacity(buffer_slots);

        loop {
            tokio::select! {
                Some(u) = slot_rx1.recv() => {
                    if u.status == SlotStatus::Finalized {
                        finalized.insert(u.slot);
                        if let Some(ref pb) = pb { pb.set_position(finalized.len() as u64); }
                    }
                    data1.add_slot_update(u);
                }
                Some(u) = slot_rx2.recv() => {
                    data2.add_slot_update(u);
                }
                Some(u) = acct_rx1.recv() => {
                    data1.add_account_update(u);
                }
                Some(u) = acct_rx2.recv() => {
                    data2.add_account_update(u);
                }
                else => break,
            }

            if finalized.len() >= buffer_slots {
                break;
            }
        }

        if let Some(pb) = pb { pb.finish_with_message("collection complete"); }
        let _ = stop_tx.send(true);

        tokio::time::sleep(Duration::from_millis(500)).await;
        handle1.abort();
        handle2.abort();

        let duration = start.elapsed().as_secs_f64();
        Ok((data1, data2, meta1, meta2, duration))
    }
}
