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

//! Raydium CPMM decoder.
//!
//! This module decodes Raydium constant product swap instructions from
//! already-projected Solana transaction instructions.

/// Raydium CPMM mainnet program id.
pub const KB_RAYDIUM_CPMM_PROGRAM_ID: &str = "CPMMoo8L3F4NbTegBCKVNunggL7H1ZpdTHKxQB5qKP1C";

/// Raydium CPMM `swap_base_input` discriminator.
const KB_RAYDIUM_CPMM_SWAP_BASE_INPUT_DISCRIMINATOR: [u8; 8] =
    [143, 190, 90, 218, 196, 30, 51, 222];

/// Raydium CPMM `swap_base_output` discriminator.
const KB_RAYDIUM_CPMM_SWAP_BASE_OUTPUT_DISCRIMINATOR: [u8; 8] =
    [55, 217, 98, 86, 163, 74, 180, 173];

/// Raydium CPMM decoded event.
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize, PartialEq, Eq)]
pub enum KbRaydiumCpmmDecodedEvent {
    /// Swap where the user fixes the input amount.
    SwapBaseInput(KbRaydiumCpmmSwapDecoded),
    /// Swap where the user fixes the output amount.
    SwapBaseOutput(KbRaydiumCpmmSwapDecoded),
}

/// Raydium CPMM swap mode.
#[derive(Clone, Copy, Debug, serde::Deserialize, serde::Serialize, PartialEq, Eq)]
pub enum KbRaydiumCpmmSwapMode {
    /// Fixed input swap.
    BaseInput,
    /// Fixed output swap.
    BaseOutput,
}

/// Raydium CPMM decoded swap.
#[derive(Clone, Debug, serde::Deserialize, serde::Serialize, PartialEq, Eq)]
pub struct KbRaydiumCpmmSwapDecoded {
    /// Instruction mode.
    pub swap_mode: KbRaydiumCpmmSwapMode,
    /// User or payer account.
    pub payer: std::string::String,
    /// Raydium authority account.
    pub authority: std::string::String,
    /// AMM config account.
    pub amm_config: std::string::String,
    /// CPMM 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,
    /// Input token program.
    pub input_token_program: std::string::String,
    /// Output token program.
    pub output_token_program: std::string::String,
    /// Input token mint.
    pub input_token_mint: std::string::String,
    /// Output token mint.
    pub output_token_mint: std::string::String,
    /// Observation state.
    pub observation_state: std::string::String,
    /// Normalized base mint.
    pub base_mint: std::string::String,
    /// Normalized quote mint.
    pub quote_mint: std::string::String,
    /// Normalized base vault.
    pub base_vault: std::string::String,
    /// Normalized quote vault.
    pub quote_vault: std::string::String,
    /// True when input mint is the normalized base mint.
    pub input_is_base: bool,
    /// Trade side from the normalized pair perspective.
    pub trade_side: std::string::String,
    /// Fixed input amount, for swap_base_input.
    pub amount_in_raw: std::option::Option<std::string::String>,
    /// Minimum output amount, for swap_base_input.
    pub minimum_amount_out_raw: std::option::Option<std::string::String>,
    /// Maximum input amount, for swap_base_output.
    pub max_amount_in_raw: std::option::Option<std::string::String>,
    /// Fixed output amount, for swap_base_output.
    pub amount_out_raw: std::option::Option<std::string::String>,
}

impl KbRaydiumCpmmDecodedEvent {
    /// Returns the storage event kind.
    pub fn event_kind(&self) -> &'static str {
        match self {
            KbRaydiumCpmmDecodedEvent::SwapBaseInput(_) => return "raydium_cpmm.swap_base_input",
            KbRaydiumCpmmDecodedEvent::SwapBaseOutput(_) => return "raydium_cpmm.swap_base_output",
        }
    }

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

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

    /// Returns the normalized quote mint.
    pub fn quote_mint(&self) -> &str {
        match self {
            KbRaydiumCpmmDecodedEvent::SwapBaseInput(event) => return event.quote_mint.as_str(),
            KbRaydiumCpmmDecodedEvent::SwapBaseOutput(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::KbRaydiumCpmmDecodedEvent::SwapBaseInput(event) => {
                let result = serde_json::to_string(event);
                match result {
                    Ok(payload) => return Some(payload),
                    Err(_) => return None,
                }
            },
            crate::KbRaydiumCpmmDecodedEvent::SwapBaseOutput(event) => {
                let result = serde_json::to_string(event);
                match result {
                    Ok(payload) => return Some(payload),
                    Err(_) => return None,
                }
            },
        }
    }
}

/// Decodes one Raydium CPMM instruction from projected instruction fields.
pub fn kb_decode_raydium_cpmm_instruction(
    accounts_json: &str,
    data_json: &str,
) -> std::vec::Vec<KbRaydiumCpmmDecodedEvent> {
    let accounts = match kb_parse_accounts_json(accounts_json) {
        Some(accounts) => accounts,
        None => return std::vec::Vec::new(),
    };
    let data_base58 = match kb_parse_data_json_as_base58(data_json) {
        Some(data_base58) => data_base58,
        None => return std::vec::Vec::new(),
    };
    let data = match bs58::decode(data_base58.as_str()).into_vec() {
        Ok(data) => data,
        Err(_) => return std::vec::Vec::new(),
    };
    if data.len() < 24 {
        return std::vec::Vec::new();
    }
    let discriminator = [data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7]];
    if discriminator == KB_RAYDIUM_CPMM_SWAP_BASE_INPUT_DISCRIMINATOR {
        let amount_in = match kb_read_u64_le(data.as_slice(), 8) {
            Some(value) => value,
            None => return std::vec::Vec::new(),
        };
        let minimum_amount_out = match kb_read_u64_le(data.as_slice(), 16) {
            Some(value) => value,
            None => return std::vec::Vec::new(),
        };
        let swap = match kb_build_raydium_cpmm_swap(
            KbRaydiumCpmmSwapMode::BaseInput,
            accounts.as_slice(),
            Some(amount_in.to_string()),
            Some(minimum_amount_out.to_string()),
            None,
            None,
        ) {
            Some(swap) => swap,
            None => return std::vec::Vec::new(),
        };
        return vec![KbRaydiumCpmmDecodedEvent::SwapBaseInput(swap)];
    }
    if discriminator == KB_RAYDIUM_CPMM_SWAP_BASE_OUTPUT_DISCRIMINATOR {
        let max_amount_in = match kb_read_u64_le(data.as_slice(), 8) {
            Some(value) => value,
            None => return std::vec::Vec::new(),
        };
        let amount_out = match kb_read_u64_le(data.as_slice(), 16) {
            Some(value) => value,
            None => return std::vec::Vec::new(),
        };
        let swap = match kb_build_raydium_cpmm_swap(
            KbRaydiumCpmmSwapMode::BaseOutput,
            accounts.as_slice(),
            None,
            None,
            Some(max_amount_in.to_string()),
            Some(amount_out.to_string()),
        ) {
            Some(swap) => swap,
            None => return std::vec::Vec::new(),
        };
        return vec![KbRaydiumCpmmDecodedEvent::SwapBaseOutput(swap)];
    }
    return std::vec::Vec::new();
}

fn kb_build_raydium_cpmm_swap(
    swap_mode: KbRaydiumCpmmSwapMode,
    accounts: &[std::string::String],
    amount_in_raw: std::option::Option<std::string::String>,
    minimum_amount_out_raw: std::option::Option<std::string::String>,
    max_amount_in_raw: std::option::Option<std::string::String>,
    amount_out_raw: std::option::Option<std::string::String>,
) -> std::option::Option<KbRaydiumCpmmSwapDecoded> {
    if accounts.len() < 13 {
        return None;
    }
    let input_token_mint = accounts[10].clone();
    let output_token_mint = accounts[11].clone();
    let input_vault = accounts[6].clone();
    let output_vault = accounts[7].clone();
    let normalized = kb_normalize_raydium_cpmm_pair(
        input_token_mint.as_str(),
        output_token_mint.as_str(),
        input_vault.as_str(),
        output_vault.as_str(),
    );
    let input_is_base = normalized.input_is_base;
    let trade_side = if input_is_base { "sell".to_string() } else { "buy".to_string() };
    return Some(KbRaydiumCpmmSwapDecoded {
        swap_mode,
        payer: accounts[0].clone(),
        authority: accounts[1].clone(),
        amm_config: accounts[2].clone(),
        pool_state: accounts[3].clone(),
        input_token_account: accounts[4].clone(),
        output_token_account: accounts[5].clone(),
        input_vault,
        output_vault,
        input_token_program: accounts[8].clone(),
        output_token_program: accounts[9].clone(),
        input_token_mint,
        output_token_mint,
        observation_state: accounts[12].clone(),
        base_mint: normalized.base_mint,
        quote_mint: normalized.quote_mint,
        base_vault: normalized.base_vault,
        quote_vault: normalized.quote_vault,
        input_is_base,
        trade_side,
        amount_in_raw,
        minimum_amount_out_raw,
        max_amount_in_raw,
        amount_out_raw,
    });
}

struct KbRaydiumCpmmNormalizedPair {
    base_mint: std::string::String,
    quote_mint: std::string::String,
    base_vault: std::string::String,
    quote_vault: std::string::String,
    input_is_base: bool,
}

fn kb_normalize_raydium_cpmm_pair(
    input_mint: &str,
    output_mint: &str,
    input_vault: &str,
    output_vault: &str,
) -> KbRaydiumCpmmNormalizedPair {
    if kb_is_quote_mint(output_mint) && !kb_is_quote_mint(input_mint) {
        return KbRaydiumCpmmNormalizedPair {
            base_mint: input_mint.to_string(),
            quote_mint: output_mint.to_string(),
            base_vault: input_vault.to_string(),
            quote_vault: output_vault.to_string(),
            input_is_base: true,
        };
    }
    if kb_is_quote_mint(input_mint) && !kb_is_quote_mint(output_mint) {
        return KbRaydiumCpmmNormalizedPair {
            base_mint: output_mint.to_string(),
            quote_mint: input_mint.to_string(),
            base_vault: output_vault.to_string(),
            quote_vault: input_vault.to_string(),
            input_is_base: false,
        };
    }
    if input_mint <= output_mint {
        return KbRaydiumCpmmNormalizedPair {
            base_mint: input_mint.to_string(),
            quote_mint: output_mint.to_string(),
            base_vault: input_vault.to_string(),
            quote_vault: output_vault.to_string(),
            input_is_base: true,
        };
    }
    return KbRaydiumCpmmNormalizedPair {
        base_mint: output_mint.to_string(),
        quote_mint: input_mint.to_string(),
        base_vault: output_vault.to_string(),
        quote_vault: input_vault.to_string(),
        input_is_base: false,
    };
}

fn kb_is_quote_mint(mint: &str) -> bool {
    if mint == "So11111111111111111111111111111111111111112" {
        return true;
    }
    if mint == "EPjFWdd5AufqSSqeM2qN1xzybapC8G4wEGGkZwyTDt1v" {
        return true;
    }
    if mint == "Es9vMFrzaCERmJfrF4H2FYD4KMZbsJyNeYrBbqD6CbK" {
        return true;
    }
    if mint == "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB" {
        return true;
    }
    return false;
}

fn kb_parse_accounts_json(
    accounts_json: &str,
) -> std::option::Option<std::vec::Vec<std::string::String>> {
    let result = serde_json::from_str::<std::vec::Vec<std::string::String>>(accounts_json);
    match result {
        Ok(accounts) => return Some(accounts),
        Err(_) => return None,
    }
}

fn kb_parse_data_json_as_base58(data_json: &str) -> std::option::Option<std::string::String> {
    let json_string_result = serde_json::from_str::<std::string::String>(data_json);
    if let Ok(value) = json_string_result {
        return Some(value);
    }
    let trimmed = data_json.trim();
    if trimmed.is_empty() {
        return None;
    }
    let without_quotes = trimmed.trim_matches('"');
    if without_quotes.is_empty() {
        return None;
    }
    return Some(without_quotes.to_string());
}

fn kb_read_u64_le(data: &[u8], offset: usize) -> std::option::Option<u64> {
    if data.len() < offset + 8 {
        return None;
    }
    let bytes = [
        data[offset],
        data[offset + 1],
        data[offset + 2],
        data[offset + 3],
        data[offset + 4],
        data[offset + 5],
        data[offset + 6],
        data[offset + 7],
    ];
    return Some(u64::from_le_bytes(bytes));
}

#[cfg(test)]
mod tests {
    #[test]
    fn decodes_swap_base_input() {
        let accounts_json = r#"[
            "ARu4n5mFdZogZAravu7CcizaojWnS6oqka37gdLT5SZn",
            "GpMZbSM2GgvTKHJirzeGfMFoaZ8UR2X7F4v8vHTvxFbL",
            "GtdipqAcw8eAYxGcANs3vpN7UaN7UH7u89Kd8wPqLThd",
            "2ErXvV1tKtG3wiHqdofDjMou7Jusdsfasvfh8HrTj5oV",
            "3tBz6MmUTaf2QdnAnYYJGDpgbTXdANP9v4CtUL42GZaZ",
            "C1qMskwQhQ4cn1M6RG8uWfUuMrq6ysRPEHccGxk4tjw8",
            "8aHJfDQBmngBpFwLysD1bhj8LKh7ywbxfUxsy7VQkx22",
            "A4jPLDGrkAogPgm8KqQzZQ4FoNJuKuKMCbNqgQkpT5Ci",
            "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
            "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
            "Pf9aSicGu3g6tTUBqrRbjNsGape9HopibspX5KSbonk",
            "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB",
            "9fzQ17bEnqJSsyHL5CodJqY2sdjZJBQCQLbgFZ1BYTWn"
        ]"#;
        let events = crate::kb_decode_raydium_cpmm_instruction(
            accounts_json,
            r#""E73fXHPWvSRF4Q2ZQFvPoeJBVGDUEMmxB""#,
        );
        assert_eq!(events.len(), 1);
        match &events[0] {
            crate::KbRaydiumCpmmDecodedEvent::SwapBaseInput(event) => {
                assert_eq!(event.pool_state, "2ErXvV1tKtG3wiHqdofDjMou7Jusdsfasvfh8HrTj5oV");
                assert_eq!(event.base_mint, "Pf9aSicGu3g6tTUBqrRbjNsGape9HopibspX5KSbonk");
                assert_eq!(event.quote_mint, "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB");
                assert!(event.input_is_base);
                assert_eq!(event.trade_side, "sell");
                assert!(event.amount_in_raw.is_some());
                assert!(event.minimum_amount_out_raw.is_some());
            },
            _ => {
                panic!("expected swap base input");
            },
        }
    }
    #[test]
    fn decodes_swap_base_output() {
        let accounts_json = r#"[
            "JD6rVaerbyz6wjQ433nrw6bFTgFrp46MiYmi8EtUAfsG",
            "GpMZbSM2GgvTKHJirzeGfMFoaZ8UR2X7F4v8vHTvxFbL",
            "GtdipqAcw8eAYxGcANs3vpN7UaN7UH7u89Kd8wPqLThd",
            "2ErXvV1tKtG3wiHqdofDjMou7Jusdsfasvfh8HrTj5oV",
            "4Fg913oAKp7BiW6WvUF32RvACsBYusMTQ16nmA6Sd9tp",
            "Cvcy56pbaDms71KyxAWFTQ24n5KbCy8KNa7zMhKGcJa6",
            "A4jPLDGrkAogPgm8KqQzZQ4FoNJuKuKMCbNqgQkpT5Ci",
            "8aHJfDQBmngBpFwLysD1bhj8LKh7ywbxfUxsy7VQkx22",
            "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
            "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA",
            "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB",
            "Pf9aSicGu3g6tTUBqrRbjNsGape9HopibspX5KSbonk",
            "9fzQ17bEnqJSsyHL5CodJqY2sdjZJBQCQLbgFZ1BYTWn"
        ]"#;
        let events = crate::kb_decode_raydium_cpmm_instruction(
            accounts_json,
            r#""66JafaVu7KNEtTngqQyD7ETVurF3rxJ47""#,
        );
        assert_eq!(events.len(), 1);
        match &events[0] {
            crate::KbRaydiumCpmmDecodedEvent::SwapBaseOutput(event) => {
                assert_eq!(event.pool_state, "2ErXvV1tKtG3wiHqdofDjMou7Jusdsfasvfh8HrTj5oV");
                assert_eq!(event.base_mint, "Pf9aSicGu3g6tTUBqrRbjNsGape9HopibspX5KSbonk");
                assert_eq!(event.quote_mint, "USD1ttGY1N17NEEHLmELoaybftRBUSErhqYiQzvEmuB");
                assert!(!event.input_is_base);
                assert_eq!(event.trade_side, "buy");
                assert!(event.max_amount_in_raw.is_some());
                assert!(event.amount_out_raw.is_some());
            },
            _ => {
                panic!("expected swap base output");
            },
        }
    }
}