// file: kb_lib/src/ws_hybrid_observation.rs

//! Hybrid WebSocket technical observation service.

use std::hash::Hasher;

/// One hybrid WebSocket observation result.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct KbWsHybridObservationResult {
    /// Stable observation name.
    pub observation_name: std::string::String,
    /// Stable deduplication key.
    pub dedupe_key: std::string::String,
    /// Optional watched address.
    pub watched_address: std::option::Option<std::string::String>,
    /// Optional observed slot.
    pub slot: std::option::Option<u64>,
    /// Whether this observation was newly recorded during this service lifetime.
    pub created_observation: bool,
}

/// Hybrid WebSocket technical observation service.
#[derive(Debug, Clone)]
pub struct KbWsHybridObservationService {
    persistence: crate::KbDetectionPersistenceService,
    seen_dedupe_keys:
        std::sync::Arc<tokio::sync::Mutex<std::collections::HashSet<std::string::String>>>,
}

impl KbWsHybridObservationService {
    /// Creates a new hybrid WebSocket technical observation service.
    pub fn new(database: std::sync::Arc<crate::KbDatabase>) -> Self {
        let persistence = crate::KbDetectionPersistenceService::new(database.clone());
        Self {
            persistence,
            seen_dedupe_keys: std::sync::Arc::new(tokio::sync::Mutex::new(
                std::collections::HashSet::<std::string::String>::new(),
            )),
        }
    }

    /// Records one `logsNotification` payload.
    pub async fn record_logs_notification(
        &self,
        endpoint_name: std::option::Option<std::string::String>,
        payload: &serde_json::Value,
    ) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
        let signature = kb_extract_string_by_candidate_keys(payload, &["signature"]);
        self.record_observation_inner(
            "ws.hybrid.logs_notification".to_string(),
            "signal.ws.hybrid.logs_notification".to_string(),
            endpoint_name,
            signature,
            payload,
        )
        .await
    }

    /// Records one `programNotification` payload for one watched program id.
    pub async fn record_program_notification(
        &self,
        endpoint_name: std::option::Option<std::string::String>,
        watched_program_id: std::string::String,
        payload: &serde_json::Value,
    ) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
        let pubkey = kb_extract_string_by_candidate_keys(payload, &["pubkey"]);
        let watched_address = match pubkey {
            Some(pubkey) => Some(pubkey),
            None => Some(watched_program_id),
        };
        self.record_observation_inner(
            "ws.hybrid.program_notification".to_string(),
            "signal.ws.hybrid.program_notification".to_string(),
            endpoint_name,
            watched_address,
            payload,
        )
        .await
    }

    /// Records one `accountNotification` payload for one watched account address.
    pub async fn record_account_notification(
        &self,
        endpoint_name: std::option::Option<std::string::String>,
        watched_account: std::string::String,
        payload: &serde_json::Value,
    ) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
        self.record_observation_inner(
            "ws.hybrid.account_notification".to_string(),
            "signal.ws.hybrid.account_notification".to_string(),
            endpoint_name,
            Some(watched_account),
            payload,
        )
        .await
    }

    async fn record_observation_inner(
        &self,
        observation_name: std::string::String,
        signal_name: std::string::String,
        endpoint_name: std::option::Option<std::string::String>,
        watched_address: std::option::Option<std::string::String>,
        payload: &serde_json::Value,
    ) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
        let slot = kb_extract_slot(payload);
        let payload_hash = kb_hash_payload(payload);
        let dedupe_key = kb_build_ws_observation_dedupe_key(
            observation_name.as_str(),
            endpoint_name.as_deref(),
            watched_address.as_deref(),
            slot,
            payload_hash.as_str(),
        );
        let mut seen_guard = self.seen_dedupe_keys.lock().await;
        let already_seen = seen_guard.contains(&dedupe_key);
        if !already_seen {
            seen_guard.insert(dedupe_key.clone());
        }
        drop(seen_guard);
        if already_seen {
            return Ok(crate::KbWsHybridObservationResult {
                observation_name,
                dedupe_key,
                watched_address,
                slot,
                created_observation: false,
            });
        }
        let observation_payload = payload.clone();
        let observation_result = self
            .persistence
            .record_observation(&crate::KbDetectionObservationInput::new(
                observation_name.clone(),
                crate::KbObservationSourceKind::WsRpc,
                endpoint_name.clone(),
                dedupe_key.clone(),
                slot,
                observation_payload.clone(),
            ))
            .await;
        if let Err(error) = observation_result {
            return Err(error);
        }
        let signal_result = self
            .persistence
            .record_signal(&crate::KbDetectionSignalInput::new(
                signal_name,
                crate::KbAnalysisSignalSeverity::Low,
                dedupe_key.clone(),
                None,
                None,
                observation_payload,
            ))
            .await;
        if let Err(error) = signal_result {
            return Err(error);
        }
        Ok(crate::KbWsHybridObservationResult {
            observation_name,
            dedupe_key,
            watched_address,
            slot,
            created_observation: true,
        })
    }
}

fn kb_extract_slot(value: &serde_json::Value) -> std::option::Option<u64> {
    kb_extract_u64_by_candidate_keys(value, &["slot"])
}

fn kb_extract_string_by_candidate_keys(
    value: &serde_json::Value,
    candidate_keys: &[&str],
) -> std::option::Option<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 {
                    return Some(text.to_string());
                }
            }
        }
        for nested_value in object.values() {
            let nested_result = kb_extract_string_by_candidate_keys(nested_value, candidate_keys);
            if nested_result.is_some() {
                return nested_result;
            }
        }
        return None;
    }
    if let Some(array) = value.as_array() {
        for nested_value in array {
            let nested_result = kb_extract_string_by_candidate_keys(nested_value, candidate_keys);
            if nested_result.is_some() {
                return nested_result;
            }
        }
    }
    None
}

fn kb_extract_u64_by_candidate_keys(
    value: &serde_json::Value,
    candidate_keys: &[&str],
) -> std::option::Option<u64> {
    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 number_option = direct.as_u64();
                if let Some(number) = number_option {
                    return Some(number);
                }
            }
        }
        for nested_value in object.values() {
            let nested_result = kb_extract_u64_by_candidate_keys(nested_value, candidate_keys);
            if nested_result.is_some() {
                return nested_result;
            }
        }
        return None;
    }
    if let Some(array) = value.as_array() {
        for nested_value in array {
            let nested_result = kb_extract_u64_by_candidate_keys(nested_value, candidate_keys);
            if nested_result.is_some() {
                return nested_result;
            }
        }
    }
    None
}

fn kb_hash_payload(payload: &serde_json::Value) -> std::string::String {
    let payload_text_result = serde_json::to_string(payload);
    let payload_text = match payload_text_result {
        Ok(payload_text) => payload_text,
        Err(_) => return "serde_error".to_string(),
    };
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    std::hash::Hash::hash(&payload_text, &mut hasher);
    hasher.finish().to_string()
}

fn kb_build_ws_observation_dedupe_key(
    source_method: &str,
    endpoint_name: std::option::Option<&str>,
    address: std::option::Option<&str>,
    slot: std::option::Option<u64>,
    payload_hash: &str,
) -> std::string::String {
    format!(
        "{}:{}:{}:{}:{}",
        source_method,
        endpoint_name.unwrap_or_default(),
        address.unwrap_or_default(),
        slot.unwrap_or_default(),
        payload_hash
    )
}

#[cfg(test)]
mod tests {
    async fn make_database() -> std::sync::Arc<crate::KbDatabase> {
        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("ws_hybrid_observation.sqlite3");
        let config = crate::KbDatabaseConfig {
            enabled: true,
            backend: crate::KbDatabaseBackend::Sqlite,
            sqlite: crate::KbSqliteDatabaseConfig {
                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::KbDatabase::connect_and_initialize(&config).await;
        let database = match database_result {
            Ok(database) => database,
            Err(error) => panic!("database init must succeed: {}", error),
        };
        std::sync::Arc::new(database)
    }

    #[tokio::test]
    async fn record_program_notification_is_idempotent_for_same_payload() {
        let database = make_database().await;
        let service = crate::KbWsHybridObservationService::new(database);
        let payload = serde_json::json!({
            "method": "programNotification",
            "params": {
                "result": {
                    "context": {
                        "slot": 111
                    },
                    "value": {
                        "pubkey": "ProgramOwnedAccount111",
                        "account": {
                            "lamports": 1,
                            "owner": "Program111"
                        }
                    }
                }
            }
        });
        let first_result = service
            .record_program_notification(
                Some("helius_primary_ws_programs".to_string()),
                "Program111".to_string(),
                &payload,
            )
            .await;
        let first = match first_result {
            Ok(first) => first,
            Err(error) => panic!("first program notification must succeed: {}", error),
        };
        assert!(first.created_observation);
        let second_result = service
            .record_program_notification(
                Some("helius_primary_ws_programs".to_string()),
                "Program111".to_string(),
                &payload,
            )
            .await;
        let second = match second_result {
            Ok(second) => second,
            Err(error) => panic!("second program notification must succeed: {}", error),
        };
        assert!(!second.created_observation);
        assert_eq!(first.dedupe_key, second.dedupe_key);
    }

    #[tokio::test]
    async fn record_account_notification_is_idempotent_for_same_payload() {
        let database = make_database().await;
        let service = crate::KbWsHybridObservationService::new(database);
        let payload = serde_json::json!({
            "method": "accountNotification",
            "params": {
                "result": {
                    "context": {
                        "slot": 222
                    },
                    "value": {
                        "lamports": 10,
                        "owner": "Program111"
                    }
                }
            }
        });
        let first_result = service
            .record_account_notification(
                Some("mainnet_public_ws_accounts".to_string()),
                "PoolAccount111".to_string(),
                &payload,
            )
            .await;
        let first = match first_result {
            Ok(first) => first,
            Err(error) => panic!("first account notification must succeed: {}", error),
        };
        assert!(first.created_observation);
        let second_result = service
            .record_account_notification(
                Some("mainnet_public_ws_accounts".to_string()),
                "PoolAccount111".to_string(),
                &payload,
            )
            .await;
        let second = match second_result {
            Ok(second) => second,
            Err(error) => panic!("second account notification must succeed: {}", error),
        };
        assert!(!second.created_observation);
        assert_eq!(first.dedupe_key, second.dedupe_key);
    }

    #[tokio::test]
    async fn record_logs_notification_uses_signature_for_dedupe() {
        let database = make_database().await;
        let service = crate::KbWsHybridObservationService::new(database);
        let payload = serde_json::json!({
            "method": "logsNotification",
            "params": {
                "result": {
                    "context": {
                        "slot": 333
                    },
                    "value": {
                        "signature": "LogsSignature111",
                        "err": null,
                        "logs": [
                            "Program log: Instruction: InitializePool"
                        ]
                    }
                }
            }
        });
        let first_result = service
            .record_logs_notification(Some("mainnet_public_ws_logs".to_string()), &payload)
            .await;
        let first = match first_result {
            Ok(first) => first,
            Err(error) => panic!("first logs notification must succeed: {}", error),
        };
        assert!(first.created_observation);
        assert_eq!(first.watched_address, Some("LogsSignature111".to_string()));
        let second_result = service
            .record_logs_notification(Some("mainnet_public_ws_logs".to_string()), &payload)
            .await;
        let second = match second_result {
            Ok(second) => second,
            Err(error) => panic!("second logs notification must succeed: {}", error),
        };
        assert!(!second.created_observation);
        assert_eq!(first.dedupe_key, second.dedupe_key);
    }
}