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

//! PumpSwap AMM transaction decoder.

const PUMP_SWAP_BUY_DISCRIMINATOR: [u8; 8] = [102, 6, 61, 18, 1, 218, 235, 234];
const PUMP_SWAP_SELL_DISCRIMINATOR: [u8; 8] = [51, 230, 133, 164, 1, 127, 131, 173];

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum PumpSwapInstructionKind {
    Buy,
    Sell,
}

#[derive(Debug, Clone)]
struct PumpSwapInstructionData {
    instruction_kind: PumpSwapInstructionKind,
    base_amount_raw: std::option::Option<std::string::String>,
    quote_amount_bound_raw: std::option::Option<std::string::String>,
}

/// Decoded PumpSwap trade event.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PumpSwapTradeDecoded {
    /// Parent transaction id.
    pub transaction_id: i64,
    /// Parent instruction id.
    pub instruction_id: i64,
    /// Transaction signature.
    pub signature: std::string::String,
    /// Program id.
    pub program_id: std::string::String,
    /// Trade side.
    pub trade_side: crate::SwapTradeSide,
    /// Optional heuristic pool account.
    pub pool_account: std::option::Option<std::string::String>,
    /// Optional token A mint.
    pub token_a_mint: std::option::Option<std::string::String>,
    /// Optional token B mint.
    pub token_b_mint: std::option::Option<std::string::String>,
    /// Optional pool_v2 account or upgraded pool account.
    pub pool_v2: std::option::Option<std::string::String>,
    /// Decoded payload.
    pub payload_json: serde_json::Value,
}

/// Decoded PumpSwap event.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum PumpSwapDecodedEvent {
    /// Buy trade.
    BuyTrade(PumpSwapTradeDecoded),
    /// Sell trade.
    SellTrade(PumpSwapTradeDecoded),
}

/// PumpSwap decoder.
#[derive(Debug, Clone, Default)]
pub struct PumpSwapDecoder;

impl PumpSwapDecoder {
    /// Creates a new decoder.
    pub fn new() -> Self {
        return Self;
    }

    /// Decodes one projected transaction into zero or more PumpSwap events.
    pub fn decode_transaction(
        &self,
        transaction: &crate::ChainTransactionDto,
        instructions: &[crate::ChainInstructionDto],
    ) -> Result<std::vec::Vec<crate::PumpSwapDecodedEvent>, crate::Error> {
        let transaction_id_option = transaction.id;
        let transaction_id = match transaction_id_option {
            Some(transaction_id) => transaction_id,
            None => {
                return Err(crate::Error::InvalidState(format!(
                    "chain transaction '{}' has no internal id",
                    transaction.signature
                )));
            },
        };
        let transaction_json_result =
            serde_json::from_str::<serde_json::Value>(transaction.transaction_json.as_str());
        let transaction_json = match transaction_json_result {
            Ok(transaction_json) => transaction_json,
            Err(error) => {
                return Err(crate::Error::Json(format!(
                    "cannot parse transaction_json for signature '{}': {}",
                    transaction.signature, error
                )));
            },
        };
        let log_messages = extract_log_messages(&transaction_json);
        let mut decoded_events = std::vec::Vec::new();
        for instruction in instructions {
            let program_id_option = &instruction.program_id;
            let program_id = match program_id_option {
                Some(program_id) => program_id,
                None => continue,
            };
            if program_id.as_str() != crate::PUMP_SWAP_PROGRAM_ID {
                continue;
            }
            let instruction_id_option = instruction.id;
            let instruction_id = match instruction_id_option {
                Some(instruction_id) => instruction_id,
                None => continue,
            };
            let accounts_result = parse_accounts_json(instruction.accounts_json.as_str());
            let accounts = match accounts_result {
                Ok(accounts) => accounts,
                Err(error) => return Err(error),
            };
            let parsed_json_result = parse_optional_parsed_json(instruction.parsed_json.as_ref());
            let parsed_json = match parsed_json_result {
                Ok(parsed_json) => parsed_json,
                Err(error) => return Err(error),
            };
            let pool_account = extract_string_by_candidate_keys(
                parsed_json.as_ref(),
                &["pool", "poolAccount", "amm", "ammPool", "poolState"],
            )
            .or_else(|| return extract_account(&accounts, 0));
            let token_a_mint = extract_string_by_candidate_keys(
                parsed_json.as_ref(),
                &["tokenAMint", "baseMint", "mintA", "coinMint", "token0Mint", "mint0"],
            )
            .or_else(|| return extract_account(&accounts, 3));
            let token_b_mint = extract_string_by_candidate_keys(
                parsed_json.as_ref(),
                &["tokenBMint", "quoteMint", "mintB", "pcMint", "token1Mint", "mint1"],
            )
            .or_else(|| return extract_account(&accounts, 4));
            let pool_v2 = extract_string_by_candidate_keys(
                parsed_json.as_ref(),
                &["poolV2", "pool_v2", "ammV2", "bondingCurveV2", "bonding_curve_v2"],
            );
            let user_base_token_account = extract_account(&accounts, 5);
            let user_quote_token_account = extract_account(&accounts, 6);
            let pool_base_token_account = extract_account(&accounts, 7);
            let pool_quote_token_account = extract_account(&accounts, 8);
            let instruction_data_result =
                parse_pump_swap_instruction_data(instruction.data_json.as_ref());
            let instruction_data_option = match instruction_data_result {
                Ok(instruction_data_option) => instruction_data_option,
                Err(error) => return Err(error),
            };
            let instruction_data = match instruction_data_option {
                Some(instruction_data) => instruction_data,
                None => continue,
            };
            let instruction_kind = match instruction_data.instruction_kind {
                PumpSwapInstructionKind::Buy => "buy",
                PumpSwapInstructionKind::Sell => "sell",
            };
            let max_quote_amount_in_raw = match instruction_data.instruction_kind {
                PumpSwapInstructionKind::Buy => instruction_data.quote_amount_bound_raw.clone(),
                PumpSwapInstructionKind::Sell => None,
            };
            let min_quote_amount_out_raw = match instruction_data.instruction_kind {
                PumpSwapInstructionKind::Buy => None,
                PumpSwapInstructionKind::Sell => instruction_data.quote_amount_bound_raw.clone(),
            };
            let payload_json = serde_json::json!({
                "decoder": "pump_swap",
                "signature": transaction.signature,
                "instructionId": instruction_id,
                "parentInstructionId": instruction.parent_instruction_id,
                "instructionIndex": instruction.instruction_index,
                "innerInstructionIndex": instruction.inner_instruction_index,
                "instructionKind": instruction_kind,
                "instructionBaseAmountRaw": instruction_data.base_amount_raw,
                "instructionQuoteAmountBoundRaw": instruction_data.quote_amount_bound_raw,
                "maxQuoteAmountInRaw": max_quote_amount_in_raw,
                "minQuoteAmountOutRaw": min_quote_amount_out_raw,
                "userBaseTokenAccount": user_base_token_account,
                "userQuoteTokenAccount": user_quote_token_account,
                "accounts": accounts,
                "parsed": parsed_json,
                "logMessages": log_messages,
                "poolAccount": pool_account,
                "tokenAMint": token_a_mint,
                "tokenBMint": token_b_mint,
                "poolV2": pool_v2,
                "poolBaseTokenAccount": pool_base_token_account,
                "poolQuoteTokenAccount": pool_quote_token_account
            });
            if instruction_data.instruction_kind == PumpSwapInstructionKind::Buy {
                decoded_events.push(crate::PumpSwapDecodedEvent::BuyTrade(
                    crate::PumpSwapTradeDecoded {
                        transaction_id,
                        instruction_id,
                        signature: transaction.signature.clone(),
                        program_id: program_id.clone(),
                        trade_side: crate::SwapTradeSide::BuyBase,
                        pool_account: pool_account.clone(),
                        token_a_mint: token_a_mint.clone(),
                        token_b_mint: token_b_mint.clone(),
                        pool_v2: pool_v2.clone(),
                        payload_json: payload_json.clone(),
                    },
                ));
            }
            if instruction_data.instruction_kind == PumpSwapInstructionKind::Sell {
                decoded_events.push(crate::PumpSwapDecodedEvent::SellTrade(
                    crate::PumpSwapTradeDecoded {
                        transaction_id,
                        instruction_id,
                        signature: transaction.signature.clone(),
                        program_id: program_id.clone(),
                        trade_side: crate::SwapTradeSide::SellBase,
                        pool_account,
                        token_a_mint,
                        token_b_mint,
                        pool_v2,
                        payload_json,
                    },
                ));
            }
        }
        return Ok(decoded_events);
    }
}

fn extract_log_messages(
    transaction_json: &serde_json::Value,
) -> std::vec::Vec<std::string::String> {
    let mut messages = std::vec::Vec::new();

    let meta_option = transaction_json.get("meta");
    let meta = match meta_option {
        Some(meta) => meta,
        None => return messages,
    };
    let logs_option = meta.get("logMessages");
    let logs = match logs_option {
        Some(logs) => logs,
        None => return messages,
    };
    let logs_array_option = logs.as_array();
    let logs_array = match logs_array_option {
        Some(logs_array) => logs_array,
        None => return messages,
    };
    for value in logs_array {
        let text_option = value.as_str();
        if let Some(text) = text_option {
            messages.push(text.to_string());
        }
    }
    return messages;
}

fn parse_accounts_json(
    accounts_json: &str,
) -> Result<std::vec::Vec<std::string::String>, crate::Error> {
    let values_result = serde_json::from_str::<std::vec::Vec<serde_json::Value>>(accounts_json);
    let values = match values_result {
        Ok(values) => values,
        Err(error) => {
            return Err(crate::Error::Json(format!(
                "cannot parse instruction accounts_json '{}': {}",
                accounts_json, error
            )));
        },
    };
    let mut accounts = std::vec::Vec::new();
    for value in values {
        let text_option = value.as_str();
        if let Some(text) = text_option {
            accounts.push(text.to_string());
        }
    }
    return Ok(accounts);
}

fn extract_account(
    extract_account: &[std::string::String],
    index: usize,
) -> std::option::Option<std::string::String> {
    if index >= extract_account.len() {
        return None;
    }
    return Some(extract_account[index].clone());
}

fn parse_optional_parsed_json(
    parsed_json: std::option::Option<&std::string::String>,
) -> Result<std::option::Option<serde_json::Value>, crate::Error> {
    let parsed_json = match parsed_json {
        Some(parsed_json) => parsed_json,
        None => return Ok(None),
    };
    let value_result = serde_json::from_str::<serde_json::Value>(parsed_json.as_str());
    match value_result {
        Ok(value) => return Ok(Some(value)),
        Err(error) => {
            return Err(crate::Error::Json(format!(
                "cannot parse instruction parsed_json '{}': {}",
                parsed_json, error
            )));
        },
    }
}

fn extract_string_by_candidate_keys(
    value: std::option::Option<&serde_json::Value>,
    candidate_keys: &[&str],
) -> std::option::Option<std::string::String> {
    let value = match value {
        Some(value) => value,
        None => return None,
    };
    return extract_string_by_candidate_keys_inner(value, candidate_keys);
}

fn extract_string_by_candidate_keys_inner(
    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 direct_text_option = direct.as_str();
                if let Some(direct_text) = direct_text_option {
                    return Some(direct_text.to_string());
                }
            }
        }
        for nested_value in object.values() {
            let nested_result =
                extract_string_by_candidate_keys_inner(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 =
                extract_string_by_candidate_keys_inner(nested_value, candidate_keys);
            if nested_result.is_some() {
                return nested_result;
            }
        }
    }
    return None;
}

fn parse_pump_swap_instruction_data(
    data_json: std::option::Option<&std::string::String>,
) -> Result<std::option::Option<PumpSwapInstructionData>, crate::Error> {
    let data_json = match data_json {
        Some(data_json) => data_json,
        None => return Ok(None),
    };
    let data_text_result = parse_instruction_data_text(data_json.as_str());
    let data_text = match data_text_result {
        Ok(data_text) => data_text,
        Err(error) => return Err(error),
    };
    if data_text.trim().is_empty() {
        return Ok(None);
    }
    let decoded_result = bs58::decode(data_text.trim()).into_vec();
    let decoded = match decoded_result {
        Ok(decoded) => decoded,
        Err(error) => {
            return Err(crate::Error::Json(format!(
                "cannot base58-decode PumpSwap instruction data '{}': {}",
                data_text, error
            )));
        },
    };
    if decoded.len() < 8 {
        return Ok(None);
    }
    let discriminator = [
        decoded[0], decoded[1], decoded[2], decoded[3], decoded[4], decoded[5], decoded[6],
        decoded[7],
    ];
    let instruction_kind = if discriminator == PUMP_SWAP_BUY_DISCRIMINATOR {
        PumpSwapInstructionKind::Buy
    } else if discriminator == PUMP_SWAP_SELL_DISCRIMINATOR {
        PumpSwapInstructionKind::Sell
    } else {
        return Ok(None);
    };
    let base_amount_raw = decode_u64_le_as_string(&decoded, 8);
    let quote_amount_bound_raw = decode_u64_le_as_string(&decoded, 16);
    return Ok(Some(PumpSwapInstructionData {
        instruction_kind,
        base_amount_raw,
        quote_amount_bound_raw,
    }));
}

fn parse_instruction_data_text(data_json: &str) -> Result<std::string::String, crate::Error> {
    let trimmed = data_json.trim();
    if trimmed.starts_with('"') {
        let parsed_result = serde_json::from_str::<std::string::String>(trimmed);
        return match parsed_result {
            Ok(parsed) => Ok(parsed),
            Err(error) => Err(crate::Error::Json(format!(
                "cannot parse instruction data JSON string '{}': {}",
                data_json, error
            ))),
        };
    }
    return Ok(trimmed.to_string());
}

fn decode_u64_le_as_string(data: &[u8], offset: usize) -> std::option::Option<std::string::String> {
    if data.len() < offset + 8 {
        return None;
    }
    let value = u64::from_le_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(value.to_string());
}

#[cfg(test)]
mod tests {
    fn make_pump_swap_buy_instruction_data_json() -> std::string::String {
        return serde_json::json!("AJTQ2h9DXrBfqJi53PDQG2Fvki5tkaTU3").to_string();
    }

    fn make_transaction_with_buy_log() -> crate::ChainTransactionDto {
        let mut dto = crate::ChainTransactionDto::new(
            "sig-pump-swap-test-1".to_string(),
            Some(777002),
            Some(1779200002),
            Some("helius_primary_http".to_string()),
            Some("0".to_string()),
            None,
            None,
            serde_json::json!({
                "slot": 777002,
                "meta": {
                    "logMessages": [
                        "Program log: Instruction: Buy"
                    ]
                },
                "transaction": {
                    "message": {
                        "instructions": []
                    }
                }
            })
            .to_string(),
        );
        dto.id = Some(92);
        return dto;
    }

    fn make_instruction() -> crate::ChainInstructionDto {
        let mut dto = crate::ChainInstructionDto::new(
            92,
            None,
            0,
            None,
            Some(crate::PUMP_SWAP_PROGRAM_ID.to_string()),
            Some("pump-amm".to_string()),
            Some(1),
            serde_json::json!([
                "PumpPool111",
                "User111",
                "GlobalConfig111",
                "TokenA111",
                "TokenB111",
                "UserBaseAta111",
                "UserQuoteAta111",
                "PoolBaseVault111",
                "PoolQuoteVault111"
            ])
            .to_string(),
            Some(make_pump_swap_buy_instruction_data_json()),
            None,
            Some(
                serde_json::json!({
                    "info": {
                        "pool": "PumpPool111",
                        "baseMint": "TokenA111",
                        "quoteMint": "TokenB111",
                        "poolV2": "PoolV2_111"
                    }
                })
                .to_string(),
            ),
        );
        dto.id = Some(18);
        return dto;
    }

    #[test]
    fn pump_swap_buy_is_detected() {
        let decoder = crate::PumpSwapDecoder::new();
        let transaction = make_transaction_with_buy_log();
        let instructions = vec![make_instruction()];
        let decoded_result = decoder.decode_transaction(&transaction, &instructions);
        let decoded = match decoded_result {
            Ok(decoded) => decoded,
            Err(error) => panic!("decode must succeed: {}", error),
        };
        assert_eq!(decoded.len(), 1);
        match &decoded[0] {
            crate::PumpSwapDecodedEvent::BuyTrade(event) => {
                assert_eq!(event.transaction_id, 92);
                assert_eq!(event.instruction_id, 18);
                assert_eq!(event.pool_account, Some("PumpPool111".to_string()));
                assert_eq!(event.token_a_mint, Some("TokenA111".to_string()));
                assert_eq!(event.token_b_mint, Some("TokenB111".to_string()));
                assert_eq!(event.pool_v2, Some("PoolV2_111".to_string()));
                assert_eq!(
                    event.payload_json.get("poolBaseTokenAccount"),
                    Some(&serde_json::Value::String("PoolBaseVault111".to_string()))
                );
                assert_eq!(
                    event.payload_json.get("poolQuoteTokenAccount"),
                    Some(&serde_json::Value::String("PoolQuoteVault111".to_string()))
                );
                assert_eq!(event.trade_side, crate::SwapTradeSide::BuyBase);
            },
            crate::PumpSwapDecodedEvent::SellTrade(_) => {
                panic!("unexpected sell event")
            },
        }
    }

    #[test]
    fn pump_swap_returns_none_without_buy_or_sell_log() {
        let decoder = crate::PumpSwapDecoder::new();
        let mut transaction = make_transaction_with_buy_log();
        transaction.transaction_json = serde_json::json!({
            "slot": 777002,
            "meta": {
                "logMessages": [
                    "Program log: Instruction: Deposit"
                ]
            },
            "transaction": {
                "message": {
                    "instructions": []
                }
            }
        })
        .to_string();
        let mut instruction = make_instruction();
        instruction.data_json = None;
        let instructions = vec![instruction];
        let decoded_result = decoder.decode_transaction(&transaction, &instructions);
        let decoded = match decoded_result {
            Ok(decoded) => decoded,
            Err(error) => panic!("decode must succeed: {}", error),
        };
        assert_eq!(decoded.len(), 0);
    }

    #[test]
    fn pump_swap_inner_buy_is_detected() {
        let decoder = crate::PumpSwapDecoder::new();
        let transaction = make_transaction_with_buy_log();
        let mut instruction = make_instruction();
        instruction.parent_instruction_id = Some(17);
        instruction.inner_instruction_index = Some(0);
        let instructions = vec![instruction];
        let decoded_result = decoder.decode_transaction(&transaction, &instructions);
        let decoded = match decoded_result {
            Ok(decoded) => decoded,
            Err(error) => panic!("decode must succeed: {}", error),
        };
        assert_eq!(decoded.len(), 1);
        match &decoded[0] {
            crate::PumpSwapDecodedEvent::BuyTrade(event) => {
                assert_eq!(event.transaction_id, 92);
                assert_eq!(event.instruction_id, 18);
                assert_eq!(event.pool_account, Some("PumpPool111".to_string()));
                assert_eq!(event.token_a_mint, Some("TokenA111".to_string()));
                assert_eq!(event.token_b_mint, Some("TokenB111".to_string()));
                assert_eq!(
                    event.payload_json.get("parentInstructionId"),
                    Some(&serde_json::Value::Number(serde_json::Number::from(17)))
                );
                assert_eq!(
                    event.payload_json.get("innerInstructionIndex"),
                    Some(&serde_json::Value::Number(serde_json::Number::from(0)))
                );
            },
            crate::PumpSwapDecodedEvent::SellTrade(_) => {
                panic!("unexpected sell event")
            },
        }
    }
}
