// file: kb_lib/src/wallet_holding_observation.rs

//! Wallet-holding observation service.

/// One wallet-holding observation result.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WalletHoldingObservationResult {
    /// Persisted wallet-holding id.
    pub wallet_holding_id: i64,
    /// Related wallet id.
    pub wallet_id: i64,
    /// Related token id.
    pub token_id: i64,
    /// Whether the holding row was newly created.
    pub created_holding: bool,
}

/// Wallet-holding observation service.
#[derive(Debug, Clone)]
pub struct WalletHoldingObservationService {
    database: std::sync::Arc<crate::Database>,
    persistence: crate::DetectionPersistenceService,
}

impl WalletHoldingObservationService {
    /// Creates a new wallet-holding observation service.
    pub fn new(database: std::sync::Arc<crate::Database>) -> Self {
        let persistence = crate::DetectionPersistenceService::new(database.clone());
        return Self { database, persistence };
    }

    /// Records observed wallet holdings for one resolved transaction signature.
    pub async fn record_transaction_by_signature(
        &self,
        signature: &str,
    ) -> Result<std::vec::Vec<crate::WalletHoldingObservationResult>, crate::Error> {
        let transaction_result =
            crate::query_chain_transactions_get_by_signature(self.database.as_ref(), 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 record wallet holdings 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(
            self.database.as_ref(),
            transaction_id,
        )
        .await;
        let decoded_events = match decoded_events_result {
            Ok(decoded_events) => decoded_events,
            Err(error) => return Err(error),
        };
        let mut results = std::vec::Vec::new();
        let mut seen_pairs =
            std::collections::HashSet::<(i64, i64, std::option::Option<i64>)>::new();
        for decoded_event in &decoded_events {
            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 pool_id = match decoded_event.pool_account.clone() {
                Some(pool_address) => {
                    let pool_result = crate::query_pools_get_by_address(
                        self.database.as_ref(),
                        pool_address.as_str(),
                    )
                    .await;
                    let pool_option = match pool_result {
                        Ok(pool_option) => pool_option,
                        Err(error) => return Err(error),
                    };
                    match pool_option {
                        Some(pool) => pool.id,
                        None => None,
                    }
                },
                None => None,
            };
            let pair_id = match pool_id {
                Some(pool_id) => {
                    let pair_result =
                        crate::query_pairs_get_by_pool_id(self.database.as_ref(), pool_id).await;
                    let pair_option = match pair_result {
                        Ok(pair_option) => pair_option,
                        Err(error) => return Err(error),
                    };
                    match pair_option {
                        Some(pair) => pair.id,
                        None => None,
                    }
                },
                None => None,
            };
            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(error) => {
                    return Err(crate::Error::Json(format!(
                        "cannot parse decoded_event payload_json '{}': {}",
                        decoded_event.payload_json, error
                    )));
                },
            };
            let wallet_roles = collect_wallet_roles(&payload);
            let token_mints = collect_token_mints(decoded_event);
            for (role, wallet_address) in wallet_roles {
                let wallet_result = crate::query_wallets_get_by_address(
                    self.database.as_ref(),
                    wallet_address.as_str(),
                )
                .await;
                let wallet_option = match wallet_result {
                    Ok(wallet_option) => wallet_option,
                    Err(error) => return Err(error),
                };
                let wallet = match wallet_option {
                    Some(wallet) => wallet,
                    None => continue,
                };
                let wallet_id = match wallet.id {
                    Some(wallet_id) => wallet_id,
                    None => {
                        return Err(crate::Error::InvalidState(
                            "wallet has no internal id".to_string(),
                        ));
                    },
                };
                for token_mint in &token_mints {
                    let token_result = crate::query_tokens_get_by_mint(
                        self.database.as_ref(),
                        token_mint.as_str(),
                    )
                    .await;
                    let token_option = match token_result {
                        Ok(token_option) => token_option,
                        Err(error) => return Err(error),
                    };
                    let token = match token_option {
                        Some(token) => token,
                        None => continue,
                    };
                    let token_id = match token.id {
                        Some(token_id) => token_id,
                        None => {
                            return Err(crate::Error::InvalidState(
                                "token has no internal id".to_string(),
                            ));
                        },
                    };
                    let dedupe_key = (wallet_id, token_id, Some(decoded_event_id));
                    if seen_pairs.contains(&dedupe_key) {
                        continue;
                    }
                    seen_pairs.insert(dedupe_key);
                    let existing_result = crate::query_wallet_holdings_get_by_wallet_and_token(
                        self.database.as_ref(),
                        wallet_id,
                        token_id,
                    )
                    .await;
                    let existing_option = match existing_result {
                        Ok(existing_option) => existing_option,
                        Err(error) => return Err(error),
                    };
                    let created_holding = existing_option.is_none();
                    let first_transaction_id = match existing_option.clone() {
                        Some(existing) => existing.first_transaction_id,
                        None => transaction_id,
                    };
                    let slot_i64 = convert_slot_to_i64(transaction.slot);
                    let dto = crate::WalletHoldingDto::new(
                        wallet_id,
                        token_id,
                        first_transaction_id,
                        transaction_id,
                        Some(decoded_event_id),
                        pool_id,
                        pair_id,
                        Some(role.clone()),
                        None,
                        slot_i64,
                        crate::ObservationSourceKind::Dex,
                        transaction.source_endpoint_name.clone(),
                    );
                    let upsert_result =
                        crate::query_wallet_holdings_upsert(self.database.as_ref(), &dto).await;
                    let wallet_holding_id = match upsert_result {
                        Ok(wallet_holding_id) => wallet_holding_id,
                        Err(error) => return Err(error),
                    };
                    if created_holding {
                        let payload = serde_json::json!({
                            "walletAddress": wallet_address,
                            "tokenMint": token_mint,
                            "role": role,
                            "transactionSignature": transaction.signature,
                            "poolId": pool_id,
                            "pairId": pair_id
                        });
                        let observation_result = self
                            .persistence
                            .record_observation(&crate::DetectionObservationInput::new(
                                "wallet.holding".to_string(),
                                crate::ObservationSourceKind::Dex,
                                transaction.source_endpoint_name.clone(),
                                transaction.signature.clone(),
                                transaction.slot,
                                payload.clone(),
                            ))
                            .await;
                        let observation_id = match observation_result {
                            Ok(observation_id) => observation_id,
                            Err(error) => return Err(error),
                        };
                        let signal_result = self
                            .persistence
                            .record_signal(&crate::DetectionSignalInput::new(
                                "signal.wallet.holding.observed".to_string(),
                                crate::AnalysisSignalSeverity::Low,
                                transaction.signature.clone(),
                                Some(observation_id),
                                None,
                                payload,
                            ))
                            .await;
                        if let Err(error) = signal_result {
                            return Err(error);
                        }
                    }
                    results.push(crate::WalletHoldingObservationResult {
                        wallet_holding_id,
                        wallet_id,
                        token_id,
                        created_holding,
                    });
                }
            }
        }
        return Ok(results);
    }
}

fn collect_wallet_roles(
    payload: &serde_json::Value,
) -> std::vec::Vec<(std::string::String, std::string::String)> {
    let role_mappings = [
        ("payer", vec!["payer", "funder"]),
        ("creator", vec!["creator", "poolCreator"]),
        ("owner", vec!["owner"]),
        ("user", vec!["user"]),
    ];
    let mut seen = std::collections::HashSet::<std::string::String>::new();
    let mut results = std::vec::Vec::new();
    for (role, keys) in role_mappings {
        let addresses = extract_strings_for_candidate_keys(payload, &keys);
        for address in addresses {
            let dedupe_key = format!("{}:{}", role, address);
            if seen.contains(&dedupe_key) {
                continue;
            }
            seen.insert(dedupe_key);
            results.push((role.to_string(), address));
        }
    }
    return results;
}

fn collect_token_mints(
    decoded_event: &crate::DexDecodedEventDto,
) -> std::vec::Vec<std::string::String> {
    let mut seen = std::collections::HashSet::<std::string::String>::new();
    let mut values = std::vec::Vec::new();
    if let Some(token_a_mint) = decoded_event.token_a_mint.clone() {
        if !seen.contains(token_a_mint.as_str()) {
            seen.insert(token_a_mint.clone());
            values.push(token_a_mint);
        }
    }
    if let Some(token_b_mint) = decoded_event.token_b_mint.clone() {
        if !seen.contains(token_b_mint.as_str()) {
            seen.insert(token_b_mint.clone());
            values.push(token_b_mint);
        }
    }
    return values;
}

fn extract_strings_for_candidate_keys(
    value: &serde_json::Value,
    candidate_keys: &[&str],
) -> std::vec::Vec<std::string::String> {
    let mut values = std::vec::Vec::new();
    extract_strings_for_candidate_keys_inner(value, candidate_keys, &mut values);
    return values;
}

fn extract_strings_for_candidate_keys_inner(
    value: &serde_json::Value,
    candidate_keys: &[&str],
    values: &mut std::vec::Vec<std::string::String>,
) {
    if let Some(object) = value.as_object() {
        for candidate_key in candidate_keys {
            let direct_option = object.get(*candidate_key);
            if let Some(direct) = direct_option {
                let text_option = direct.as_str();
                if let Some(text) = text_option {
                    values.push(text.to_string());
                }
            }
        }
        for nested_value in object.values() {
            extract_strings_for_candidate_keys_inner(nested_value, candidate_keys, values);
        }
        return;
    }
    if let Some(array) = value.as_array() {
        for nested_value in array {
            extract_strings_for_candidate_keys_inner(nested_value, candidate_keys, values);
        }
    }
}

fn convert_slot_to_i64(slot: std::option::Option<u64>) -> std::option::Option<i64> {
    match slot {
        Some(slot) => match i64::try_from(slot) {
            Ok(slot) => return Some(slot),
            Err(_) => return None,
        },
        None => return None,
    }
}

#[cfg(test)]
mod tests {
    async fn make_database() -> std::sync::Arc<crate::Database> {
        let tempdir_result = tempfile::tempdir();
        let tempdir = match tempdir_result {
            Ok(tempdir) => tempdir,
            Err(error) => panic!("tempdir must succeed: {}", error),
        };
        let database_path = tempdir.path().join("wallet_holding_observation.sqlite3");
        let config = crate::DatabaseConfig {
            enabled: true,
            backend: crate::DatabaseBackend::Sqlite,
            sqlite: crate::SqliteDatabaseConfig {
                path: database_path.to_string_lossy().to_string(),
                create_if_missing: true,
                busy_timeout_ms: 5000,
                max_connections: 1,
                auto_initialize_schema: true,
                use_wal: true,
            },
        };
        let database_result = crate::Database::connect_and_initialize(&config).await;
        let database = match database_result {
            Ok(database) => database,
            Err(error) => panic!("database init must succeed: {}", error),
        };
        return std::sync::Arc::new(database);
    }

    async fn seed_fluxbeam_create_transaction(
        database: std::sync::Arc<crate::Database>,
        signature: &str,
    ) {
        let transaction_model = crate::TransactionModelService::new(database.clone());
        let dex_decode = crate::DexDecodeService::new(database.clone());
        let dex_detect = crate::DexDetectService::new(database.clone());
        let wallet_observation = crate::WalletObservationService::new(database.clone());
        let resolved_transaction = serde_json::json!({
            "slot": 950001,
            "blockTime": 1779600001,
            "version": 0,
            "transaction": {
                "message": {
                    "instructions": [
                        {
                            "programId": crate::FLUXBEAM_PROGRAM_ID,
                            "program": "fluxbeam",
                            "stackHeight": 1,
                            "accounts": [
                                "HoldingPool111",
                                "HoldingLpMint111",
                                "HoldingTokenA111",
                                crate::WSOL_MINT_ID,
                                "HoldingWallet111"
                            ],
                            "parsed": {
                                "info": {
                                    "instruction": "create_pool",
                                    "pool": "HoldingPool111",
                                    "lpMint": "HoldingLpMint111",
                                    "tokenA": "HoldingTokenA111",
                                    "tokenB": crate::WSOL_MINT_ID,
                                    "payer": "HoldingWallet111"
                                }
                            },
                            "data": "opaque"
                        }
                    ]
                }
            },
            "meta": {
                "err": null,
                "logMessages": [
                    "Program log: Instruction: CreatePool"
                ]
            }
        });
        let project_result = transaction_model
            .persist_resolved_transaction(
                signature,
                Some("helius_primary_http".to_string()),
                &resolved_transaction,
            )
            .await;
        if let Err(error) = project_result {
            panic!("projection must succeed: {}", error);
        }
        let decode_result = dex_decode.decode_transaction_by_signature(signature).await;
        if let Err(error) = decode_result {
            panic!("dex decode must succeed: {}", error);
        }
        let detect_result = dex_detect.detect_transaction_by_signature(signature).await;
        if let Err(error) = detect_result {
            panic!("dex detect must succeed: {}", error);
        }
        let wallet_result = wallet_observation.record_transaction_by_signature(signature).await;
        if let Err(error) = wallet_result {
            panic!("wallet observation must succeed: {}", error);
        }
    }

    #[tokio::test]
    async fn record_transaction_by_signature_creates_wallet_holdings() {
        let database = make_database().await;
        seed_fluxbeam_create_transaction(database.clone(), "sig-wallet-holding-1").await;
        let service = crate::WalletHoldingObservationService::new(database.clone());
        let record_result = service.record_transaction_by_signature("sig-wallet-holding-1").await;
        let results = match record_result {
            Ok(results) => results,
            Err(error) => panic!("wallet holding observation must succeed: {}", error),
        };
        assert_eq!(results.len(), 2);
        let wallet_result =
            crate::query_wallets_get_by_address(database.as_ref(), "HoldingWallet111").await;
        let wallet_option = match wallet_result {
            Ok(wallet_option) => wallet_option,
            Err(error) => panic!("wallet fetch must succeed: {}", error),
        };
        let wallet = match wallet_option {
            Some(wallet) => wallet,
            None => panic!("wallet must exist"),
        };
        let wallet_id = match wallet.id {
            Some(wallet_id) => wallet_id,
            None => panic!("wallet must have an id"),
        };
        let holdings_result =
            crate::query_wallet_holdings_list_by_wallet_id(database.as_ref(), wallet_id).await;
        let holdings = match holdings_result {
            Ok(holdings) => holdings,
            Err(error) => panic!("wallet holding list must succeed: {}", error),
        };
        assert_eq!(holdings.len(), 2);
        let mut token_ids = std::vec::Vec::new();
        for holding in &holdings {
            token_ids.push(holding.token_id);
            assert_eq!(holding.last_role, Some("payer".to_string()));
        }
        token_ids.sort();
        assert_eq!(token_ids.len(), 2);
    }

    #[tokio::test]
    async fn record_transaction_by_signature_is_idempotent() {
        let database = make_database().await;
        seed_fluxbeam_create_transaction(database.clone(), "sig-wallet-holding-2").await;
        let service = crate::WalletHoldingObservationService::new(database.clone());
        let first_result = service.record_transaction_by_signature("sig-wallet-holding-2").await;
        let first_results = match first_result {
            Ok(first_results) => first_results,
            Err(error) => panic!("first wallet holding observation must succeed: {}", error),
        };
        assert_eq!(first_results.len(), 2);
        assert!(first_results[0].created_holding);
        assert!(first_results[1].created_holding);
        let second_result = service.record_transaction_by_signature("sig-wallet-holding-2").await;
        let second_results = match second_result {
            Ok(second_results) => second_results,
            Err(error) => panic!("second wallet holding observation must succeed: {}", error),
        };
        assert_eq!(second_results.len(), 2);
        assert!(!second_results[0].created_holding);
        assert!(!second_results[1].created_holding);
        let wallet_result =
            crate::query_wallets_get_by_address(database.as_ref(), "HoldingWallet111").await;
        let wallet_option = match wallet_result {
            Ok(wallet_option) => wallet_option,
            Err(error) => panic!("wallet fetch must succeed: {}", error),
        };
        let wallet = match wallet_option {
            Some(wallet) => wallet,
            None => panic!("wallet must exist"),
        };
        let wallet_id = wallet.id.unwrap_or_default();
        let holdings_result =
            crate::query_wallet_holdings_list_by_wallet_id(database.as_ref(), wallet_id).await;
        let holdings = match holdings_result {
            Ok(holdings) => holdings,
            Err(error) => panic!("wallet holding list must succeed: {}", error),
        };
        assert_eq!(holdings.len(), 2);
    }
}
