// file: kb_lib/src/dex/raydium_clmm.rs

//! Raydium CLMM instruction decoder.

const RAYDIUM_CLMM_SWAP_V2_DISCRIMINATOR: [u8; 8] = [43, 4, 237, 11, 26, 201, 30, 98];

const RAYDIUM_CLMM_SWAP_LEGACY_DISCRIMINATOR: [u8; 8] = [248, 198, 158, 145, 225, 117, 135, 200];

/// Decoded Raydium CLMM event.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq)]
pub enum RaydiumClmmDecodedEvent {
    /// Raydium CLMM swap_v2 event.
    SwapV2(crate::RaydiumClmmSwapV2Decoded),
}

impl RaydiumClmmDecodedEvent {
    /// Returns the normalized event kind.
    pub fn event_kind(&self) -> &'static str {
        match self {
            crate::RaydiumClmmDecodedEvent::SwapV2(_) => return "raydium_clmm.swap_v2",
        }
    }

    /// Returns the pool account.
    pub fn pool_account(&self) -> &str {
        match self {
            crate::RaydiumClmmDecodedEvent::SwapV2(event) => return event.pool_state.as_str(),
        }
    }

    /// Returns the normalized base mint.
    pub fn base_mint(&self) -> &str {
        match self {
            crate::RaydiumClmmDecodedEvent::SwapV2(event) => return event.base_mint.as_str(),
        }
    }

    /// Returns the normalized quote mint.
    pub fn quote_mint(&self) -> &str {
        match self {
            crate::RaydiumClmmDecodedEvent::SwapV2(event) => return event.quote_mint.as_str(),
        }
    }

    /// Converts the decoded event to JSON payload.
    pub fn to_payload_json(&self) -> std::option::Option<std::string::String> {
        match self {
            crate::RaydiumClmmDecodedEvent::SwapV2(event) => {
                let result = serde_json::to_string(event);
                match result {
                    Ok(payload_json) => return Some(payload_json),
                    Err(_) => return None,
                }
            },
        }
    }
}

/// Decoded Raydium CLMM swap_v2 instruction.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq)]
pub struct RaydiumClmmSwapV2Decoded {
    /// User performing the swap.
    pub payer: std::string::String,
    /// AMM config account.
    pub amm_config: std::string::String,
    /// CLMM pool state account.
    pub pool_state: std::string::String,
    /// User input token account.
    pub input_token_account: std::string::String,
    /// User output token account.
    pub output_token_account: std::string::String,
    /// Pool input vault.
    pub input_vault: std::string::String,
    /// Pool output vault.
    pub output_vault: std::string::String,
    /// Pool oracle observation state.
    pub observation_state: std::string::String,
    /// Input vault mint.
    pub input_vault_mint: std::string::String,
    /// Output vault mint.
    pub output_vault_mint: std::string::String,
    /// Canonical base mint.
    pub base_mint: std::string::String,
    /// Canonical quote mint.
    pub quote_mint: std::string::String,
    /// Canonical base vault.
    pub base_vault: std::string::String,
    /// Canonical quote vault.
    pub quote_vault: std::string::String,
    /// Trade side relative to the canonical base mint.
    #[serde(rename = "tradeSide")]
    pub trade_side: std::string::String,
    /// Amount argument.
    pub amount: u64,
    /// Other amount threshold argument.
    pub other_amount_threshold: u64,
    /// Sqrt price limit as decimal string.
    pub sqrt_price_limit_x64: std::string::String,
    /// Whether the instruction uses exact input mode.
    pub is_base_input: bool,
}

/// Decodes a Raydium CLMM instruction.
pub fn decode_raydium_clmm_instruction(
    accounts_json: &str,
    data_json: &str,
) -> std::vec::Vec<crate::RaydiumClmmDecodedEvent> {
    let mut decoded = std::vec::Vec::new();
    let accounts_result = serde_json::from_str::<std::vec::Vec<std::string::String>>(accounts_json);
    let accounts = match accounts_result {
        Ok(accounts) => accounts,
        Err(_) => return decoded,
    };
    let data_base58_result = serde_json::from_str::<std::string::String>(data_json);
    let data_base58 = match data_base58_result {
        Ok(data_base58) => data_base58,
        Err(_) => data_json.to_string(),
    };
    let data_option = decode_base58(data_base58.as_str());
    let data = match data_option {
        Some(data) => data,
        None => return decoded,
    };
    if data.len() < 41 {
        return decoded;
    }
    let discriminator_option = read_discriminator(data.as_slice());
    let discriminator = match discriminator_option {
        Some(discriminator) => discriminator,
        None => return decoded,
    };
    if discriminator == RAYDIUM_CLMM_SWAP_LEGACY_DISCRIMINATOR {
        return decoded;
    }
    if discriminator != RAYDIUM_CLMM_SWAP_V2_DISCRIMINATOR {
        return decoded;
    }
    let event_option = decode_swap_v2(accounts.as_slice(), data.as_slice());
    let event = match event_option {
        Some(event) => event,
        None => return decoded,
    };
    decoded.push(crate::RaydiumClmmDecodedEvent::SwapV2(event));
    return decoded;
}

fn decode_swap_v2(
    accounts: &[std::string::String],
    data: &[u8],
) -> std::option::Option<crate::RaydiumClmmSwapV2Decoded> {
    let payer = match clone_account(accounts, 0) {
        Some(value) => value,
        None => return None,
    };
    let amm_config = match clone_account(accounts, 1) {
        Some(value) => value,
        None => return None,
    };
    let pool_state = match clone_account(accounts, 2) {
        Some(value) => value,
        None => return None,
    };
    let input_token_account = match clone_account(accounts, 3) {
        Some(value) => value,
        None => return None,
    };
    let output_token_account = match clone_account(accounts, 4) {
        Some(value) => value,
        None => return None,
    };
    let input_vault = match clone_account(accounts, 5) {
        Some(value) => value,
        None => return None,
    };
    let output_vault = match clone_account(accounts, 6) {
        Some(value) => value,
        None => return None,
    };
    let observation_state = match clone_account(accounts, 7) {
        Some(value) => value,
        None => return None,
    };
    let input_vault_mint = match clone_account(accounts, 11) {
        Some(value) => value,
        None => return None,
    };
    let output_vault_mint = match clone_account(accounts, 12) {
        Some(value) => value,
        None => return None,
    };
    let amount = match read_u64_le(data, 8) {
        Some(value) => value,
        None => return None,
    };
    let other_amount_threshold = match read_u64_le(data, 16) {
        Some(value) => value,
        None => return None,
    };
    let sqrt_price_limit_x64 = match read_u128_le(data, 24) {
        Some(value) => value,
        None => return None,
    };
    let is_base_input = match read_bool(data, 40) {
        Some(value) => value,
        None => return None,
    };
    let mut base_mint = input_vault_mint.clone();
    let mut quote_mint = output_vault_mint.clone();
    let mut base_vault = input_vault.clone();
    let mut quote_vault = output_vault.clone();
    let mut trade_side = "SellBase".to_string();
    if output_vault_mint.as_str() < input_vault_mint.as_str() {
        base_mint = output_vault_mint.clone();
        quote_mint = input_vault_mint.clone();
        base_vault = output_vault.clone();
        quote_vault = input_vault.clone();
        trade_side = "BuyBase".to_string();
    }
    return Some(crate::RaydiumClmmSwapV2Decoded {
        payer,
        amm_config,
        pool_state,
        input_token_account,
        output_token_account,
        input_vault,
        output_vault,
        observation_state,
        input_vault_mint,
        output_vault_mint,
        base_mint,
        quote_mint,
        base_vault,
        quote_vault,
        trade_side,
        amount,
        other_amount_threshold,
        sqrt_price_limit_x64: sqrt_price_limit_x64.to_string(),
        is_base_input,
    });
}

fn clone_account(
    accounts: &[std::string::String],
    index: usize,
) -> std::option::Option<std::string::String> {
    let account_option = accounts.get(index);
    match account_option {
        Some(account) => return Some(account.clone()),
        None => return None,
    }
}

fn read_discriminator(data: &[u8]) -> std::option::Option<[u8; 8]> {
    if data.len() < 8 {
        return None;
    }
    let mut bytes = [0_u8; 8];
    let mut index = 0_usize;
    while index < 8 {
        bytes[index] = data[index];
        index += 1;
    }
    return Some(bytes);
}

fn read_u64_le(data: &[u8], offset: usize) -> std::option::Option<u64> {
    if data.len() < offset + 8 {
        return None;
    }
    let mut bytes = [0_u8; 8];
    let mut index = 0_usize;
    while index < 8 {
        bytes[index] = data[offset + index];
        index += 1;
    }
    return Some(u64::from_le_bytes(bytes));
}

fn read_u128_le(data: &[u8], offset: usize) -> std::option::Option<u128> {
    if data.len() < offset + 16 {
        return None;
    }
    let mut bytes = [0_u8; 16];
    let mut index = 0_usize;
    while index < 16 {
        bytes[index] = data[offset + index];
        index += 1;
    }
    return Some(u128::from_le_bytes(bytes));
}

fn read_bool(data: &[u8], offset: usize) -> std::option::Option<bool> {
    if data.len() <= offset {
        return None;
    }
    match data[offset] {
        0 => return Some(false),
        1 => return Some(true),
        _ => return None,
    }
}

fn decode_base58(input: &str) -> std::option::Option<std::vec::Vec<u8>> {
    match bs58::decode(input).into_vec() {
        Ok(decoded) => return Some(decoded),
        Err(_) => return None,
    }
}

#[cfg(test)]
mod tests {
    fn sample_swap_v2_accounts_json() -> &'static str {
        return r#"[
            "8NQ32SyFKD1d5kenq4oM8Da6C6J9TQSMW1uAgFRveEQr",
            "A1BBtTYJd4i3xU8D6Tc2FzU6ZN4oXZWXKZnCxwbHXr8x",
            "GUrRxvnWVQSnbcz1eP9D5BqXwPZtRhmrqVfm5wY9meWR",
            "D2frZyyQ7NQaXRiEoBUM9S64Ckv7KZ7wuqupqdMhpsHy",
            "H7qe6sAyEyqztyMtRrDf5J1gugLx6yuyKPy5veVmR14W",
            "AvRzvwpSVnxsinLGQS3vZLqkZxhXZDM8F2qKccAo7rSq",
            "CTkc4xDrpzjWcFLC1cxmUZZjZLSRV46HZa8wu5eKTbuh",
            "8QtFSxNzD3zmEX8nzQKZB83TH4WGUAkLkQoRHAw5fuhn",
            "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
            "TokenzQdBNbLqP5VEhdkAS6EPFLC1PHnBqCXEpPxuEb",
            "MemoSq4gqABAXKb96qnH8TysNcWxMyWCqXgDLGmfcHr",
            "CKvjP8FrZpaKXjASEtX2nEU9w7M4RKskfnLQbKJBodV",
            "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs",
            "8ovxZR2Gv9Mr73aoXLQYTMaZvHCSpnEohzgjVHQwmyHr",
            "9MssDxndh2Rn8DmGWL94hXVv22zxfDYHV7tvzfPgcaWe"
        ]"#;
    }

    #[test]
    fn decodes_swap_v2() {
        let events = crate::decode_raydium_clmm_instruction(
            sample_swap_v2_accounts_json(),
            r#""ASCsAbe1UnDnCsnGLPALJUXSS5JREycfhGyTzKh7xRWNyRHCqBuzR23S""#,
        );
        assert_eq!(events.len(), 1);
        match &events[0] {
            crate::RaydiumClmmDecodedEvent::SwapV2(event) => {
                assert_eq!(events[0].event_kind(), "raydium_clmm.swap_v2");
                assert_eq!(
                    events[0].pool_account(),
                    "GUrRxvnWVQSnbcz1eP9D5BqXwPZtRhmrqVfm5wY9meWR"
                );
                assert_eq!(event.pool_state, "GUrRxvnWVQSnbcz1eP9D5BqXwPZtRhmrqVfm5wY9meWR");
                assert_eq!(event.input_vault, "AvRzvwpSVnxsinLGQS3vZLqkZxhXZDM8F2qKccAo7rSq");
                assert_eq!(event.output_vault, "CTkc4xDrpzjWcFLC1cxmUZZjZLSRV46HZa8wu5eKTbuh");
                assert_eq!(event.input_vault_mint, "CKvjP8FrZpaKXjASEtX2nEU9w7M4RKskfnLQbKJBodV");
                assert_eq!(event.output_vault_mint, "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs");
                assert_eq!(event.base_mint, "7vfCXTUXx5WJV5JADk17DUJ4ksgau7utNKj4b963voxs");
                assert_eq!(event.quote_mint, "CKvjP8FrZpaKXjASEtX2nEU9w7M4RKskfnLQbKJBodV");
                assert_eq!(event.base_vault, "CTkc4xDrpzjWcFLC1cxmUZZjZLSRV46HZa8wu5eKTbuh");
                assert_eq!(event.quote_vault, "AvRzvwpSVnxsinLGQS3vZLqkZxhXZDM8F2qKccAo7rSq");
                assert_eq!(event.trade_side, "BuyBase");
                assert_eq!(event.amount, 148441657491969);
                assert_eq!(event.other_amount_threshold, 0);
                assert_eq!(event.sqrt_price_limit_x64, "0");
                assert!(event.is_base_input);
            },
        }
    }

    #[test]
    fn serializes_swap_v2_payload_json() {
        let events = crate::decode_raydium_clmm_instruction(
            sample_swap_v2_accounts_json(),
            r#""ASCsAbe1UnDnCsnGLPALJUXSS5JREycfhGyTzKh7xRWNyRHCqBuzR23S""#,
        );
        assert_eq!(events.len(), 1);
        let payload_option = events[0].to_payload_json();
        let payload = match payload_option {
            Some(payload) => payload,
            None => panic!("payload json must be available"),
        };
        assert!(payload.contains("GUrRxvnWVQSnbcz1eP9D5BqXwPZtRhmrqVfm5wY9meWR"));
        assert!(payload.contains("input_vault"));
        assert!(payload.contains("output_vault"));
        assert!(payload.contains("tradeSide"));
    }

    #[test]
    fn ignores_invalid_data() {
        let events = crate::decode_raydium_clmm_instruction(
            sample_swap_v2_accounts_json(),
            r#""not-base58-data-0""#,
        );
        assert_eq!(events.len(), 0);
    }

    #[test]
    fn ignores_incomplete_accounts() {
        let accounts_json = r#"[
            "8NQ32SyFKD1d5kenq4oM8Da6C6J9TQSMW1uAgFRveEQr",
            "A1BBtTYJd4i3xU8D6Tc2FzU6ZN4oXZWXKZnCxwbHXr8x",
            "GUrRxvnWVQSnbcz1eP9D5BqXwPZtRhmrqVfm5wY9meWR"
        ]"#;
        let events = crate::decode_raydium_clmm_instruction(
            accounts_json,
            r#""ASCsAbe1UnDnCsnGLPALJUXSS5JREycfhGyTzKh7xRWNyRHCqBuzR23S""#,
        );
        assert_eq!(events.len(), 0);
    }

    #[test]
    fn ignores_legacy_swap_for_now() {
        let mut data: std::vec::Vec<u8> = vec![248, 198, 158, 145, 225, 117, 135, 200];
        while data.len() < 41 {
            data.push(0);
        }
        data[40] = 1;
        let encoded = bs58::encode(data).into_string();
        let data_json = format!("\"{}\"", encoded);
        let events = crate::decode_raydium_clmm_instruction(
            sample_swap_v2_accounts_json(),
            data_json.as_str(),
        );
        assert_eq!(events.len(), 0);
    }
}
