//! Polymarket orderbook WebSocket feed.
//!
//! Connects to Polymarket's CLOB WebSocket and streams
//! best bid/ask + mid price into a shared FeedSnapshot,
//! plus full L2 depth into a shared OrderbookSnapshot.

use crate::feeds::{FeedSnapshot, MetricsStore, OrderbookStore};
use crate::orderbook::OrderbookSnapshot;
use dashmap::DashMap;
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio_tungstenite::connect_async;
use tracing::{debug, warn};

const POLYMARKET_WS_URL: &str = "wss://ws-subscriptions-clob.polymarket.com/ws/market";

/// Run the Polymarket orderbook feed until shutdown is signaled.
pub async fn run_polymarket_book(
    name: String,
    market_slug: String,
    snapshots: Arc<DashMap<String, FeedSnapshot>>,
    orderbooks: OrderbookStore,
    metrics: MetricsStore,
    global_shutdown: Arc<tokio::sync::Notify>,
    feed_shutdown: Arc<tokio::sync::Notify>,
) {
    // Pre-compute source string to avoid allocation per message
    let source = format!("polymarket:{}", market_slug);
    let mut reconnect_delay_ms: u64 = 1000;

    loop {
        let connect_result = tokio::select! {
            _ = global_shutdown.notified() => return,
            _ = feed_shutdown.notified() => return,
            result = connect_async(POLYMARKET_WS_URL) => result,
        };

        match connect_result {
            Ok((ws_stream, _)) => {
                reconnect_delay_ms = 1000; // Reset on successful connect
                let (mut write, mut read) = ws_stream.split();

                // Subscribe to the market's orderbook channel (re-subscribes on every reconnect)
                // Polymarket CLOB WS expects: assets_ids (plural, array), type "market"
                let subscribe_msg = serde_json::json!({
                    "assets_ids": [&market_slug],
                    "type": "market",
                });

                if let Err(e) = write
                    .send(tokio_tungstenite::tungstenite::Message::Text(
                        subscribe_msg.to_string(),
                    ))
                    .await
                {
                    warn!(
                        market = %market_slug,
                        error = %e,
                        "polymarket feed: subscribe failed"
                    );
                    if let Some(mut m) = metrics.get_mut(&name) {
                        m.error_count += 1;
                        m.last_error = e.to_string();
                    }
                    // Fall through to reconnect
                } else {
                    debug!(market = %market_slug, "polymarket feed: subscribed");
                    if let Some(mut m) = metrics.get_mut(&name) {
                        m.is_connected = true;
                    }

                    loop {
                        let msg = tokio::select! {
                            _ = global_shutdown.notified() => return,
                            _ = feed_shutdown.notified() => return,
                            msg = read.next() => msg,
                            _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {
                                tracing::warn!("polymarket feed '{}': idle timeout, reconnecting", name);
                                break;
                            }
                        };

                        match msg {
                            Some(Ok(tokio_tungstenite::tungstenite::Message::Text(text))) => {
                                if let Ok(json) =
                                    serde_json::from_str::<serde_json::Value>(&text)
                                {
                                    // Parse the message based on its format
                                    let (best_bid, best_ask, bid_levels, ask_levels) =
                                        parse_polymarket_msg(&json, &market_slug);

                                    // Also check for a direct price field
                                    let direct_price = json
                                        .get("price")
                                        .or_else(|| json.get("last_trade_price"))
                                        .and_then(|v| {
                                            v.as_f64()
                                                .or_else(|| v.as_str().and_then(|s| s.parse().ok()))
                                        });

                                    let has_book = best_bid > 0.0 || best_ask > 0.0;
                                    let has_price = direct_price.is_some();

                                    if has_book || has_price {
                                        let mid = if best_bid > 0.0 && best_ask > 0.0 {
                                            (best_bid + best_ask) / 2.0
                                        } else if let Some(p) = direct_price {
                                            p
                                        } else if best_bid > 0.0 {
                                            best_bid
                                        } else {
                                            best_ask
                                        };

                                        let now = std::time::SystemTime::now()
                                            .duration_since(std::time::UNIX_EPOCH)
                                            .unwrap_or_default()
                                            .as_secs_f64();

                                        let snap = FeedSnapshot {
                                            price: mid,
                                            timestamp: now,
                                            source: source.clone(),
                                            bid: if best_bid > 0.0 {
                                                best_bid
                                            } else {
                                                direct_price.unwrap_or(0.0)
                                            },
                                            ask: if best_ask > 0.0 {
                                                best_ask
                                            } else {
                                                direct_price.unwrap_or(0.0)
                                            },
                                            volume_24h: 0.0,
                                            last_trade_size: 0.0,
                                            last_trade_is_buy: false,
                                        };

                                        snapshots.insert(name.clone(), snap);
                                        if let Some(mut m) = metrics.get_mut(&name) {
                                            m.update_count += 1;
                                            m.last_update_time = now;
                                        }

                                        // Store L2 orderbook if we have depth data
                                        if !bid_levels.is_empty() || !ask_levels.is_empty() {
                                            let book = OrderbookSnapshot::from_levels(
                                                bid_levels,
                                                ask_levels,
                                                now,
                                                source.clone(),
                                            );
                                            orderbooks.insert(name.clone(), book);
                                        }
                                    }
                                }
                            }
                            Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_))) => {
                                if let Some(mut m) = metrics.get_mut(&name) {
                                    m.is_connected = false;
                                }
                                break;
                            }
                            Some(Err(e)) => {
                                warn!(
                                    market = %market_slug,
                                    error = %e,
                                    "polymarket feed: connection error"
                                );
                                if let Some(mut m) = metrics.get_mut(&name) {
                                    m.error_count += 1;
                                    m.last_error = e.to_string();
                                    m.is_connected = false;
                                }
                                break;
                            }
                            None => break,
                            _ => {} // Ping/Pong/Binary — ignore
                        }
                    }
                }
            }
            Err(e) => {
                warn!(
                    market = %market_slug,
                    error = %e,
                    "polymarket feed: connection failed"
                );
                if let Some(mut m) = metrics.get_mut(&name) {
                    m.error_count += 1;
                    m.last_error = e.to_string();
                }
            }
        }

        // Reconnect with exponential backoff (1s, 2s, 4s, 8s, ... max 30s)
        if let Some(mut m) = metrics.get_mut(&name) {
            m.reconnect_count += 1;
            m.is_connected = false;
        }
        warn!(market = %market_slug, delay_ms = reconnect_delay_ms, "polymarket feed: reconnecting");
        tokio::select! {
            _ = global_shutdown.notified() => return,
            _ = feed_shutdown.notified() => return,
            _ = tokio::time::sleep(tokio::time::Duration::from_millis(reconnect_delay_ms)) => {},
        }
        reconnect_delay_ms = (reconnect_delay_ms * 2).min(30_000);
    }
}

/// Parse a Polymarket WS message into best bid/ask and L2 levels.
///
/// Handles three message formats:
/// 1. Initial book snapshot: `[{bids: [...], asks: [...], ...}]` (JSON array)
/// 2. Price change: `{price_changes: [{best_bid, best_ask, ...}], ...}`
/// 3. Direct book: `{bids: [...], asks: [...], ...}` (fallback)
fn parse_polymarket_msg(
    json: &serde_json::Value,
    token_id: &str,
) -> (f64, f64, Vec<(f64, f64)>, Vec<(f64, f64)>) {
    // Format 1: Initial snapshot is a JSON array wrapping book objects
    if let Some(arr) = json.as_array() {
        for item in arr {
            // Filter: only parse snapshots for our subscribed token
            let asset_match = item
                .get("asset_id")
                .and_then(|v| v.as_str())
                .map(|a| a == token_id)
                .unwrap_or(true); // If no asset_id field, accept it
            if !asset_match {
                continue;
            }
            let bid_levels = extract_all_levels(item, "bids");
            let ask_levels = extract_all_levels(item, "asks");
            if !bid_levels.is_empty() || !ask_levels.is_empty() {
                let best_bid = bid_levels.first().map(|&(p, _)| p).unwrap_or(0.0);
                let best_ask = ask_levels.first().map(|&(p, _)| p).unwrap_or(0.0);
                return (best_bid, best_ask, bid_levels, ask_levels);
            }
        }
        return (0.0, 0.0, Vec::new(), Vec::new());
    }

    // Format 2: Price change messages with best_bid/best_ask
    if let Some(changes) = json.get("price_changes").and_then(|v| v.as_array()) {
        for change in changes {
            // Filter: only use price changes for our subscribed token
            let asset_match = change
                .get("asset_id")
                .and_then(|v| v.as_str())
                .map(|a| a == token_id)
                .unwrap_or(false);
            if !asset_match {
                continue;
            }
            let best_bid = parse_num(change.get("best_bid")).unwrap_or(0.0);
            let best_ask = parse_num(change.get("best_ask")).unwrap_or(0.0);
            if best_bid > 0.0 || best_ask > 0.0 {
                return (best_bid, best_ask, Vec::new(), Vec::new());
            }
        }
        return (0.0, 0.0, Vec::new(), Vec::new());
    }

    // Format 3: Direct bids/asks at top level
    let bid_levels = extract_all_levels(json, "bids");
    let ask_levels = extract_all_levels(json, "asks");
    let best_bid = bid_levels.first().map(|&(p, _)| p).unwrap_or(0.0);
    let best_ask = ask_levels.first().map(|&(p, _)| p).unwrap_or(0.0);
    (best_bid, best_ask, bid_levels, ask_levels)
}

/// Parse a JSON value as an f64 (handles both number and string formats).
fn parse_num(val: Option<&serde_json::Value>) -> Option<f64> {
    val.and_then(|v| v.as_f64().or_else(|| v.as_str().and_then(|s| s.parse().ok())))
        .filter(|f| f.is_finite())
}

/// Extract all price levels from an orderbook side.
/// Returns Vec<(price, size)>, sorted descending for bids, ascending for asks.
fn extract_all_levels(json: &serde_json::Value, side: &str) -> Vec<(f64, f64)> {
    let arr = match json.get(side).and_then(|v| v.as_array()) {
        Some(a) => a,
        None => return Vec::new(),
    };

    let mut levels: Vec<(f64, f64)> = arr
        .iter()
        .filter_map(|entry| {
            // Format 1: {"price": "0.55", "size": "100"}
            let price = entry
                .get("price")
                .and_then(|p| {
                    p.as_f64()
                        .or_else(|| p.as_str().and_then(|s| s.parse::<f64>().ok()))
                })
                .or_else(|| {
                    // Format 2: [price, size]
                    entry
                        .as_array()
                        .and_then(|pair| pair.first())
                        .and_then(|p| {
                            p.as_f64()
                                .or_else(|| p.as_str().and_then(|s| s.parse::<f64>().ok()))
                        })
                })?;

            let size = entry
                .get("size")
                .and_then(|s| {
                    s.as_f64()
                        .or_else(|| s.as_str().and_then(|s| s.parse::<f64>().ok()))
                })
                .or_else(|| {
                    entry
                        .as_array()
                        .and_then(|pair| pair.get(1))
                        .and_then(|s| {
                            s.as_f64()
                                .or_else(|| s.as_str().and_then(|s| s.parse::<f64>().ok()))
                        })
                })
                .unwrap_or(0.0);

            if price > 0.0 {
                Some((price, size))
            } else {
                None
            }
        })
        .collect();

    // Sort: bids descending, asks ascending
    if side == "bids" {
        levels.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
    } else {
        levels.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
    }

    levels
}
