mod postgres_sink;

use anyhow::Result;
use chrono::{DateTime, Utc};
use clap::{Parser, Subcommand};
use serde::Serialize;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::time::Duration;
use tokio::sync::broadcast;
use tracing::{error, info, warn};
use tracing_subscriber::EnvFilter;
use uforia_ms_book::{BboTracker, CrossVenueBook, L2OrderBook};
use uforia_ms_config::{load_config, AppConfig};
use uforia_ms_feed::manager::FeedManager;
use uforia_ms_model::load_artifact;
use uforia_ms_types::{Asset, BboQuote, FundingEvent, IvEvent, MarketEvent, Venue};

use crate::postgres_sink::{load_health, DbSinkConfig, PostgresSink};

const BOOK_TRACKER_CAPACITY: usize = 600;

#[derive(Parser)]
#[command(name = "uforia-ms-app", about = "Uforia microstructure live engine")]
struct Cli {
    #[arg(
        short,
        long,
        default_value = "engines/microstructure/config/default.toml"
    )]
    config: PathBuf,
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    Run {
        #[arg(long, default_value = "paper")]
        mode: String,
    },
    Health,
    InspectModel {
        #[arg(long)]
        path: PathBuf,
    },
}

#[derive(Debug, Serialize, Default, Clone)]
pub struct HealthAssetSummary {
    books_tracked: usize,
    complete_books: usize,
    partial_books: usize,
    cross_venue_spread: Option<f64>,
    cross_venue_depth_imbalance: Option<f64>,
}

#[derive(Debug, Serialize, Default, Clone)]
pub struct EngineHealth {
    status: String,
    engine: String,
    mode: String,
    updated_at: DateTime<Utc>,
    data_dir: String,
    assets: Vec<String>,
    books_tracked: usize,
    complete_books: usize,
    partial_books: usize,
    event_counts: HashMap<String, u64>,
    last_market_event_ts: Option<DateTime<Utc>>,
    last_funding_event_ts: Option<DateTime<Utc>>,
    last_iv_event_ts: Option<DateTime<Utc>>,
    asset_summaries: HashMap<String, HealthAssetSummary>,
    warnings: Vec<String>,
}

#[derive(Debug, Default)]
pub struct BookRuntimeState {
    pub book: L2OrderBook,
    tracker: BboTracker,
    complete: bool,
    is_seeded: bool,
    is_sequence_valid: bool,
    is_gap_detected: bool,
    reseed_count: u64,
    last_mark_price: Option<f64>,
    last_index_price: Option<f64>,
    last_event_ts: Option<DateTime<Utc>>,
}

impl BookRuntimeState {
    fn new() -> Self {
        Self {
            book: L2OrderBook::new(),
            tracker: BboTracker::new(BOOK_TRACKER_CAPACITY),
            complete: false,
            is_seeded: false,
            is_sequence_valid: false,
            is_gap_detected: false,
            reseed_count: 0,
            last_mark_price: None,
            last_index_price: None,
            last_event_ts: None,
        }
    }

    fn maybe_push_bbo(&mut self, asset: Asset, venue: Venue, timestamp: DateTime<Utc>) {
        if let (Some((bid_price, bid_size)), Some((ask_price, ask_size))) =
            (self.book.best_bid(), self.book.best_ask())
        {
            self.tracker.push(BboQuote {
                timestamp,
                asset,
                venue,
                bid_price,
                bid_size,
                ask_price,
                ask_size,
            });
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let config = load_config(&cli.config)?;
    init_tracing(&config.general.log_level)?;

    match cli.command {
        Commands::Run { mode } => run_engine(config, mode).await,
        Commands::Health => show_health(&config).await,
        Commands::InspectModel { path } => {
            let artifact = load_artifact(&path)?;
            println!("{}", serde_json::to_string_pretty(&artifact)?);
            Ok(())
        }
    }
}

fn init_tracing(log_level: &str) -> Result<()> {
    let filter =
        EnvFilter::try_from_default_env().or_else(|_| EnvFilter::try_new(log_level.to_string()))?;
    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_target(false)
        .without_time()
        .try_init()
        .ok();
    Ok(())
}

async fn run_engine(config: AppConfig, mode: String) -> Result<()> {
    let sink_config = sink_config(&config);
    let mut sink = PostgresSink::connect(&sink_config).await?;
    let mut health = initial_health(&config, &mode);

    let (market_tx, mut market_rx) = broadcast::channel(4096);
    let (funding_tx, mut funding_rx) = broadcast::channel(1024);
    let (iv_tx, mut iv_rx) = broadcast::channel(1024);
    let manager = FeedManager::new(config.clone(), market_tx, funding_tx, iv_tx);
    let feed_handles = manager.start();

    let mut books: HashMap<(Asset, Venue), BookRuntimeState> = HashMap::new();
    let mut cross_venue = CrossVenueBook::new();
    let mut dirty_books: HashSet<(Asset, Venue)> = HashSet::new();

    let tick_ms = config.bar.flush_delay_ms.max(1000);
    let mut flush_interval = tokio::time::interval(Duration::from_millis(tick_ms));
    flush_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    info!(
        mode = %mode,
        schema = %sink_config.schema,
        flush_ms = tick_ms,
        "microstructure live runtime starting"
    );

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {
                info!("ctrl-c received, shutting down microstructure engine");
                break;
            }
            _ = flush_interval.tick() => {
                flush_dirty_books(&mut sink, &books, &mut dirty_books).await?;
                sink.flush().await?;
                refresh_health(&mut health, &books, &cross_venue);
            }
            event = market_rx.recv() => {
                match event {
                    Ok(event) => {
                        handle_market_event(&mut sink, &mut health, &mut books, &mut cross_venue, &mut dirty_books, event).await?;
                    }
                    Err(broadcast::error::RecvError::Lagged(skipped)) => {
                        warn!(skipped, "market channel lagged");
                    }
                    Err(broadcast::error::RecvError::Closed) => {
                        warn!("market channel closed");
                        break;
                    }
                }
            }
            event = funding_rx.recv() => {
                match event {
                    Ok(event) => handle_funding_event(&mut sink, &mut health, event).await?,
                    Err(broadcast::error::RecvError::Lagged(skipped)) => {
                        warn!(skipped, "funding channel lagged");
                    }
                    Err(broadcast::error::RecvError::Closed) => {
                        warn!("funding channel closed");
                        break;
                    }
                }
            }
            event = iv_rx.recv() => {
                match event {
                    Ok(event) => handle_iv_event(&mut sink, &mut health, event).await?,
                    Err(broadcast::error::RecvError::Lagged(skipped)) => {
                        warn!(skipped, "iv channel lagged");
                    }
                    Err(broadcast::error::RecvError::Closed) => {
                        warn!("iv channel closed");
                        break;
                    }
                }
            }
        }
    }

    manager.shutdown();
    for handle in feed_handles {
        if let Err(join_error) = handle.await {
            error!(error = %join_error, "feed task join failed");
        }
    }

    flush_dirty_books(&mut sink, &books, &mut dirty_books).await?;
    sink.flush().await?;
    refresh_health(&mut health, &books, &cross_venue);
    health.status = "stopped".to_string();
    health.updated_at = Utc::now();
    info!("microstructure live runtime stopped");
    Ok(())
}

async fn show_health(config: &AppConfig) -> Result<()> {
    let health = load_health(
        &sink_config(config),
        "microstructure_live_engine",
        "paper",
        &config.general.assets,
    )
    .await
    .unwrap_or_else(|_| initial_health(config, "paper"));
    println!("{}", serde_json::to_string_pretty(&health)?);
    Ok(())
}

fn initial_health(config: &AppConfig, mode: &str) -> EngineHealth {
    EngineHealth {
        status: "online".to_string(),
        engine: "microstructure_live_engine".to_string(),
        mode: mode.to_string(),
        updated_at: Utc::now(),
        data_dir: "postgres".to_string(),
        assets: config.general.assets.clone(),
        event_counts: HashMap::from([
            ("market".to_string(), 0),
            ("funding".to_string(), 0),
            ("iv".to_string(), 0),
        ]),
        warnings: vec![
            "Books are downgraded whenever snapshot seeding is missing or sequence validation detects a gap.".to_string(),
        ],
        ..EngineHealth::default()
    }
}

fn sink_config(config: &AppConfig) -> DbSinkConfig {
    DbSinkConfig {
        dsn: config.storage.postgres.dsn.clone(),
        schema: config.storage.postgres.schema.clone(),
        buffer_size: config.storage.postgres.buffer_size,
    }
}

async fn handle_market_event(
    sink: &mut PostgresSink,
    health: &mut EngineHealth,
    books: &mut HashMap<(Asset, Venue), BookRuntimeState>,
    cross_venue: &mut CrossVenueBook,
    dirty_books: &mut HashSet<(Asset, Venue)>,
    event: MarketEvent,
) -> Result<()> {
    *health.event_counts.entry("market".to_string()).or_default() += 1;
    let received_at = Utc::now();

    match &event {
        MarketEvent::Trade(trade) => {
            health.last_market_event_ts = Some(trade.timestamp);
        }
        MarketEvent::L2Snapshot(snapshot) => {
            let key = (snapshot.asset, snapshot.venue);
            let state = books.entry(key).or_insert_with(BookRuntimeState::new);
            state.book.apply_snapshot(snapshot);
            state.complete = true;
            state.is_seeded = true;
            state.is_sequence_valid = true;
            state.is_gap_detected = false;
            state.last_event_ts = Some(snapshot.timestamp);
            state.maybe_push_bbo(snapshot.asset, snapshot.venue, snapshot.timestamp);
            cross_venue.update_book(snapshot.asset, snapshot.venue, &state.book);
            dirty_books.insert(key);
            health.last_market_event_ts = Some(snapshot.timestamp);
        }
        MarketEvent::L2Update(update) => {
            let key = (update.asset, update.venue);
            let state = books.entry(key).or_insert_with(BookRuntimeState::new);
            state.book.apply_update(update);
            state.is_gap_detected = false;
            state.last_event_ts = Some(update.timestamp);
            state.maybe_push_bbo(update.asset, update.venue, update.timestamp);
            cross_venue.update_book(update.asset, update.venue, &state.book);
            dirty_books.insert(key);
            health.last_market_event_ts = Some(update.timestamp);
        }
        MarketEvent::Bbo(quote) => {
            let key = (quote.asset, quote.venue);
            let state = books.entry(key).or_insert_with(BookRuntimeState::new);
            state.tracker.push(quote.clone());
            state.last_event_ts = Some(quote.timestamp);
            dirty_books.insert(key);
            health.last_market_event_ts = Some(quote.timestamp);
        }
        MarketEvent::MarkIndex(mark) => {
            let key = (mark.asset, mark.venue);
            let state = books.entry(key).or_insert_with(BookRuntimeState::new);
            state.last_mark_price = Some(mark.mark_price);
            state.last_index_price = Some(mark.index_price);
            state.last_event_ts = Some(mark.timestamp);
            dirty_books.insert(key);
            health.last_market_event_ts = Some(mark.timestamp);
        }
        MarketEvent::BookStatus(status) => {
            let key = (status.asset, status.venue);
            let state = books.entry(key).or_insert_with(BookRuntimeState::new);
            state.is_seeded = status.is_seeded;
            state.is_sequence_valid = status.is_sequence_valid;
            state.is_gap_detected = status.is_gap_detected;
            state.reseed_count = status.reseed_count;
            state.complete = status.is_seeded && status.is_sequence_valid;
            state.last_event_ts = Some(status.timestamp);
            dirty_books.insert(key);
            health.last_market_event_ts = Some(status.timestamp);
        }
    }

    sink.push_market_event(&event, received_at).await?;
    Ok(())
}

async fn handle_funding_event(
    sink: &mut PostgresSink,
    health: &mut EngineHealth,
    event: FundingEvent,
) -> Result<()> {
    *health
        .event_counts
        .entry("funding".to_string())
        .or_default() += 1;
    match &event {
        FundingEvent::Funding(funding) => {
            health.last_funding_event_ts = Some(funding.timestamp);
        }
        FundingEvent::Oi(oi) => {
            health.last_funding_event_ts = Some(oi.timestamp);
        }
        FundingEvent::Liquidation(liq) => {
            health.last_funding_event_ts = Some(liq.timestamp);
        }
    }
    sink.push_funding_event(&event).await?;
    Ok(())
}

async fn handle_iv_event(
    sink: &mut PostgresSink,
    health: &mut EngineHealth,
    event: IvEvent,
) -> Result<()> {
    *health.event_counts.entry("iv".to_string()).or_default() += 1;
    match &event {
        IvEvent::IvUpdate(iv) => {
            health.last_iv_event_ts = Some(iv.timestamp);
        }
    }
    sink.push_iv_event(&event).await?;
    Ok(())
}

async fn flush_dirty_books(
    sink: &mut PostgresSink,
    books: &HashMap<(Asset, Venue), BookRuntimeState>,
    dirty_books: &mut HashSet<(Asset, Venue)>,
) -> Result<()> {
    if dirty_books.is_empty() {
        return Ok(());
    }

    let now = Utc::now();
    let keys: Vec<(Asset, Venue)> = dirty_books.drain().collect();
    for (asset, venue) in keys {
        if let Some(state) = books.get(&(asset, venue)) {
            sink.push_book_state(now, asset, venue, state).await?;
        }
    }
    Ok(())
}

fn refresh_health(
    health: &mut EngineHealth,
    books: &HashMap<(Asset, Venue), BookRuntimeState>,
    cross_venue: &CrossVenueBook,
) {
    health.updated_at = Utc::now();
    health.books_tracked = books.len();
    health.complete_books = books.values().filter(|state| state.complete).count();
    health.partial_books = books.len().saturating_sub(health.complete_books);
    health.asset_summaries.clear();

    for asset in [Asset::Btc, Asset::Eth] {
        let mut summary = HealthAssetSummary::default();
        for ((book_asset, _), state) in books {
            if *book_asset != asset {
                continue;
            }
            summary.books_tracked += 1;
            if state.complete {
                summary.complete_books += 1;
            } else {
                summary.partial_books += 1;
            }
        }
        summary.cross_venue_spread = cross_venue.cross_venue_spread(asset);
        summary.cross_venue_depth_imbalance = cross_venue.cross_venue_depth_imbalance(asset);
        health.asset_summaries.insert(asset.to_string(), summary);
    }
}
