use barter_data::{
    streams::{
        builder::dynamic::indexed::init_indexed_multi_exchange_market_stream,
        reconnect::Event,
    },
    subscription::SubKind,
};
use barter_instrument::{
    Underlying,
    asset::{Asset, name::AssetNameExchange},
    exchange::ExchangeId,
    index::IndexedInstruments,
    instrument::{
        Instrument,
        name::{InstrumentNameExchange, InstrumentNameInternal},
    },
};
use futures::StreamExt;
use sqlx::PgPool;
use std::collections::{HashMap, HashSet};
use tokio::sync::mpsc;
use tracing::{error, info, warn};

use crate::ingest::ResolvedEvent;
use swym_dal::models::market_subscription::ActiveSubscriptionRow;
use swym_dal::repo::market_subscription;

pub struct SubscriptionManager {
    pool: PgPool,
    event_tx: mpsc::Sender<ResolvedEvent>,
    /// Set of (instrument_id, sub_kind) keys that already have a live stream.
    live: HashSet<(i32, String)>,
}

impl SubscriptionManager {
    pub fn new(pool: PgPool, event_tx: mpsc::Sender<ResolvedEvent>) -> Self {
        Self {
            pool,
            event_tx,
            live: HashSet::new(),
        }
    }

    /// Poll DB for active subscriptions, diff against live set, and spawn streams
    /// for any new subscriptions found.
    pub async fn poll_and_spawn(&mut self) {
        let rows = match market_subscription::list_active(&self.pool).await {
            Ok(r) => r,
            Err(e) => {
                error!(?e, "failed to poll active subscriptions");
                return;
            }
        };

        // Find subscriptions not yet live
        let new_rows: Vec<&ActiveSubscriptionRow> = rows
            .iter()
            .filter(|r| !self.live.contains(&(r.instrument_id, r.sub_kind.clone())))
            .collect();

        if new_rows.is_empty() {
            return;
        }

        // Collect unique instruments and sub_kinds from the new batch
        let mut instruments_by_id: HashMap<i32, &ActiveSubscriptionRow> = HashMap::new();
        let mut sub_kind_strs: HashSet<&str> = HashSet::new();

        for row in &new_rows {
            instruments_by_id.entry(row.instrument_id).or_insert(row);
            sub_kind_strs.insert(&row.sub_kind);
        }

        // Parse sub_kinds into barter SubKind
        let sub_kinds: Vec<SubKind> = sub_kind_strs
            .iter()
            .filter_map(|s| match *s {
                "PublicTrades" => Some(SubKind::PublicTrades),
                "OrderBooksL1" => Some(SubKind::OrderBooksL1),
                other => {
                    warn!(kind = other, "unsupported sub_kind, skipping");
                    None
                }
            })
            .collect();

        if sub_kinds.is_empty() {
            return;
        }

        // Build Instrument<ExchangeId, Asset> for each unique instrument
        // and a lookup to map back from (exchange_str, name_exchange) -> db instrument_id
        let mut barter_instruments: Vec<Instrument<ExchangeId, Asset>> = Vec::new();
        let mut lookup: HashMap<(String, String), i32> = HashMap::new();

        for (&db_id, row) in &instruments_by_id {
            let exchange_id: ExchangeId = match serde_json::from_value(
                serde_json::Value::String(row.exchange_name.clone()),
            ) {
                Ok(id) => id,
                Err(e) => {
                    error!(exchange = %row.exchange_name, ?e, "cannot parse exchange from DB");
                    continue;
                }
            };

            let instrument = Instrument::spot(
                exchange_id,
                InstrumentNameInternal::new_from_exchange(
                    exchange_id,
                    InstrumentNameExchange::new(&row.name_exchange),
                ),
                InstrumentNameExchange::new(&row.name_exchange),
                Underlying::new(
                    Asset::new_from_exchange(AssetNameExchange::new(&row.base_asset)),
                    Asset::new_from_exchange(AssetNameExchange::new(&row.quote_asset)),
                ),
                None,
            );

            lookup.insert(
                (exchange_id.as_str().to_owned(), row.name_exchange.clone()),
                db_id,
            );
            barter_instruments.push(instrument);
        }

        if barter_instruments.is_empty() {
            return;
        }

        // Build IndexedInstruments and map InstrumentIndex -> DB id
        let indexed = IndexedInstruments::new(barter_instruments);

        let mut index_to_db_id: HashMap<usize, i32> = HashMap::new();
        for keyed in indexed.instruments() {
            let idx = keyed.key.index();
            let exchange_str = keyed.value.exchange.value.as_str().to_owned();
            let name_exchange: &str = keyed.value.name_exchange.as_ref();

            if let Some(&db_id) = lookup.get(&(exchange_str, name_exchange.to_owned())) {
                index_to_db_id.insert(idx, db_id);
            }
        }

        // Init market stream for this batch
        let mut stream =
            match init_indexed_multi_exchange_market_stream(&indexed, &sub_kinds).await {
                Ok(s) => s,
                Err(e) => {
                    error!(?e, "failed to init market stream for new subscriptions");
                    return;
                }
            };

        // Mark all new subscriptions as live
        for row in &new_rows {
            self.live
                .insert((row.instrument_id, row.sub_kind.clone()));
        }

        let instrument_count = instruments_by_id.len();
        let sub_kind_count = sub_kinds.len();
        info!(
            instruments = instrument_count,
            sub_kinds = sub_kind_count,
            "spawned stream for new subscriptions"
        );

        // Spawn a task to consume the stream and forward resolved events
        let tx = self.event_tx.clone();
        tokio::spawn(async move {
            while let Some(event) = stream.next().await {
                match event {
                    Event::Item(market_event) => {
                        let Some(&db_id) =
                            index_to_db_id.get(&market_event.instrument.index())
                        else {
                            continue;
                        };
                        let resolved = ResolvedEvent {
                            instrument_id: db_id,
                            time_exchange: market_event.time_exchange,
                            time_received: market_event.time_received,
                            kind: market_event.kind,
                        };
                        if tx.send(resolved).await.is_err() {
                            break; // receiver dropped, shut down
                        }
                    }
                    Event::Reconnecting(exchange) => {
                        warn!(?exchange, "stream reconnecting");
                    }
                }
            }
            info!("subscription stream task ended");
        });
    }
}
