// file: kb_lib/src/wallet_observation.rs

//! Wallet-observation service.

/// One wallet-observation result.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WalletObservationResult {
    /// Related wallet id.
    pub wallet_id: i64,
    /// Wallet address.
    pub wallet_address: std::string::String,
    /// Persisted wallet-participation id.
    pub wallet_participation_id: i64,
    /// Stable observed role.
    pub role: std::string::String,
    /// Optional related pool id.
    pub pool_id: std::option::Option<i64>,
    /// Optional related pair id.
    pub pair_id: std::option::Option<i64>,
    /// Whether the participation row was newly created.
    pub created_participation: bool,
}

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

impl WalletObservationService {
    /// Creates a new wallet-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 wallets and participations for one resolved transaction signature.
    pub async fn record_transaction_by_signature(
        &self,
        signature: &str,
    ) -> Result<std::vec::Vec<crate::WalletObservationResult>, 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 observations 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();
        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 observed_roles = collect_wallet_roles(&payload);
            for (role, wallet_address) in observed_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_id = match wallet_option {
                    Some(existing) => match existing.id {
                        Some(wallet_id) => {
                            let dto = crate::WalletDto::new(
                                wallet_address.clone(),
                                existing.label.clone(),
                            );
                            let upsert_result =
                                crate::query_wallets_upsert(self.database.as_ref(), &dto).await;
                            match upsert_result {
                                Ok(_) => wallet_id,
                                Err(error) => return Err(error),
                            }
                        },
                        None => {
                            return Err(crate::Error::InvalidState(
                                "wallet has no internal id".to_string(),
                            ));
                        },
                    },
                    None => {
                        let dto = crate::WalletDto::new(wallet_address.clone(), None);
                        let upsert_result =
                            crate::query_wallets_upsert(self.database.as_ref(), &dto).await;
                        match upsert_result {
                            Ok(wallet_id) => wallet_id,
                            Err(error) => return Err(error),
                        }
                    },
                };
                let participation_dto = crate::WalletParticipationDto::new(
                    wallet_id,
                    transaction_id,
                    Some(decoded_event_id),
                    pool_id,
                    pair_id,
                    role.clone(),
                    crate::ObservationSourceKind::HttpRpc,
                    transaction.source_endpoint_name.clone(),
                );
                let existing_participation_result =
                    crate::query_wallet_participations_get_by_unique_key(
                        self.database.as_ref(),
                        participation_dto.unique_key.as_str(),
                    )
                    .await;
                let existing_participation_option = match existing_participation_result {
                    Ok(existing_participation_option) => existing_participation_option,
                    Err(error) => return Err(error),
                };
                let created_participation = existing_participation_option.is_none();
                let participation_id_result = crate::query_wallet_participations_upsert(
                    self.database.as_ref(),
                    &participation_dto,
                )
                .await;
                let wallet_participation_id = match participation_id_result {
                    Ok(wallet_participation_id) => wallet_participation_id,
                    Err(error) => return Err(error),
                };
                if created_participation {
                    let payload = serde_json::json!({
                        "walletAddress": wallet_address,
                        "role": role,
                        "protocolName": decoded_event.protocol_name,
                        "eventKind": decoded_event.event_kind,
                        "poolId": pool_id,
                        "pairId": pair_id,
                        "transactionSignature": transaction.signature
                    });
                    let observation_result = self
                        .persistence
                        .record_observation(&crate::DetectionObservationInput::new(
                            "wallet.participation".to_string(),
                            crate::ObservationSourceKind::HttpRpc,
                            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.participation.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::WalletObservationResult {
                    wallet_id,
                    wallet_address,
                    wallet_participation_id,
                    role,
                    pool_id,
                    pair_id,
                    created_participation,
                });
            }
        }
        return Ok(results);
    }
}

fn collect_wallet_roles(
    payload: &serde_json::Value,
) -> std::vec::Vec<(std::string::String, std::string::String)> {
    let role_mappings = [
        ("creator", vec!["creator", "poolCreator"]),
        ("payer", vec!["payer", "funder"]),
        ("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 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);
        }
    }
}

#[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_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_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);
        let resolved_transaction = serde_json::json!({
            "slot": 930001,
            "blockTime": 1779300001,
            "version": 0,
            "transaction": {
                "message": {
                    "instructions": [
                        {
                            "programId": crate::FLUXBEAM_PROGRAM_ID,
                            "program": "fluxbeam",
                            "stackHeight": 1,
                            "accounts": [
                                "WalletObsPool111",
                                "WalletObsLpMint111",
                                "WalletObsTokenA111",
                                crate::WSOL_MINT_ID,
                                "WalletObserved111"
                            ],
                            "parsed": {
                                "info": {
                                    "instruction": "create_pool",
                                    "pool": "WalletObsPool111",
                                    "lpMint": "WalletObsLpMint111",
                                    "tokenA": "WalletObsTokenA111",
                                    "tokenB": crate::WSOL_MINT_ID,
                                    "payer": "WalletObserved111"
                                }
                            },
                            "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);
        }
    }

    #[tokio::test]
    async fn record_transaction_by_signature_creates_wallet_and_participations() {
        let database = make_database().await;
        seed_fluxbeam_transaction(database.clone(), "sig-wallet-observation-1").await;
        let service = crate::WalletObservationService::new(database.clone());
        let record_result =
            service.record_transaction_by_signature("sig-wallet-observation-1").await;
        let results = match record_result {
            Ok(results) => results,
            Err(error) => panic!("wallet observation must succeed: {}", error),
        };
        assert_eq!(results.len(), 2);
        let wallets_result = crate::query_wallets_list(database.as_ref()).await;
        let wallets = match wallets_result {
            Ok(wallets) => wallets,
            Err(error) => panic!("wallet list must succeed: {}", error),
        };
        assert_eq!(wallets.len(), 1);
        let wallet = &wallets[0];
        let wallet_id = match wallet.id {
            Some(wallet_id) => wallet_id,
            None => panic!("wallet must have an id"),
        };
        let participations_result =
            crate::query_wallet_participations_list_by_wallet_id(database.as_ref(), wallet_id)
                .await;
        let participations = match participations_result {
            Ok(participations) => participations,
            Err(error) => panic!("participation list must succeed: {}", error),
        };
        assert_eq!(participations.len(), 2);
        let mut roles = std::vec::Vec::new();
        for participation in &participations {
            roles.push(participation.role.clone());
        }
        roles.sort();
        assert_eq!(roles, vec!["creator".to_string(), "payer".to_string()]);
    }

    #[tokio::test]
    async fn record_transaction_by_signature_is_idempotent() {
        let database = make_database().await;
        seed_fluxbeam_transaction(database.clone(), "sig-wallet-observation-2").await;
        let service = crate::WalletObservationService::new(database.clone());
        let first_result =
            service.record_transaction_by_signature("sig-wallet-observation-2").await;
        let first_results = match first_result {
            Ok(first_results) => first_results,
            Err(error) => panic!("first wallet observation must succeed: {}", error),
        };
        assert_eq!(first_results.len(), 2);
        assert!(first_results[0].created_participation);
        assert!(first_results[1].created_participation);
        let second_result =
            service.record_transaction_by_signature("sig-wallet-observation-2").await;
        let second_results = match second_result {
            Ok(second_results) => second_results,
            Err(error) => panic!("second wallet observation must succeed: {}", error),
        };
        assert_eq!(second_results.len(), 2);
        assert!(!second_results[0].created_participation);
        assert!(!second_results[1].created_participation);
        let wallets_result = crate::query_wallets_list(database.as_ref()).await;
        let wallets = match wallets_result {
            Ok(wallets) => wallets,
            Err(error) => panic!("wallet list must succeed: {}", error),
        };
        assert_eq!(wallets.len(), 1);
        let wallet_id = wallets[0].id.unwrap_or_default();
        let participations_result =
            crate::query_wallet_participations_list_by_wallet_id(database.as_ref(), wallet_id)
                .await;
        let participations = match participations_result {
            Ok(participations) => participations,
            Err(error) => panic!("participation list must succeed: {}", error),
        };
        assert_eq!(participations.len(), 2);
    }
}