// file: crates/ksp-raw-transaction-lib/src/wire.rs
// version: 3

use base64::Engine; // rust-rules: trait-import

const MAX_LEGACY_SHORT_VECTOR_VALUE: usize = u16::MAX as usize;
const MAX_V1_ADDRESS_COUNT: usize = 64;
const MAX_V1_INSTRUCTION_COUNT: usize = 64;
const MAX_V1_SIGNATURE_COUNT: usize = 12;
const MAX_V1_TRANSACTION_BYTES: usize = 4_096;

/// Source-neutral Solana transaction message version used by the exact wire serializer.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RawSolanaMessageVersion {
    /// Legacy Solana message without a version prefix.
    Legacy,
    /// Versioned transaction message v0, prefixed with `0x80`.
    V0,
    /// Transaction V1 / SIMD-0385 message, prefixed with `0x81`.
    V1,
}

/// Exact three-byte Solana message header.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct RawSolanaMessageHeader {
    num_required_signatures: u8,
    num_readonly_signed_accounts: u8,
    num_readonly_unsigned_accounts: u8,
}

impl crate::RawSolanaMessageHeader {
    /// Creates an exact Solana message header.
    #[must_use]
    pub const fn new(num_required_signatures: u8, num_readonly_signed_accounts: u8, num_readonly_unsigned_accounts: u8) -> Self {
        return Self { num_required_signatures, num_readonly_signed_accounts, num_readonly_unsigned_accounts };
    }

    /// Returns the required signature count.
    #[must_use]
    pub const fn num_required_signatures(self) -> u8 {
        return self.num_required_signatures;
    }

    /// Returns the readonly signed-account count.
    #[must_use]
    pub const fn num_readonly_signed_accounts(self) -> u8 {
        return self.num_readonly_signed_accounts;
    }

    /// Returns the readonly unsigned-account count.
    #[must_use]
    pub const fn num_readonly_unsigned_accounts(self) -> u8 {
        return self.num_readonly_unsigned_accounts;
    }
}

/// One source-neutral compiled Solana instruction.
#[derive(Clone, Eq, PartialEq)]
pub struct RawSolanaCompiledInstruction {
    program_id_index: u8,
    accounts: std::vec::Vec<u8>,
    data: std::vec::Vec<u8>,
}

impl crate::RawSolanaCompiledInstruction {
    /// Creates one compiled instruction while preserving exact ordered account indexes and data bytes.
    #[must_use]
    pub fn new(program_id_index: u8, accounts: std::vec::Vec<u8>, data: std::vec::Vec<u8>) -> Self {
        return Self { program_id_index, accounts, data };
    }

    /// Returns the program account index.
    #[must_use]
    pub const fn program_id_index(&self) -> u8 {
        return self.program_id_index;
    }

    /// Returns ordered account indexes.
    #[must_use]
    pub fn accounts(&self) -> &[u8] {
        return self.accounts.as_slice();
    }

    /// Returns exact instruction data bytes.
    #[must_use]
    pub fn data(&self) -> &[u8] {
        return self.data.as_slice();
    }
}

impl std::fmt::Debug for crate::RawSolanaCompiledInstruction {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter
            .debug_struct("RawSolanaCompiledInstruction")
            .field("program_id_index", &self.program_id_index)
            .field("account_index_count", &self.accounts.len())
            .field("data_length", &self.data.len())
            .finish();
    }
}

/// One source-neutral v0 address-table lookup.
#[derive(Clone, Eq, PartialEq)]
pub struct RawSolanaAddressTableLookup {
    account_key: [u8; 32],
    writable_indexes: std::vec::Vec<u8>,
    readonly_indexes: std::vec::Vec<u8>,
}

impl crate::RawSolanaAddressTableLookup {
    /// Creates one exact v0 address-table lookup.
    #[must_use]
    pub fn new(account_key: [u8; 32], writable_indexes: std::vec::Vec<u8>, readonly_indexes: std::vec::Vec<u8>) -> Self {
        return Self { account_key, writable_indexes, readonly_indexes };
    }

    /// Returns the lookup-table account key bytes.
    #[must_use]
    pub const fn account_key(&self) -> &[u8; 32] {
        return &self.account_key;
    }

    /// Returns ordered writable lookup indexes.
    #[must_use]
    pub fn writable_indexes(&self) -> &[u8] {
        return self.writable_indexes.as_slice();
    }

    /// Returns ordered readonly lookup indexes.
    #[must_use]
    pub fn readonly_indexes(&self) -> &[u8] {
        return self.readonly_indexes.as_slice();
    }
}

impl std::fmt::Debug for crate::RawSolanaAddressTableLookup {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter
            .debug_struct("RawSolanaAddressTableLookup")
            .field("writable_index_count", &self.writable_indexes.len())
            .field("readonly_index_count", &self.readonly_indexes.len())
            .finish_non_exhaustive();
    }
}

/// Optional inline budget configuration carried by Solana Transaction V1.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct RawSolanaTransactionConfig {
    priority_fee: std::option::Option<u64>,
    compute_unit_limit: std::option::Option<u32>,
    loaded_accounts_data_size_limit: std::option::Option<u32>,
    heap_size: std::option::Option<u32>,
}

impl crate::RawSolanaTransactionConfig {
    /// Creates an explicit Transaction V1 configuration, including the meaningful all-`None` case.
    #[must_use]
    pub const fn new(
        priority_fee: std::option::Option<u64>,
        compute_unit_limit: std::option::Option<u32>,
        loaded_accounts_data_size_limit: std::option::Option<u32>,
        heap_size: std::option::Option<u32>,
    ) -> Self {
        return Self { priority_fee, compute_unit_limit, loaded_accounts_data_size_limit, heap_size };
    }

    /// Returns the optional priority fee.
    #[must_use]
    pub const fn priority_fee(self) -> std::option::Option<u64> {
        return self.priority_fee;
    }

    /// Returns the optional compute-unit limit.
    #[must_use]
    pub const fn compute_unit_limit(self) -> std::option::Option<u32> {
        return self.compute_unit_limit;
    }

    /// Returns the optional loaded-account-data-size limit.
    #[must_use]
    pub const fn loaded_accounts_data_size_limit(self) -> std::option::Option<u32> {
        return self.loaded_accounts_data_size_limit;
    }

    /// Returns the optional heap-size override.
    #[must_use]
    pub const fn heap_size(self) -> std::option::Option<u32> {
        return self.heap_size;
    }
}

/// Complete source-neutral Solana transaction message required for exact wire serialization.
#[derive(Clone, Eq, PartialEq)]
pub struct RawSolanaTransactionMessage {
    version: crate::RawSolanaMessageVersion,
    header: crate::RawSolanaMessageHeader,
    account_keys: std::vec::Vec<[u8; 32]>,
    recent_blockhash: [u8; 32],
    instructions: std::vec::Vec<crate::RawSolanaCompiledInstruction>,
    address_table_lookups: std::vec::Vec<crate::RawSolanaAddressTableLookup>,
    config: std::option::Option<crate::RawSolanaTransactionConfig>,
}

impl crate::RawSolanaTransactionMessage {
    /// Creates one exact source-neutral message.
    #[must_use]
    pub fn new(
        version: crate::RawSolanaMessageVersion,
        header: crate::RawSolanaMessageHeader,
        account_keys: std::vec::Vec<[u8; 32]>,
        recent_blockhash: [u8; 32],
        instructions: std::vec::Vec<crate::RawSolanaCompiledInstruction>,
        address_table_lookups: std::vec::Vec<crate::RawSolanaAddressTableLookup>,
        config: std::option::Option<crate::RawSolanaTransactionConfig>,
    ) -> Self {
        return Self { version, header, account_keys, recent_blockhash, instructions, address_table_lookups, config };
    }

    /// Returns the message version.
    #[must_use]
    pub const fn version(&self) -> crate::RawSolanaMessageVersion {
        return self.version;
    }

    /// Returns the message header.
    #[must_use]
    pub const fn header(&self) -> crate::RawSolanaMessageHeader {
        return self.header;
    }

    /// Returns ordered static account key bytes.
    #[must_use]
    pub fn account_keys(&self) -> &[[u8; 32]] {
        return self.account_keys.as_slice();
    }

    /// Returns the recent blockhash bytes.
    #[must_use]
    pub const fn recent_blockhash(&self) -> &[u8; 32] {
        return &self.recent_blockhash;
    }

    /// Returns ordered compiled instructions.
    #[must_use]
    pub fn instructions(&self) -> &[crate::RawSolanaCompiledInstruction] {
        return self.instructions.as_slice();
    }

    /// Returns ordered v0 address-table lookups.
    #[must_use]
    pub fn address_table_lookups(&self) -> &[crate::RawSolanaAddressTableLookup] {
        return self.address_table_lookups.as_slice();
    }

    /// Returns optional V1 inline configuration.
    #[must_use]
    pub const fn config(&self) -> std::option::Option<crate::RawSolanaTransactionConfig> {
        return self.config;
    }
}

impl std::fmt::Debug for crate::RawSolanaTransactionMessage {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter
            .debug_struct("RawSolanaTransactionMessage")
            .field("version", &self.version)
            .field("header", &self.header)
            .field("account_key_count", &self.account_keys.len())
            .field("instruction_count", &self.instructions.len())
            .field("address_table_lookup_count", &self.address_table_lookups.len())
            .field("has_config", &self.config.is_some())
            .finish();
    }
}

/// Complete source-neutral Solana transaction wire material.
#[derive(Clone, Eq, PartialEq)]
pub struct RawSolanaTransactionWire {
    signatures: std::vec::Vec<[u8; 64]>,
    message: crate::RawSolanaTransactionMessage,
}

impl crate::RawSolanaTransactionWire {
    /// Creates one complete transaction wire.
    #[must_use]
    pub fn new(signatures: std::vec::Vec<[u8; 64]>, message: crate::RawSolanaTransactionMessage) -> Self {
        return Self { signatures, message };
    }

    /// Returns ordered signature bytes.
    #[must_use]
    pub fn signatures(&self) -> &[[u8; 64]] {
        return self.signatures.as_slice();
    }

    /// Returns the exact source-neutral message.
    #[must_use]
    pub const fn message(&self) -> &crate::RawSolanaTransactionMessage {
        return &self.message;
    }
}

impl std::fmt::Debug for crate::RawSolanaTransactionWire {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter.debug_struct("RawSolanaTransactionWire").field("signature_count", &self.signatures.len()).field("message", &self.message).finish();
    }
}

/// Serializes one source-neutral Solana transaction to exact canonical wire bytes.
pub fn serialize_solana_transaction_wire(transaction: &crate::RawSolanaTransactionWire) -> ksp_core_lib::Result<std::vec::Vec<u8>> {
    if let std::result::Result::Err(error) = validate_transaction(transaction) {
        return std::result::Result::Err(error);
    }
    let mut output = std::vec::Vec::new();
    match transaction.message().version() {
        crate::RawSolanaMessageVersion::Legacy | crate::RawSolanaMessageVersion::V0 => {
            if let std::result::Result::Err(error) = encode_short_vec(transaction.signatures().len(), &mut output) {
                return std::result::Result::Err(error);
            }
            for signature in transaction.signatures() {
                output.extend_from_slice(signature);
            }
            if transaction.message().version() == crate::RawSolanaMessageVersion::V0 {
                output.push(0x80);
            }
            if let std::result::Result::Err(error) = encode_legacy_or_v0_message(transaction.message(), &mut output) {
                return std::result::Result::Err(error);
            }
        },
        crate::RawSolanaMessageVersion::V1 => {
            if let std::result::Result::Err(error) = encode_v1_message(transaction, &mut output) {
                return std::result::Result::Err(error);
            }
        },
    }
    if output.len() > ksp_store_api::MAX_RAW_PAYLOAD_BYTES {
        return std::result::Result::Err(crate::material_error("wire.bytes").with_context("actual_len", output.len().to_string()));
    }
    if transaction.message().version() == crate::RawSolanaMessageVersion::V1 && output.len() > MAX_V1_TRANSACTION_BYTES {
        return std::result::Result::Err(crate::material_error("wire.bytes").with_context("actual_len", output.len().to_string()));
    }
    return std::result::Result::Ok(output);
}

/// Serializes one source-neutral Solana transaction and returns canonical padded standard Base64.
pub fn serialize_solana_transaction_wire_base64(transaction: &crate::RawSolanaTransactionWire) -> ksp_core_lib::Result<std::string::String> {
    let bytes = match crate::serialize_solana_transaction_wire(transaction) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    return std::result::Result::Ok(base64::engine::general_purpose::STANDARD.encode(bytes));
}

fn validate_transaction(transaction: &crate::RawSolanaTransactionWire) -> ksp_core_lib::Result<()> {
    let message = transaction.message();
    let required_signatures = usize::from(message.header().num_required_signatures());
    if required_signatures == 0 || transaction.signatures().len() != required_signatures {
        return std::result::Result::Err(crate::material_error("wire"));
    }
    if usize::from(message.header().num_readonly_signed_accounts()) >= required_signatures {
        return std::result::Result::Err(crate::material_error("wire"));
    }
    if message.account_keys().len() < required_signatures + usize::from(message.header().num_readonly_unsigned_accounts()) {
        return std::result::Result::Err(crate::material_error("wire"));
    }
    let loaded_count = message.address_table_lookups().iter().fold(0_usize, |count, lookup| {
        return count.saturating_add(lookup.writable_indexes().len()).saturating_add(lookup.readonly_indexes().len());
    });
    let account_count = message.account_keys().len().saturating_add(loaded_count);
    for instruction in message.instructions() {
        if usize::from(instruction.program_id_index()) >= account_count {
            return std::result::Result::Err(crate::material_error("wire"));
        }
        for account_index in instruction.accounts() {
            if usize::from(*account_index) >= account_count {
                return std::result::Result::Err(crate::material_error("wire"));
            }
        }
    }
    match message.version() {
        crate::RawSolanaMessageVersion::Legacy => {
            if !message.address_table_lookups().is_empty() || message.config().is_some() {
                return std::result::Result::Err(crate::material_error("wire"));
            }
        },
        crate::RawSolanaMessageVersion::V0 => {
            if message.config().is_some() {
                return std::result::Result::Err(crate::material_error("wire"));
            }
        },
        crate::RawSolanaMessageVersion::V1 => {
            if message.config().is_none() || !message.address_table_lookups().is_empty() {
                return std::result::Result::Err(crate::material_error("wire"));
            }
            if transaction.signatures().len() > MAX_V1_SIGNATURE_COUNT
                || message.account_keys().len() > MAX_V1_ADDRESS_COUNT
                || message.instructions().len() > MAX_V1_INSTRUCTION_COUNT
            {
                return std::result::Result::Err(crate::material_error("wire"));
            }
            for (index, account) in message.account_keys().iter().enumerate() {
                if message.account_keys()[..index].contains(account) {
                    return std::result::Result::Err(crate::material_error("wire"));
                }
            }
            for instruction in message.instructions() {
                if instruction.accounts().len() > usize::from(u8::MAX) || instruction.data().len() > usize::from(u16::MAX) {
                    return std::result::Result::Err(crate::material_error("wire"));
                }
            }
            if let std::option::Option::Some(heap_size) = message.config().and_then(crate::RawSolanaTransactionConfig::heap_size)
                && (!(32_768..=262_144).contains(&heap_size) || heap_size % 1_024 != 0)
            {
                return std::result::Result::Err(crate::material_error("wire"));
            }
        },
    }
    return std::result::Result::Ok(());
}

fn encode_legacy_or_v0_message(message: &crate::RawSolanaTransactionMessage, output: &mut std::vec::Vec<u8>) -> ksp_core_lib::Result<()> {
    encode_header(message.header(), output);
    if let std::result::Result::Err(error) = encode_short_vec(message.account_keys().len(), output) {
        return std::result::Result::Err(error);
    }
    for account_key in message.account_keys() {
        output.extend_from_slice(account_key);
    }
    output.extend_from_slice(message.recent_blockhash());
    if let std::result::Result::Err(error) = encode_short_vec(message.instructions().len(), output) {
        return std::result::Result::Err(error);
    }
    for instruction in message.instructions() {
        output.push(instruction.program_id_index());
        if let std::result::Result::Err(error) = encode_short_vec(instruction.accounts().len(), output) {
            return std::result::Result::Err(error);
        }
        output.extend_from_slice(instruction.accounts());
        if let std::result::Result::Err(error) = encode_short_vec(instruction.data().len(), output) {
            return std::result::Result::Err(error);
        }
        output.extend_from_slice(instruction.data());
    }
    if message.version() == crate::RawSolanaMessageVersion::V0 {
        if let std::result::Result::Err(error) = encode_short_vec(message.address_table_lookups().len(), output) {
            return std::result::Result::Err(error);
        }
        for lookup in message.address_table_lookups() {
            output.extend_from_slice(lookup.account_key());
            if let std::result::Result::Err(error) = encode_short_vec(lookup.writable_indexes().len(), output) {
                return std::result::Result::Err(error);
            }
            output.extend_from_slice(lookup.writable_indexes());
            if let std::result::Result::Err(error) = encode_short_vec(lookup.readonly_indexes().len(), output) {
                return std::result::Result::Err(error);
            }
            output.extend_from_slice(lookup.readonly_indexes());
        }
    }
    return std::result::Result::Ok(());
}

fn encode_v1_message(transaction: &crate::RawSolanaTransactionWire, output: &mut std::vec::Vec<u8>) -> ksp_core_lib::Result<()> {
    let message = transaction.message();
    let config = match message.config() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => return std::result::Result::Err(crate::material_error("wire")),
    };
    output.push(0x81);
    encode_header(message.header(), output);
    let mut mask = 0_u32;
    if config.priority_fee().is_some() {
        mask |= 0b11;
    }
    if config.compute_unit_limit().is_some() {
        mask |= 1_u32 << 2;
    }
    if config.loaded_accounts_data_size_limit().is_some() {
        mask |= 1_u32 << 3;
    }
    if config.heap_size().is_some() {
        mask |= 1_u32 << 4;
    }
    output.extend_from_slice(&mask.to_le_bytes());
    output.extend_from_slice(message.recent_blockhash());
    output.push(message.instructions().len() as u8);
    output.push(message.account_keys().len() as u8);
    for account_key in message.account_keys() {
        output.extend_from_slice(account_key);
    }
    if let std::option::Option::Some(value) = config.priority_fee() {
        output.extend_from_slice(&value.to_le_bytes());
    }
    if let std::option::Option::Some(value) = config.compute_unit_limit() {
        output.extend_from_slice(&value.to_le_bytes());
    }
    if let std::option::Option::Some(value) = config.loaded_accounts_data_size_limit() {
        output.extend_from_slice(&value.to_le_bytes());
    }
    if let std::option::Option::Some(value) = config.heap_size() {
        output.extend_from_slice(&value.to_le_bytes());
    }
    for instruction in message.instructions() {
        output.push(instruction.program_id_index());
        output.push(instruction.accounts().len() as u8);
        output.extend_from_slice(&(instruction.data().len() as u16).to_le_bytes());
    }
    for instruction in message.instructions() {
        output.extend_from_slice(instruction.accounts());
        output.extend_from_slice(instruction.data());
    }
    for signature in transaction.signatures() {
        output.extend_from_slice(signature);
    }
    return std::result::Result::Ok(());
}

fn encode_header(header: crate::RawSolanaMessageHeader, output: &mut std::vec::Vec<u8>) {
    output.push(header.num_required_signatures());
    output.push(header.num_readonly_signed_accounts());
    output.push(header.num_readonly_unsigned_accounts());
}

fn encode_short_vec(value: usize, output: &mut std::vec::Vec<u8>) -> ksp_core_lib::Result<()> {
    if value > MAX_LEGACY_SHORT_VECTOR_VALUE {
        return std::result::Result::Err(crate::material_error("wire"));
    }
    let mut remaining = value;
    loop {
        let mut byte = (remaining & 0x7f) as u8;
        remaining >>= 7;
        if remaining != 0 {
            byte |= 0x80;
        }
        output.push(byte);
        if remaining == 0 {
            break;
        }
    }
    return std::result::Result::Ok(());
}

#[cfg(test)]
#[path = "../unit_tests/wire.rs"]
mod tests;