use std::collections::HashMap;
use std::time::{Duration, Instant, SystemTime};
use anyhow::Result;
use futures::StreamExt;
use tokio::sync::mpsc;
use yellowstone_grpc_client::{GeyserGrpcClient, ClientTlsConfig};
use yellowstone_grpc_proto::geyser::{
    subscribe_update::UpdateOneof, CommitmentLevel, SubscribeRequest,
    SubscribeRequestFilterAccounts, SubscribeRequestFilterSlots,
};

use super::config::ConnectionParams;
use super::types::{AccountUpdate, SlotStatus, SlotUpdate};

pub struct StreamClient {
    params: ConnectionParams,
    track_accounts: bool,
    owner_filter: Option<String>,
}

impl StreamClient {
    pub fn new(params: ConnectionParams, track_accounts: bool, owner_filter: Option<String>) -> Self {
        Self { params, track_accounts, owner_filter }
    }

    async fn build_connection(&self) -> Result<GeyserGrpcClient<impl yellowstone_grpc_client::Interceptor>> {
        let mut builder = GeyserGrpcClient::build_from_shared(self.params.url.clone())?;

        if let Some(ref token) = self.params.token {
            builder = builder.x_token(Some(token.clone()))?;
        }

        builder = builder
            .max_decoding_message_size(self.params.max_msg_size)
            .connect_timeout(self.params.connection_timeout)
            .timeout(self.params.request_timeout)
            .http2_adaptive_window(self.params.adaptive_window)
            .tcp_nodelay(self.params.no_delay);

        if let Some(interval) = self.params.keepalive_interval {
            builder = builder.http2_keep_alive_interval(interval);
        }
        if let Some(size) = self.params.conn_window_size {
            builder = builder.initial_connection_window_size(size);
        }
        if let Some(size) = self.params.stream_window_size {
            builder = builder.initial_stream_window_size(size);
        }
        if let Some(keepalive) = self.params.tcp_keepalive {
            builder = builder.tcp_keepalive(Some(keepalive));
        }
        if let Some(size) = self.params.channel_buffer {
            builder = builder.buffer_size(size);
        }

        if self.params.tls_enabled {
            builder = builder.tls_config(ClientTlsConfig::new().with_native_roots())?;
        }

        Ok(builder.connect().await?)
    }

    pub async fn get_version(&self) -> Result<String> {
        let mut client = self.build_connection().await?;
        let response = client.get_version().await?;

        // Try to parse as JSON and extract version, otherwise return raw string
        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&response.version) {
            if let Some(ver) = json.get("version").and_then(|v| v.get("version")).and_then(|v| v.as_str()) {
                return Ok(ver.to_string());
            }
        }
        Ok(response.version)
    }

    pub async fn measure_ping(&self, samples: usize) -> Result<Vec<Duration>> {
        let mut client = self.build_connection().await?;
        let mut latencies = Vec::with_capacity(samples);

        for _ in 0..samples {
            let start = Instant::now();
            let _ = client.get_version().await?;
            latencies.push(start.elapsed());
        }

        Ok(latencies)
    }

    pub async fn stream_slots(
        &self,
        slot_tx: mpsc::Sender<SlotUpdate>,
        account_tx: Option<mpsc::Sender<AccountUpdate>>,
        mut stop_rx: tokio::sync::watch::Receiver<bool>,
    ) -> Result<()> {
        let mut client = self.build_connection().await?;

        let mut slots_filter = HashMap::new();
        slots_filter.insert(
            "slots".to_string(),
            SubscribeRequestFilterSlots {
                filter_by_commitment: Some(false),
                interslot_updates: Some(true),
            },
        );

        let mut accounts_filter = HashMap::new();
        if self.track_accounts && account_tx.is_some() {
            let owners = self.owner_filter.clone().map(|o| vec![o]).unwrap_or_default();
            accounts_filter.insert(
                "accounts".to_string(),
                SubscribeRequestFilterAccounts {
                    account: vec![],
                    owner: owners,
                    filters: vec![],
                    nonempty_txn_signature: Some(true),
                },
            );
        }

        let request = SubscribeRequest {
            slots: slots_filter,
            accounts: accounts_filter,
            transactions: HashMap::new(),
            transactions_status: HashMap::new(),
            blocks: HashMap::new(),
            blocks_meta: HashMap::new(),
            entry: HashMap::new(),
            commitment: Some(CommitmentLevel::Processed as i32),
            accounts_data_slice: vec![],
            ping: None,
            from_slot: None,
        };

        let (_, mut stream) = client.subscribe_with_request(Some(request)).await?;

        loop {
            tokio::select! {
                _ = stop_rx.changed() => {
                    if *stop_rx.borrow() { break; }
                }
                msg = stream.next() => {
                    match msg {
                        Some(Ok(update)) => {
                            let now = Instant::now();
                            let wall = SystemTime::now();

                            if let Some(payload) = update.update_oneof {
                                match payload {
                                    UpdateOneof::Slot(s) => {
                                        if let Some(status) = SlotStatus::from_proto(s.status) {
                                            let msg = SlotUpdate {
                                                slot: s.slot,
                                                status,
                                                timestamp: now,
                                                wall_time: wall,
                                                parent: s.parent,
                                            };
                                            if slot_tx.send(msg).await.is_err() { break; }
                                        }
                                    }
                                    UpdateOneof::Account(a) => {
                                        if let Some(ref tx) = account_tx {
                                            if let Some(acct) = a.account {
                                                let pubkey = bs58::encode(&acct.pubkey).into_string();
                                                let sig = acct.txn_signature.as_ref()
                                                    .map(|s| bs58::encode(s).into_string());
                                                let msg = AccountUpdate {
                                                    slot: a.slot,
                                                    pubkey,
                                                    write_version: acct.write_version,
                                                    txn_signature: sig,
                                                    timestamp: now,
                                                    wall_time: wall,
                                                };
                                                if tx.send(msg).await.is_err() { break; }
                                            }
                                        }
                                    }
                                    _ => {}
                                }
                            }
                        }
                        Some(Err(e)) => return Err(anyhow::anyhow!("stream error: {}", e)),
                        None => break,
                    }
                }
            }
        }

        Ok(())
    }
}

pub async fn ping_endpoint(endpoint: &str, x_token: Option<&str>) -> Result<Duration> {
    let start = Instant::now();

    let builder = GeyserGrpcClient::build_from_shared(endpoint.to_string())?
        .x_token(x_token.map(|s| s.to_string()))?
        .tls_config(ClientTlsConfig::new().with_native_roots())?;

    let mut client = builder.connect().await?;
    let _ = client.get_version().await?;

    Ok(start.elapsed())
}
