// file: kb_lib/src/trade_aggregation_context.rs

//! Database context loading for trade aggregation.
//!
//! This module resolves the normalized database context required by
//! `TradeAggregationService`: transaction, decoded events, pool, pair,
//! base/quote token metadata and pool token vault addresses.

/// Transaction-level context used by trade aggregation.
pub(crate) struct TradeAggregationTransactionContext {
    /// Persisted transaction row.
    pub(crate) transaction: crate::ChainTransactionDto,
    /// Internal transaction id.
    pub(crate) transaction_id: i64,
    /// Decoded DEX events attached to the transaction.
    pub(crate) decoded_events: std::vec::Vec<crate::DexDecodedEventDto>,
}

/// Decoded-event-level context used by trade aggregation.
pub(crate) struct TradeAggregationDecodedEventContext {
    /// Internal decoded event id.
    pub(crate) decoded_event_id: i64,
    /// Existing trade event, when this decoded event was already materialized.
    pub(crate) existing_trade_event: std::option::Option<crate::TradeEventDto>,
    /// Pool account address from the decoded event.
    pub(crate) pool_address: std::string::String,
    /// Persisted pool row.
    pub(crate) pool: crate::PoolDto,
    /// Internal pool id.
    pub(crate) pool_id: i64,
    /// Persisted pair row.
    pub(crate) pair: crate::PairDto,
    /// Internal pair id.
    pub(crate) pair_id: i64,
    /// Base token mint, when the token row exists.
    pub(crate) base_token_mint: std::option::Option<std::string::String>,
    /// Base token decimals, when known.
    pub(crate) base_token_decimals: std::option::Option<u8>,
    /// Quote token mint, when the token row exists.
    pub(crate) quote_token_mint: std::option::Option<std::string::String>,
    /// Quote token decimals, when known.
    pub(crate) quote_token_decimals: std::option::Option<u8>,
    /// Base token vault address, when known.
    pub(crate) base_vault_address: std::option::Option<std::string::String>,
    /// Quote token vault address, when known.
    pub(crate) quote_vault_address: std::option::Option<std::string::String>,
}

/// Loads a transaction and its decoded DEX events from one signature.
pub(crate) async fn load_trade_aggregation_transaction_context(
    database: &crate::Database,
    signature: &str,
) -> Result<crate::trade_aggregation_context::TradeAggregationTransactionContext, crate::Error> {
    let transaction_result =
        crate::query_chain_transactions_get_by_signature(database, signature).await;
    let transaction_option = match transaction_result {
        Ok(transaction_option) => transaction_option,
        Err(error) => return Err(error),
    };
    let transaction = match transaction_option {
        Some(transaction) => transaction,
        None => {
            return Err(crate::Error::InvalidState(format!(
                "cannot aggregate trades for unknown transaction '{}'",
                signature
            )));
        },
    };
    let transaction_id = match transaction.id {
        Some(transaction_id) => transaction_id,
        None => {
            return Err(crate::Error::InvalidState(format!(
                "transaction '{}' has no internal id",
                signature
            )));
        },
    };
    let decoded_events_result =
        crate::query_dex_decoded_events_list_by_transaction_id(database, transaction_id).await;
    let decoded_events = match decoded_events_result {
        Ok(decoded_events) => decoded_events,
        Err(error) => return Err(error),
    };
    return Ok(crate::trade_aggregation_context::TradeAggregationTransactionContext {
        transaction,
        transaction_id,
        decoded_events,
    });
}

/// Loads the normalized DB context for one decoded event.
///
/// Returns `Ok(None)` when the decoded event is not materializable yet:
/// missing pool account, pool row or pair row.
pub(crate) async fn load_trade_aggregation_decoded_event_context(
    database: &crate::Database,
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<
    std::option::Option<crate::trade_aggregation_context::TradeAggregationDecodedEventContext>,
    crate::Error,
> {
    let decoded_event_id = match decoded_event.id {
        Some(decoded_event_id) => decoded_event_id,
        None => {
            return Err(crate::Error::InvalidState("decoded event has no internal id".to_string()));
        },
    };
    let existing_trade_result =
        crate::query_trade_events_get_by_decoded_event_id(database, decoded_event_id).await;
    let existing_trade_event = match existing_trade_result {
        Ok(existing_trade_event) => existing_trade_event,
        Err(error) => return Err(error),
    };
    let pool_address = match decoded_event.pool_account.clone() {
        Some(pool_address) => pool_address,
        None => return Ok(None),
    };
    let pool_result = crate::query_pools_get_by_address(database, pool_address.as_str()).await;
    let pool_option = match pool_result {
        Ok(pool_option) => pool_option,
        Err(error) => return Err(error),
    };
    let pool = match pool_option {
        Some(pool) => pool,
        None => {
            let materialized =
                crate::trade_aggregation_context::try_materialize_pump_swap_trade_pool(
                    database,
                    decoded_event,
                )
                .await;
            match materialized {
                Ok(true) => {
                    let refreshed_result =
                        crate::query_pools_get_by_address(database, pool_address.as_str()).await;
                    match refreshed_result {
                        Ok(Some(pool)) => pool,
                        Ok(None) => return Ok(None),
                        Err(error) => return Err(error),
                    }
                },
                Ok(false) => return Ok(None),
                Err(error) => return Err(error),
            }
        },
    };
    let pool_id = match pool.id {
        Some(pool_id) => pool_id,
        None => {
            return Err(crate::Error::InvalidState(format!(
                "pool '{}' has no internal id",
                pool.address
            )));
        },
    };
    let pair_result = crate::query_pairs_get_by_pool_id(database, pool_id).await;
    let pair_option = match pair_result {
        Ok(pair_option) => pair_option,
        Err(error) => return Err(error),
    };
    let pair = match pair_option {
        Some(pair) => pair,
        None => {
            let materialized =
                crate::trade_aggregation_context::try_materialize_pump_swap_trade_pool(
                    database,
                    decoded_event,
                )
                .await;
            match materialized {
                Ok(true) => {
                    let refreshed_pair_result = crate::query_pairs_get_by_pool_id(database, pool_id).await;
                    match refreshed_pair_result {
                        Ok(Some(pair)) => pair,
                        Ok(None) => return Ok(None),
                        Err(error) => return Err(error),
                    }
                },
                Ok(false) => return Ok(None),
                Err(error) => return Err(error),
            }
        },
    };
    let pair_id = match pair.id {
        Some(pair_id) => pair_id,
        None => {
            return Err(crate::Error::InvalidState(format!(
                "pair for pool '{}' has no internal id",
                pool_id
            )));
        },
    };
    let base_token_result = crate::query_tokens_get_by_id(database, pair.base_token_id).await;
    let (base_token_mint, base_token_decimals) = match base_token_result {
        Ok(Some(token)) => (Some(token.mint), token.decimals),
        Ok(None) => (None, None),
        Err(error) => return Err(error),
    };
    let quote_token_result = crate::query_tokens_get_by_id(database, pair.quote_token_id).await;
    let (quote_token_mint, quote_token_decimals) = match quote_token_result {
        Ok(Some(token)) => (Some(token.mint), token.decimals),
        Ok(None) => (None, None),
        Err(error) => return Err(error),
    };
    let pool_tokens_result = crate::query_pool_tokens_list_by_pool_id(database, pool_id).await;
    let pool_tokens = match pool_tokens_result {
        Ok(pool_tokens) => pool_tokens,
        Err(error) => return Err(error),
    };
    let base_vault_address =
        crate::trade_aggregation_context::find_pool_token_vault_address_by_token_id(
            &pool_tokens,
            pair.base_token_id,
        );
    let quote_vault_address =
        crate::trade_aggregation_context::find_pool_token_vault_address_by_token_id(
            &pool_tokens,
            pair.quote_token_id,
        );
    return Ok(Some(crate::trade_aggregation_context::TradeAggregationDecodedEventContext {
        decoded_event_id,
        existing_trade_event,
        pool_address,
        pool,
        pool_id,
        pair,
        pair_id,
        base_token_mint,
        base_token_decimals,
        quote_token_mint,
        quote_token_decimals,
        base_vault_address,
        quote_vault_address,
    }));
}

async fn try_materialize_pump_swap_trade_pool(
    database: &crate::Database,
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<bool, crate::Error> {
    if decoded_event.protocol_name != "pump_swap" {
        return Ok(false);
    }
    if !crate::is_dex_trade_event_kind(decoded_event.event_kind.as_str()) {
        return Ok(false);
    }
    if decoded_event.pool_account.is_none()
        || decoded_event.token_a_mint.is_none()
        || decoded_event.token_b_mint.is_none()
    {
        return Ok(false);
    }
    let dex_result = crate::query_dexs_get_by_code(database, decoded_event.protocol_name.as_str()).await;
    let dex = match dex_result {
        Ok(Some(dex)) => dex,
        Ok(None) => return Ok(false),
        Err(error) => return Err(error),
    };
    let dex_id = match dex.id {
        Some(dex_id) => dex_id,
        None => return Ok(false),
    };
    let payload_result = serde_json::from_str::<serde_json::Value>(decoded_event.payload_json.as_str());
    let payload = match payload_result {
        Ok(payload) => payload,
        Err(_) => return Ok(false),
    };
    let token_a_vault_address = crate::trade_aggregation_context::extract_payload_string(
        &payload,
        &["poolBaseTokenAccount", "pool_base_token_account", "baseVault", "base_vault"],
    );
    let token_b_vault_address = crate::trade_aggregation_context::extract_payload_string(
        &payload,
        &["poolQuoteTokenAccount", "pool_quote_token_account", "quoteVault", "quote_vault"],
    );
    let materialization_input_result =
        crate::dex_pool_materialization::DexPoolMaterializationInput::from_decoded_event(
            decoded_event,
            dex_id,
            crate::PoolKind::Amm,
            crate::PoolStatus::Active,
            crate::dex_pool_materialization::DexPoolTokenOrder::AlreadyBaseQuote,
            token_a_vault_address,
            token_b_vault_address,
            None,
        );
    let materialization_input = match materialization_input_result {
        Ok(materialization_input) => materialization_input,
        Err(_) => return Ok(false),
    };
    let materialization_result =
        crate::dex_pool_materialization::materialize_dex_pool(database, &materialization_input).await;
    match materialization_result {
        Ok(_) => return Ok(true),
        Err(error) => return Err(error),
    }
}

fn extract_payload_string(
    payload: &serde_json::Value,
    candidate_keys: &[&str],
) -> std::option::Option<std::string::String> {
    if let Some(object) = payload.as_object() {
        for candidate_key in candidate_keys {
            if let Some(value) = object.get(*candidate_key) {
                if let Some(text) = value.as_str() {
                    let trimmed = text.trim();
                    if !trimmed.is_empty() {
                        return Some(trimmed.to_string());
                    }
                }
            }
        }
        for nested_value in object.values() {
            let nested = crate::trade_aggregation_context::extract_payload_string(
                nested_value,
                candidate_keys,
            );
            if nested.is_some() {
                return nested;
            }
        }
    }
    if let Some(array) = payload.as_array() {
        for nested_value in array {
            let nested = crate::trade_aggregation_context::extract_payload_string(
                nested_value,
                candidate_keys,
            );
            if nested.is_some() {
                return nested;
            }
        }
    }
    return None;
}

fn find_pool_token_vault_address_by_token_id(
    pool_tokens: &[crate::PoolTokenDto],
    token_id: i64,
) -> std::option::Option<std::string::String> {
    for pool_token in pool_tokens {
        if pool_token.token_id != token_id {
            continue;
        }
        let vault_address = match pool_token.vault_address.clone() {
            Some(vault_address) => vault_address.trim().to_string(),
            None => continue,
        };
        if vault_address.is_empty() {
            continue;
        }
        return Some(vault_address);
    }
    return None;
}
