#![cfg(feature = "solana-instruction-convert-v2")]
use std::collections::HashMap;

use serde::Deserialize;
use solana_instruction_v2::{AccountMeta, Instruction};
use solana_pubkey_v2::Pubkey;

use titan_api_types::ws::v1;

mod common;

#[derive(Debug, Deserialize, Clone, Copy, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
struct OldAccountMeta {
    #[serde(with = "common::serialize_as_str")]
    pub pubkey: Pubkey,
    pub is_signer: bool,
    pub is_writable: bool,
}

impl From<OldAccountMeta> for AccountMeta {
    fn from(value: OldAccountMeta) -> Self {
        Self {
            pubkey: value.pubkey,
            is_signer: value.is_signer,
            is_writable: value.is_writable,
        }
    }
}

impl From<OldAccountMeta> for titan_api_types::common::AccountMeta {
    fn from(value: OldAccountMeta) -> Self {
        Self {
            pubkey: value.pubkey.into(),
            is_signer: value.is_signer,
            is_writable: value.is_writable,
        }
    }
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
struct OldQuoteInstruction {
    #[serde(with = "common::serialize_as_str")]
    pub program_id: Pubkey,
    pub accounts: Vec<OldAccountMeta>,
    pub data: Vec<u8>,
}

impl From<OldQuoteInstruction> for Instruction {
    fn from(value: OldQuoteInstruction) -> Self {
        Self {
            program_id: value.program_id,
            accounts: value.accounts.into_iter().map(Into::into).collect(),
            data: value.data,
        }
    }
}

impl From<OldQuoteInstruction> for titan_api_types::common::Instruction {
    fn from(value: OldQuoteInstruction) -> Self {
        Self {
            program_id: value.program_id.into(),
            accounts: value.accounts.into_iter().map(Into::into).collect(),
            data: value.data,
        }
    }
}

// Old quote response for a single provider.
// The ndjson files contain one of these per line.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct OldQuoteResponse {
    pub exchange: String,
    // JSON-encoded quote data.
    pub data: String,
    //pub success: bool,
    //pub request_timestamp: String,
    //pub response_timestamp: String,
}

#[derive(Deserialize, Clone, Copy, Debug)]
#[serde(rename_all = "camelCase")]
struct PlatformFee {
    #[serde(with = "common::serialize_as_str")]
    pub amount: u64,
    pub fee_bps: u8,
}

impl From<PlatformFee> for v1::PlatformFee {
    fn from(value: PlatformFee) -> Self {
        Self {
            amount: value.amount,
            fee_bps: value.fee_bps,
        }
    }
}

type RoutePlanWithMetadata = Vec<RoutePlanStep>;

#[derive(Deserialize, Debug, PartialEq, Clone)]
#[serde(rename_all = "camelCase")]
struct RoutePlanStep {
    pub swap_info: SwapInfo,
    pub percent: f32,
}

#[derive(Deserialize, Clone, Debug, Default, PartialEq)]
#[serde(rename_all = "camelCase")]
struct SwapInfo {
    #[serde(with = "common::serialize_as_str")]
    pub amm_key: Pubkey,
    pub label: String,
    #[serde(with = "common::serialize_as_str")]
    pub input_mint: Pubkey,
    #[serde(with = "common::serialize_as_str")]
    pub output_mint: Pubkey,
    /// An estimation of the input amount into the AMM
    #[serde(with = "common::serialize_as_str")]
    pub in_amount: u64,
    /// An estimation of the output amount into the AMM
    #[serde(with = "common::serialize_as_str")]
    pub out_amount: u64,
    pub context_slot: Option<u64>,
}

impl From<RoutePlanStep> for v1::RoutePlanStep {
    fn from(value: RoutePlanStep) -> Self {
        Self {
            amm_key: value.swap_info.amm_key.into(),
            label: value.swap_info.label,
            input_mint: value.swap_info.input_mint.into(),
            output_mint: value.swap_info.output_mint.into(),
            in_amount: value.swap_info.in_amount,
            out_amount: value.swap_info.out_amount,
            alloc_ppb: (value.percent * 10_000_000.0) as u32,
            fee_mint: None,
            fee_amount: None,
            context_slot: value.swap_info.context_slot,
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct OldQuoteInfo {
    #[serde(with = "common::serialize_as_str")]
    pub input_mint: Pubkey,
    #[serde(with = "common::serialize_as_str")]
    pub in_amount: u64,
    #[serde(with = "common::serialize_as_str")]
    pub output_mint: Pubkey,
    #[serde(with = "common::serialize_as_str")]
    pub out_amount: u64,
    pub swap_mode: v1::SwapMode,
    pub slippage_bps: u16,
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub computed_auto_slippage: Option<u16>,
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub uses_quote_minimizing_slippage: Option<bool>,
    pub platform_fee: Option<PlatformFee>,
    //#[serde(with = "rust_decimal::serde::str")]
    //pub price_impact_pct: Decimal,
    pub route_plan: RoutePlanWithMetadata,
    #[serde(default)]
    pub context_slot: u64,
    #[serde(default)]
    pub time_taken: f64,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct OldQuoteInstructions {
    pub instructions: Vec<OldQuoteInstruction>,
    #[serde(with = "common::serialize_vec_as_strs")]
    pub address_lookup_tables: Vec<Pubkey>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct OldQuoteData {
    pub instructions: Vec<OldQuoteInstructions>,
    pub quotes: Vec<OldQuoteInfo>,
}

fn convert_old_quote(old_quote: &OldQuoteInfo, old_instrs: &OldQuoteInstructions) -> v1::SwapRoute {
    v1::SwapRoute {
        in_amount: old_quote.in_amount,
        out_amount: old_quote.out_amount,
        slippage_bps: old_quote.slippage_bps,
        platform_fee: old_quote.platform_fee.map(Into::into),
        steps: old_quote
            .route_plan
            .iter()
            .cloned()
            .map(Into::into)
            .collect(),
        instructions: old_instrs
            .instructions
            .iter()
            .cloned()
            .map(Into::into)
            .collect(),
        address_lookup_tables: old_instrs
            .address_lookup_tables
            .iter()
            .cloned()
            .map(Into::into)
            .collect(),
        context_slot: Some(old_quote.context_slot),
        time_taken_ns: Some((old_quote.time_taken * 1_000_000_000.0) as u64),
        expires_after_slot: None,
        expires_at_ms: None,
        ..Default::default()
    }
}

fn percent_difference(orig: usize, new: usize) -> f64 {
    let orig_f64 = orig as f64;
    let new_f64 = new as f64;

    (new_f64 / orig_f64) * 100.0
}

#[test]
fn compare_response_sizes() {
    let old_response_ndjson = include_str!("resources/old_quote_response.ndjson");
    let old_responses = old_response_ndjson
        .split('\n')
        .flat_map(|line| {
            if line.is_empty() {
                return None;
            }
            Some(
                serde_json::from_str::<OldQuoteResponse>(line)
                    .expect("old response should be parsable"),
            )
        })
        .collect::<Vec<_>>();

    let old_provider_to_best_quote = old_responses
        .iter()
        .flat_map(|resp| {
            let resp_data =
                serde_json::from_str::<OldQuoteData>(&resp.data).expect("data should decode");
            let best_quote_index =
                resp_data
                    .quotes
                    .iter()
                    .enumerate()
                    .fold(None, |acc, (idx, elem)| {
                        if let Some((max_idx, max_amount_out)) = acc {
                            if elem.out_amount > max_amount_out {
                                Some((idx, elem.out_amount))
                            } else {
                                Some((max_idx, max_amount_out))
                            }
                        } else {
                            Some((idx, elem.out_amount))
                        }
                    });
            if let Some((max_idx, _)) = best_quote_index {
                let old_quote = resp_data.quotes[max_idx].clone();
                let old_instructions = resp_data.instructions[max_idx].clone();
                Some((resp.exchange.clone(), (old_quote, old_instructions)))
            } else {
                None
            }
        })
        .collect::<HashMap<_, _>>();

    let (input_mint, output_mint, swap_mode, amount) =
        if let Some((quote, _)) = old_provider_to_best_quote.values().next() {
            (
                quote.input_mint,
                quote.output_mint,
                quote.swap_mode,
                quote.in_amount,
            )
        } else {
            panic!("no good quotes found!");
        };

    let our_quotes = old_provider_to_best_quote
        .iter()
        .map(|(provider, (old_quote, old_instrs))| {
            (provider.clone(), convert_old_quote(old_quote, old_instrs))
        })
        .collect::<HashMap<String, v1::SwapRoute>>();

    let our_data = v1::ServerMessage::StreamData(v1::StreamData {
        id: 0,
        seq: 0,
        payload: v1::StreamDataPayload::SwapQuotes(v1::SwapQuotes {
            id: "test_quote".into(),
            input_mint: input_mint.into(),
            output_mint: output_mint.into(),
            swap_mode,
            amount,
            quotes: our_quotes,
            metadata: None,
        }),
    });

    let our_data_serialized =
        rmp_serde::encode::to_vec_named(&our_data).expect("should be able to encode our data");

    let old_data_len = old_response_ndjson.len();
    let our_data_len = our_data_serialized.len();

    let pct_diff_from_old = percent_difference(old_data_len, our_data_len);
    println!(
        "Old ndjson size: {}B, v1 StreamData size: {}B, percent of old size: {:.4}%",
        old_data_len, our_data_len, pct_diff_from_old,
    );

    let old_encoded = zstd::bulk::compress(old_response_ndjson.as_bytes(), 3)
        .expect("should be able to compress old data");
    let new_encoded =
        zstd::bulk::compress(&our_data_serialized, 3).expect("should be able to compress new data");

    println!(
        "Zstd serialized (level 3): Old {}B ({:.4}% of orig), New {}B ({:.4}% of orig)",
        old_encoded.len(),
        percent_difference(old_response_ndjson.len(), old_encoded.len()),
        new_encoded.len(),
        percent_difference(our_data_serialized.len(), new_encoded.len()),
    );
    println!(
        "New zstd serialized is {:.4}% of old zstd serialized",
        percent_difference(old_encoded.len(), new_encoded.len())
    );
}

#[test]
fn load_streamdata_msgpack() {
    let stream_data_msgpack = include_bytes!("resources/stream_data.msgpack");
    let data = rmp_serde::from_slice::<v1::ServerMessage>(stream_data_msgpack)
        .expect("should decode to server message");

    let swap_quotes = match data {
        v1::ServerMessage::StreamData(stream_data) => {
            assert_eq!(stream_data.id, 0u32);
            assert_eq!(stream_data.seq, 0u32);
            match stream_data.payload {
                v1::StreamDataPayload::SwapQuotes(quotes) => quotes,
                _ => panic!("Got wrong kind of StreamDataPayload"),
            }
        }
        _ => panic!("Got wrong kind of ServerMessage"),
    };

    assert_eq!(swap_quotes.quotes.len(), 4);
}
