// file: kb_lib/src/dex_pool_materialization.rs

//! Shared DEX pool materialization helpers.
//!
//! This module persists normalized pool, pair, pool token and pool listing
//! records from decoded DEX events.
//!
//! It intentionally does not decode instructions and does not record analysis
//! signals. Detection services remain responsible for deciding which signals
//! to emit.

/// Token ordering strategy used when materializing a decoded DEX pool.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub(crate) enum DexPoolTokenOrder {
    /// `token_a_mint` is already the base token and `token_b_mint` is already the quote token.
    AlreadyBaseQuote,
    /// Base/quote order should be chosen from token A/B using known quote-token hints.
    ChooseBaseQuoteFromTokenAB,
}

/// Input required to materialize a normalized DEX pool.
#[derive(Debug, Clone)]
pub(crate) struct DexPoolMaterializationInput {
    /// Parent decoded event id.
    pub(crate) decoded_event_id: i64,
    /// Internal DEX id.
    pub(crate) dex_id: i64,
    /// Pool account address.
    pub(crate) pool_address: std::string::String,
    /// Token A mint, or already-base mint depending on `token_order`.
    pub(crate) token_a_mint: std::string::String,
    /// Token B mint, or already-quote mint depending on `token_order`.
    pub(crate) token_b_mint: std::string::String,
    /// Optional LP mint.
    pub(crate) lp_mint: std::option::Option<std::string::String>,
    /// Optional token A vault address, or base vault when `token_order` is `AlreadyBaseQuote`.
    pub(crate) token_a_vault_address: std::option::Option<std::string::String>,
    /// Optional token B vault address, or quote vault when `token_order` is `AlreadyBaseQuote`.
    pub(crate) token_b_vault_address: std::option::Option<std::string::String>,
    /// Pool kind to persist.
    pub(crate) pool_kind: crate::PoolKind,
    /// Pool status to persist.
    pub(crate) pool_status: crate::PoolStatus,
    /// Token ordering strategy.
    pub(crate) token_order: crate::dex_pool_materialization::DexPoolTokenOrder,
    /// Listing source kind.
    pub(crate) listing_source_kind: crate::ObservationSourceKind,
    /// Optional source endpoint logical name.
    pub(crate) source_endpoint_name: std::option::Option<std::string::String>,
}

impl DexPoolMaterializationInput {
    /// Creates a materialization input from a decoded event requiring pool and two token mints.
    pub(crate) fn from_decoded_event(
        decoded_event: &crate::DexDecodedEventDto,
        dex_id: i64,
        pool_kind: crate::PoolKind,
        pool_status: crate::PoolStatus,
        token_order: crate::dex_pool_materialization::DexPoolTokenOrder,
        token_a_vault_address: std::option::Option<std::string::String>,
        token_b_vault_address: std::option::Option<std::string::String>,
        source_endpoint_name: std::option::Option<std::string::String>,
    ) -> Result<Self, crate::Error> {
        let decoded_event_id_result =
            crate::dex_pool_materialization::required_decoded_event_id(decoded_event);
        let decoded_event_id = match decoded_event_id_result {
            Ok(decoded_event_id) => decoded_event_id,
            Err(error) => return Err(error),
        };
        let pool_address_result =
            crate::dex_pool_materialization::required_pool_account(decoded_event);
        let pool_address = match pool_address_result {
            Ok(pool_address) => pool_address,
            Err(error) => return Err(error),
        };
        let token_a_mint_result =
            crate::dex_pool_materialization::required_token_a_mint(decoded_event);
        let token_a_mint = match token_a_mint_result {
            Ok(token_a_mint) => token_a_mint,
            Err(error) => return Err(error),
        };
        let token_b_mint_result =
            crate::dex_pool_materialization::required_token_b_mint(decoded_event);
        let token_b_mint = match token_b_mint_result {
            Ok(token_b_mint) => token_b_mint,
            Err(error) => return Err(error),
        };
        return Ok(Self {
            decoded_event_id,
            dex_id,
            pool_address,
            token_a_mint,
            token_b_mint,
            lp_mint: decoded_event.lp_mint.clone(),
            token_a_vault_address,
            token_b_vault_address,
            pool_kind,
            pool_status,
            token_order,
            listing_source_kind: crate::ObservationSourceKind::Dex,
            source_endpoint_name,
        });
    }

    /// Creates a materialization input from a decoded event and explicit token mints.
    ///
    /// This is used by launch or bonding-curve events where the decoded event
    /// may expose only the launched token mint and the quote mint is inferred
    /// by the detector.
    pub(crate) fn from_decoded_event_with_mints(
        decoded_event: &crate::DexDecodedEventDto,
        dex_id: i64,
        token_a_mint: std::string::String,
        token_b_mint: std::string::String,
        lp_mint: std::option::Option<std::string::String>,
        pool_kind: crate::PoolKind,
        pool_status: crate::PoolStatus,
        token_order: crate::dex_pool_materialization::DexPoolTokenOrder,
        token_a_vault_address: std::option::Option<std::string::String>,
        token_b_vault_address: std::option::Option<std::string::String>,
        source_endpoint_name: std::option::Option<std::string::String>,
    ) -> Result<Self, crate::Error> {
        let decoded_event_id_result =
            crate::dex_pool_materialization::required_decoded_event_id(decoded_event);
        let decoded_event_id = match decoded_event_id_result {
            Ok(decoded_event_id) => decoded_event_id,
            Err(error) => return Err(error),
        };
        let pool_address_result =
            crate::dex_pool_materialization::required_pool_account(decoded_event);
        let pool_address = match pool_address_result {
            Ok(pool_address) => pool_address,
            Err(error) => return Err(error),
        };
        return Ok(Self {
            decoded_event_id,
            dex_id,
            pool_address,
            token_a_mint,
            token_b_mint,
            lp_mint,
            token_a_vault_address,
            token_b_vault_address,
            pool_kind,
            pool_status,
            token_order,
            listing_source_kind: crate::ObservationSourceKind::Dex,
            source_endpoint_name,
        });
    }
}

/// Returns the decoded event id or fails with a stable diagnostic.
pub(crate) fn required_decoded_event_id(
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<i64, crate::Error> {
    match decoded_event.id {
        Some(decoded_event_id) => return Ok(decoded_event_id),
        None => {
            return Err(crate::Error::InvalidState(
                "decoded dex event has no internal id".to_string(),
            ));
        },
    }
}

/// Returns the pool account or fails with a stable diagnostic.
pub(crate) fn required_pool_account(
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<std::string::String, crate::Error> {
    let decoded_event_id_result =
        crate::dex_pool_materialization::required_decoded_event_id(decoded_event);
    let decoded_event_id = match decoded_event_id_result {
        Ok(decoded_event_id) => decoded_event_id,
        Err(error) => return Err(error),
    };
    match decoded_event.pool_account.clone() {
        Some(pool_account) => return Ok(pool_account),
        None => {
            return Err(crate::Error::InvalidState(format!(
                "decoded event '{}' has no pool_account",
                decoded_event_id
            )));
        },
    }
}

/// Returns token A mint or fails with a stable diagnostic.
pub(crate) fn required_token_a_mint(
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<std::string::String, crate::Error> {
    let decoded_event_id_result =
        crate::dex_pool_materialization::required_decoded_event_id(decoded_event);
    let decoded_event_id = match decoded_event_id_result {
        Ok(decoded_event_id) => decoded_event_id,
        Err(error) => return Err(error),
    };
    match decoded_event.token_a_mint.clone() {
        Some(token_a_mint) => return Ok(token_a_mint),
        None => {
            return Err(crate::Error::InvalidState(format!(
                "decoded event '{}' has no token_a_mint",
                decoded_event_id
            )));
        },
    }
}

/// Returns token B mint or fails with a stable diagnostic.
pub(crate) fn required_token_b_mint(
    decoded_event: &crate::DexDecodedEventDto,
) -> Result<std::string::String, crate::Error> {
    let decoded_event_id_result =
        crate::dex_pool_materialization::required_decoded_event_id(decoded_event);
    let decoded_event_id = match decoded_event_id_result {
        Ok(decoded_event_id) => decoded_event_id,
        Err(error) => return Err(error),
    };
    match decoded_event.token_b_mint.clone() {
        Some(token_b_mint) => return Ok(token_b_mint),
        None => {
            return Err(crate::Error::InvalidState(format!(
                "decoded event '{}' has no token_b_mint",
                decoded_event_id
            )));
        },
    }
}

/// Persists pool, pair, pool tokens and listing, returning the materialized ids.
pub(crate) async fn materialize_dex_pool(
    database: &crate::Database,
    input: &crate::dex_pool_materialization::DexPoolMaterializationInput,
) -> Result<crate::DexPoolDetectionResult, crate::Error> {
    let ordered_tokens = crate::dex_pool_materialization::ordered_pool_tokens_from_input(input);
    let base_token_id_result =
        crate::dex_pool_materialization::ensure_token(database, ordered_tokens.base_mint.as_str())
            .await;
    let base_token_id = match base_token_id_result {
        Ok(base_token_id) => base_token_id,
        Err(error) => return Err(error),
    };
    let quote_token_id_result =
        crate::dex_pool_materialization::ensure_token(database, ordered_tokens.quote_mint.as_str())
            .await;
    let quote_token_id = match quote_token_id_result {
        Ok(quote_token_id) => quote_token_id,
        Err(error) => return Err(error),
    };
    let lp_token_id = match input.lp_mint.clone() {
        Some(lp_mint) => {
            let lp_token_id_result =
                crate::dex_pool_materialization::ensure_token(database, lp_mint.as_str()).await;
            match lp_token_id_result {
                Ok(lp_token_id) => Some(lp_token_id),
                Err(error) => return Err(error),
            }
        },
        None => None,
    };
    let pool_result = crate::dex_pool_materialization::ensure_pool(database, input).await;
    let pool_materialization = match pool_result {
        Ok(pool_materialization) => pool_materialization,
        Err(error) => return Err(error),
    };
    let pair_result = crate::dex_pool_materialization::ensure_pair(
        database,
        input.dex_id,
        pool_materialization.pool_id,
        base_token_id,
        quote_token_id,
        ordered_tokens.base_mint.as_str(),
        ordered_tokens.quote_mint.as_str(),
    )
    .await;
    let pair_materialization = match pair_result {
        Ok(pair_materialization) => pair_materialization,
        Err(error) => return Err(error),
    };
    let upsert_base_pool_token_result = crate::query_pool_tokens_upsert(
        database,
        &crate::PoolTokenDto::new(
            pool_materialization.pool_id,
            base_token_id,
            crate::PoolTokenRole::Base,
            ordered_tokens.base_vault_address,
            Some(0),
        ),
    )
    .await;
    if let Err(error) = upsert_base_pool_token_result {
        return Err(error);
    }
    let upsert_quote_pool_token_result = crate::query_pool_tokens_upsert(
        database,
        &crate::PoolTokenDto::new(
            pool_materialization.pool_id,
            quote_token_id,
            crate::PoolTokenRole::Quote,
            ordered_tokens.quote_vault_address,
            Some(1),
        ),
    )
    .await;
    if let Err(error) = upsert_quote_pool_token_result {
        return Err(error);
    }
    if let Some(lp_token_id) = lp_token_id {
        let upsert_lp_pool_token_result = crate::query_pool_tokens_upsert(
            database,
            &crate::PoolTokenDto::new(
                pool_materialization.pool_id,
                lp_token_id,
                crate::PoolTokenRole::LpMint,
                None,
                None,
            ),
        )
        .await;
        if let Err(error) = upsert_lp_pool_token_result {
            return Err(error);
        }
    }
    let listing_result = crate::dex_pool_materialization::ensure_pool_listing(
        database,
        input,
        pool_materialization.pool_id,
        pair_materialization.pair_id,
    )
    .await;
    let listing_materialization = match listing_result {
        Ok(listing_materialization) => listing_materialization,
        Err(error) => return Err(error),
    };
    return Ok(crate::DexPoolDetectionResult {
        decoded_event_id: input.decoded_event_id,
        dex_id: input.dex_id,
        pool_id: pool_materialization.pool_id,
        pair_id: pair_materialization.pair_id,
        pool_listing_id: listing_materialization.pool_listing_id,
        created_pool: pool_materialization.created_pool,
        created_pair: pair_materialization.created_pair,
        created_listing: listing_materialization.created_listing,
    });
}

#[derive(Debug, Clone)]
struct OrderedPoolTokens {
    base_mint: std::string::String,
    quote_mint: std::string::String,
    base_vault_address: std::option::Option<std::string::String>,
    quote_vault_address: std::option::Option<std::string::String>,
}

fn ordered_pool_tokens_from_input(
    input: &crate::dex_pool_materialization::DexPoolMaterializationInput,
) -> OrderedPoolTokens {
    match input.token_order {
        crate::dex_pool_materialization::DexPoolTokenOrder::AlreadyBaseQuote => {
            return OrderedPoolTokens {
                base_mint: input.token_a_mint.clone(),
                quote_mint: input.token_b_mint.clone(),
                base_vault_address: input.token_a_vault_address.clone(),
                quote_vault_address: input.token_b_vault_address.clone(),
            };
        },
        crate::dex_pool_materialization::DexPoolTokenOrder::ChooseBaseQuoteFromTokenAB => {
            let base_is_token_a = crate::dex_pool_materialization::choose_base_quote_order(
                input.token_a_mint.as_str(),
                input.token_b_mint.as_str(),
            );
            let base_mint = if base_is_token_a {
                input.token_a_mint.clone()
            } else {
                input.token_b_mint.clone()
            };
            let quote_mint = if base_is_token_a {
                input.token_b_mint.clone()
            } else {
                input.token_a_mint.clone()
            };
            let base_vault_address = if base_is_token_a {
                input.token_a_vault_address.clone()
            } else {
                input.token_b_vault_address.clone()
            };
            let quote_vault_address = if base_is_token_a {
                input.token_b_vault_address.clone()
            } else {
                input.token_a_vault_address.clone()
            };
            return OrderedPoolTokens {
                base_mint,
                quote_mint,
                base_vault_address,
                quote_vault_address,
            };
        },
    }
}

async fn ensure_pool(
    database: &crate::Database,
    input: &crate::dex_pool_materialization::DexPoolMaterializationInput,
) -> Result<PoolMaterialization, crate::Error> {
    let existing_pool_result =
        crate::query_pools_get_by_address(database, input.pool_address.as_str()).await;
    let existing_pool_option = match existing_pool_result {
        Ok(existing_pool_option) => existing_pool_option,
        Err(error) => return Err(error),
    };
    let created_pool = existing_pool_option.is_none();
    let pool_id = match existing_pool_option {
        Some(pool) => match pool.id {
            Some(pool_id) => pool_id,
            None => {
                return Err(crate::Error::InvalidState(format!(
                    "pool '{}' has no internal id",
                    pool.address
                )));
            },
        },
        None => {
            let pool_dto = crate::PoolDto::new(
                input.dex_id,
                input.pool_address.clone(),
                input.pool_kind,
                input.pool_status,
            );
            let upsert_result = crate::query_pools_upsert(database, &pool_dto).await;
            match upsert_result {
                Ok(pool_id) => pool_id,
                Err(error) => return Err(error),
            }
        },
    };
    return Ok(PoolMaterialization { pool_id, created_pool });
}

async fn ensure_pair(
    database: &crate::Database,
    dex_id: i64,
    pool_id: i64,
    base_token_id: i64,
    quote_token_id: i64,
    base_mint: &str,
    quote_mint: &str,
) -> Result<PairMaterialization, crate::Error> {
    let existing_pair_result = crate::query_pairs_get_by_pool_id(database, pool_id).await;
    let existing_pair_option = match existing_pair_result {
        Ok(existing_pair_option) => existing_pair_option,
        Err(error) => return Err(error),
    };
    let created_pair = existing_pair_option.is_none();
    let pair_symbol = crate::dex_pool_materialization::build_pair_symbol(base_mint, quote_mint);
    let pair_id = match existing_pair_option {
        Some(pair) => match pair.id {
            Some(pair_id) => pair_id,
            None => {
                return Err(crate::Error::InvalidState(format!(
                    "pair for pool '{}' has no internal id",
                    pool_id
                )));
            },
        },
        None => {
            let pair_dto =
                crate::PairDto::new(dex_id, pool_id, base_token_id, quote_token_id, pair_symbol);
            let upsert_result = crate::query_pairs_upsert(database, &pair_dto).await;
            match upsert_result {
                Ok(pair_id) => pair_id,
                Err(error) => return Err(error),
            }
        },
    };
    return Ok(PairMaterialization { pair_id, created_pair });
}

async fn ensure_pool_listing(
    database: &crate::Database,
    input: &crate::dex_pool_materialization::DexPoolMaterializationInput,
    pool_id: i64,
    pair_id: i64,
) -> Result<ListingMaterialization, crate::Error> {
    let existing_listing_result =
        crate::query_pool_listings_get_by_pool_id(database, pool_id).await;
    let existing_listing_option = match existing_listing_result {
        Ok(existing_listing_option) => existing_listing_option,
        Err(error) => return Err(error),
    };
    let created_listing = existing_listing_option.is_none();
    let pool_listing_id = match existing_listing_option {
        Some(pool_listing) => pool_listing.id,
        None => {
            let listing_dto = crate::PoolListingDto::new(
                input.dex_id,
                pool_id,
                Some(pair_id),
                input.listing_source_kind,
                input.source_endpoint_name.clone(),
                None,
                None,
                None,
            );
            let upsert_result = crate::query_pool_listings_upsert(database, &listing_dto).await;
            match upsert_result {
                Ok(listing_id) => Some(listing_id),
                Err(error) => return Err(error),
            }
        },
    };
    return Ok(ListingMaterialization { pool_listing_id, created_listing });
}

async fn ensure_token(database: &crate::Database, mint: &str) -> Result<i64, crate::Error> {
    let token_result = crate::query_tokens_get_by_mint(database, mint).await;
    let token_option = match token_result {
        Ok(token_option) => token_option,
        Err(error) => return Err(error),
    };
    match token_option {
        Some(token) => match token.id {
            Some(token_id) => return Ok(token_id),
            None => {
                return Err(crate::Error::InvalidState(format!(
                    "token '{}' has no internal id",
                    mint
                )));
            },
        },
        None => {
            let token_dto = crate::TokenDto::new(
                mint.to_string(),
                None,
                None,
                None,
                crate::SPL_TOKEN_PROGRAM_ID.to_string(),
                crate::dex_pool_materialization::is_quote_mint(mint),
            );
            return crate::query_tokens_upsert(database, &token_dto).await;
        },
    }
}

#[derive(Debug, Clone)]
struct PoolMaterialization {
    pool_id: i64,
    created_pool: bool,
}

#[derive(Debug, Clone)]
struct PairMaterialization {
    pair_id: i64,
    created_pair: bool,
}

#[derive(Debug, Clone)]
struct ListingMaterialization {
    pool_listing_id: std::option::Option<i64>,
    created_listing: bool,
}

fn is_quote_mint(mint: &str) -> bool {
    return mint == crate::WSOL_MINT_ID;
}

fn choose_base_quote_order(token_a_mint: &str, token_b_mint: &str) -> bool {
    let token_a_is_quote = crate::dex_pool_materialization::is_quote_mint(token_a_mint);
    let token_b_is_quote = crate::dex_pool_materialization::is_quote_mint(token_b_mint);
    if token_a_is_quote && !token_b_is_quote {
        return false;
    }
    if token_b_is_quote && !token_a_is_quote {
        return true;
    }
    return true;
}

fn build_pair_symbol(
    base_mint: &str,
    quote_mint: &str,
) -> std::option::Option<std::string::String> {
    let base_symbol = crate::dex_pool_materialization::symbol_hint_from_mint(base_mint);
    let quote_symbol = crate::dex_pool_materialization::symbol_hint_from_mint(quote_mint);
    match (base_symbol, quote_symbol) {
        (Some(base_symbol), Some(quote_symbol)) => {
            return Some(format!("{base_symbol}/{quote_symbol}"));
        },
        _ => return None,
    }
}

fn symbol_hint_from_mint(mint: &str) -> std::option::Option<std::string::String> {
    if mint == crate::WSOL_MINT_ID {
        return Some("WSOL".to_string());
    }
    return None;
}

#[cfg(test)]
mod tests {
    #[test]
    fn quote_token_is_moved_to_quote_side_when_order_is_chosen() {
        let input = crate::dex_pool_materialization::DexPoolMaterializationInput {
            decoded_event_id: 1,
            dex_id: 2,
            pool_address: "Pool111".to_string(),
            token_a_mint: crate::WSOL_MINT_ID.to_string(),
            token_b_mint: "TokenB111".to_string(),
            lp_mint: None,
            token_a_vault_address: Some("VaultA111".to_string()),
            token_b_vault_address: Some("VaultB111".to_string()),
            pool_kind: crate::PoolKind::Amm,
            pool_status: crate::PoolStatus::Active,
            token_order:
                crate::dex_pool_materialization::DexPoolTokenOrder::ChooseBaseQuoteFromTokenAB,
            listing_source_kind: crate::ObservationSourceKind::Dex,
            source_endpoint_name: Some("test".to_string()),
        };
        let ordered = super::ordered_pool_tokens_from_input(&input);
        assert_eq!(ordered.base_mint, "TokenB111");
        assert_eq!(ordered.quote_mint, crate::WSOL_MINT_ID);
        assert_eq!(ordered.base_vault_address, Some("VaultB111".to_string()));
        assert_eq!(ordered.quote_vault_address, Some("VaultA111".to_string()));
    }

    #[test]
    fn already_base_quote_order_is_preserved() {
        let input = crate::dex_pool_materialization::DexPoolMaterializationInput {
            decoded_event_id: 1,
            dex_id: 2,
            pool_address: "Pool111".to_string(),
            token_a_mint: crate::WSOL_MINT_ID.to_string(),
            token_b_mint: "TokenB111".to_string(),
            lp_mint: None,
            token_a_vault_address: Some("BaseVault111".to_string()),
            token_b_vault_address: Some("QuoteVault111".to_string()),
            pool_kind: crate::PoolKind::Amm,
            pool_status: crate::PoolStatus::Active,
            token_order: crate::dex_pool_materialization::DexPoolTokenOrder::AlreadyBaseQuote,
            listing_source_kind: crate::ObservationSourceKind::Dex,
            source_endpoint_name: Some("test".to_string()),
        };
        let ordered = super::ordered_pool_tokens_from_input(&input);
        assert_eq!(ordered.base_mint, crate::WSOL_MINT_ID);
        assert_eq!(ordered.quote_mint, "TokenB111");
        assert_eq!(ordered.base_vault_address, Some("BaseVault111".to_string()));
        assert_eq!(ordered.quote_vault_address, Some("QuoteVault111".to_string()));
    }

    #[test]
    fn pair_symbol_is_none_when_only_one_symbol_is_known() {
        let pair_symbol = super::build_pair_symbol(crate::WSOL_MINT_ID, "TokenB111");
        assert_eq!(pair_symbol, None);
    }

    #[test]
    fn wsol_symbol_hint_is_known() {
        let symbol = super::symbol_hint_from_mint(crate::WSOL_MINT_ID);
        assert_eq!(symbol, Some("WSOL".to_string()));
    }
}