// file: kb-lib/src/decoder/spl/associated_token_account/wire.rs
// version: 1

//! Bounded ATA wire, ordered account and canonical PDA decoding.

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

#[derive(Clone, Debug)]
struct Account {
    position: usize,
    account_index: u64,
    key: std::string::String,
    signer: bool,
    writable: bool,
    source: serde_json::Value,
}

pub(crate) fn entry(bytes: &[u8]) -> std::option::Option<(u8, &'static str, bool)> {
    if bytes.is_empty() {
        return std::option::Option::Some((0, "create", true));
    }
    if bytes.len() != 1 {
        return std::option::Option::None;
    }
    return crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_INSTRUCTION_ENTRIES
        .iter()
        .copied()
        .find(|(tag, _, _)| return *tag == bytes[0]);
}

pub(crate) fn payload(
    input: &crate::CoreInstructionReplayInput,
) -> kb_core::Result<std::vec::Vec<u8>> {
    let value = match input.instruction_payload_json.as_ref() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::invalid_state(
                "ATA instruction payload is not retained",
            ));
        },
    };
    let encoded = match value.get("dataBase64").and_then(serde_json::Value::as_str) {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::invalid_state(
                "ATA payload does not contain dataBase64",
            ));
        },
    };
    if encoded.len() > 92 {
        return std::result::Result::Err(kb_core::Error::invalid_state(
            "ATA base64 payload exceeds the bounded input limit",
        ));
    }
    let bytes = match base64::engine::general_purpose::STANDARD.decode(encoded.as_bytes()) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => {
            return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                "ATA payload is not valid base64: {error}"
            )));
        },
    };
    if bytes.len() > crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_MAX_INSTRUCTION_BYTES {
        return std::result::Result::Err(kb_core::Error::invalid_state(
            "ATA decoded payload exceeds 64 bytes",
        ));
    }
    return std::result::Result::Ok(bytes);
}

pub(crate) fn decode(input: &crate::CoreInstructionReplayInput) -> crate::DecoderExecutionResult {
    let bytes = match crate::spl_associated_token_account_payload(input) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => {
            return failed(std::option::Option::None, "ata_payload_unavailable", error.to_string());
        },
    };
    let instruction = match crate::spl_associated_token_account_entry(bytes.as_slice()) {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            let code = if bytes.len() > 1 {
                "malformed_ata_instruction"
            } else {
                "unknown_ata_instruction_tag"
            };
            return failed(
                std::option::Option::None,
                code,
                format!("unrecognized ATA wire; payloadPrefixHex={}", hex(bytes.as_slice())),
            );
        },
    };
    let accounts = match accounts(input) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => {
            return failed(
                std::option::Option::Some(instruction.1),
                "ata_accounts_malformed",
                error.to_string(),
            );
        },
    };
    let diagnostics = diagnostics(instruction.1, accounts.as_slice());
    let projected_accounts = project_accounts(instruction.1, accounts.as_slice());
    let addresses = addresses(instruction.1, accounts.as_slice());
    let hash = hash(bytes.as_slice());
    let committed = !input.transaction_failed;
    let confidence = crate::DecoderConfidence::ManualExact;
    let event = crate::DecodedProtocolEvent {
        signature: crate::Signature(input.signature.clone()),
        slot: crate::Slot(input.slot),
        instruction_path: crate::InstructionPath(input.instruction_path.clone()),
        program_id: crate::ProgramId(input.program_id.clone()),
        protocol_code: crate::ProtocolCode(
            crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_SURFACE_CODE.to_string(),
        ),
        surface_code: crate::SurfaceCode(
            crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_SURFACE_CODE.to_string(),
        ),
        event_code: crate::EventCode(format!(
            "{}.{}",
            crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_SURFACE_CODE,
            instruction.1
        )),
        event_name: crate::EventName(instruction.1.to_string()),
        event_family: crate::EventFamily::Lifecycle,
        source_kind: if input.instruction_path.contains('/') {
            crate::EventSourceKind::InnerInstruction
        } else {
            crate::EventSourceKind::Instruction
        },
        confidence,
    };
    let observation = crate::DecodedObservation {
        event_key: format!("ata:{}:0", instruction.1),
        event,
        payload_json: serde_json::json!({
            "eventVersion": crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_EVENT_VERSION,
            "programId": input.program_id,
            "instruction": instruction.1,
            "wireTag": instruction.0,
            "wireTagHex": format!("{:02x}", instruction.0),
            "wireEncoding": if bytes.is_empty() { "legacy_empty_create" } else { "borsh_enum_discriminant" },
            "instructionPath": input.instruction_path,
            "instructionLocation": if input.instruction_path.contains('/') { "inner" } else { "outer" },
            "transactionSucceeded": !input.transaction_failed,
            "committed": committed,
            "accounts": projected_accounts,
            "addresses": addresses,
            "semanticDiagnostics": diagnostics,
            "wire": {
                "lengthBytes": bytes.len(),
                "sha256": hash,
                "prefixHex": hex(bytes.as_slice()),
                "complete": true,
            },
            "tokenProgramExtensionsDecoded": false,
            "stateSnapshotReconstructed": false,
        }),
        transaction_failed: input.transaction_failed,
        transaction_error: input.transaction_err_json.clone(),
        observation_committed: committed,
        proof: crate::DecoderProof {
            kind: crate::DecoderProofKind::Manual,
            confidence,
            evidence: std::vec![
                "spl_associated_token_account_interface_2_0_0_and_processor_audit".to_string(),
                format!("wire_sha256:{hash}"),
            ],
        },
    };
    return crate::DecoderExecutionResult {
        status: crate::DecoderOutcomeStatus::Decoded,
        recognized_entry_code: std::option::Option::Some(instruction.1.to_string()),
        observations: std::vec![observation],
        diagnostics: std::vec::Vec::new(),
    };
}

fn accounts(input: &crate::CoreInstructionReplayInput) -> kb_core::Result<std::vec::Vec<Account>> {
    let instruction_accounts = match input.instruction_accounts_json.as_array() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::invalid_state(
                "ATA instruction accounts must be a JSON array",
            ));
        },
    };
    let keys = match input.account_keys_json.as_array() {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::invalid_state(
                "ATA transaction account keys must be a JSON array",
            ));
        },
    };
    if instruction_accounts.len() > crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_MAX_ACCOUNTS {
        return std::result::Result::Err(kb_core::Error::invalid_state(
            "ATA instruction account count exceeds 64",
        ));
    }
    if keys.len() > crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_MAX_TRANSACTION_KEYS {
        return std::result::Result::Err(kb_core::Error::invalid_state(
            "ATA transaction account key count exceeds 4096",
        ));
    }
    let mut output = std::vec::Vec::with_capacity(instruction_accounts.len());
    for (position, value) in instruction_accounts.iter().enumerate() {
        let index = match value.get("accountIndex").and_then(serde_json::Value::as_u64) {
            std::option::Option::Some(value) => value,
            std::option::Option::None => {
                return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                    "ATA instruction account {position} has no accountIndex"
                )));
            },
        };
        let key = match value.get("accountKey").and_then(serde_json::Value::as_str) {
            std::option::Option::Some(value)
                if !value.is_empty()
                    && value.len() <= crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_MAX_ACCOUNT_KEY_BYTES =>
            {
                value
            },
            _ => {
                return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                    "ATA instruction account {position} has no accountKey"
                )));
            },
        };
        let resolved = match keys.iter().find(|candidate| {
            return candidate.get("accountIndex").and_then(serde_json::Value::as_u64)
                == std::option::Option::Some(index);
        }) {
            std::option::Option::Some(value) => value,
            std::option::Option::None => {
                return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                    "ATA account index {index} is not resolved"
                )));
            },
        };
        if resolved.get("accountKey").and_then(serde_json::Value::as_str)
            != std::option::Option::Some(key)
        {
            return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                "ATA account index {index} resolves to a different key"
            )));
        }
        let signer = match resolved.get("signer").and_then(serde_json::Value::as_bool) {
            std::option::Option::Some(value) => value,
            std::option::Option::None => {
                return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                    "ATA account index {index} has no signer flag"
                )));
            },
        };
        let writable = match resolved.get("writable").and_then(serde_json::Value::as_bool) {
            std::option::Option::Some(value) => value,
            std::option::Option::None => {
                return std::result::Result::Err(kb_core::Error::invalid_state(format!(
                    "ATA account index {index} has no writable flag"
                )));
            },
        };
        output.push(Account {
            position,
            account_index: index,
            key: key.to_string(),
            signer,
            writable,
            source: resolved.get("source").cloned().unwrap_or(serde_json::Value::Null),
        });
    }
    return std::result::Result::Ok(output);
}

fn roles(code: &str) -> &'static [(&'static str, bool, bool)] {
    return if code == "recover_nested" {
        &[
            ("nested_associated_token_account", false, true),
            ("nested_mint", false, false),
            ("wallet_nested_mint_associated_token_account", false, true),
            ("owner_associated_token_account", false, false),
            ("owner_mint", false, false),
            ("wallet_owner", true, true),
            ("token_program", false, false),
        ]
    } else {
        &[
            ("funding_account", true, true),
            ("associated_token_account", false, true),
            ("wallet_owner", false, false),
            ("mint", false, false),
            ("system_program", false, false),
            ("token_program", false, false),
        ]
    };
}

fn project_accounts(code: &str, accounts: &[Account]) -> serde_json::Value {
    let roles = roles(code);
    return serde_json::Value::Array(accounts.iter().map(|account| return serde_json::json!({
        "position": account.position,
        "accountIndex": account.account_index,
        "accountKey": account.key,
        "signer": account.signer,
        "writable": account.writable,
        "source": account.source,
        "role": roles.get(account.position).map(|value| return value.0).unwrap_or("extra_account"),
    })).collect());
}

fn diagnostics(code: &str, accounts: &[Account]) -> serde_json::Value {
    let roles = roles(code);
    let mut values = std::vec::Vec::new();
    if accounts.len() != roles.len() {
        values.push(serde_json::json!({"code":"unexpected_account_count","expected":roles.len(),"observed":accounts.len()}));
    }
    for (position, (role, signer, writable)) in roles.iter().enumerate() {
        if let std::option::Option::Some(account) = accounts.get(position) {
            if account.signer != *signer || account.writable != *writable {
                values.push(serde_json::json!({"code":"account_flags_mismatch","position":position,"role":role,"expectedSigner":signer,"observedSigner":account.signer,"expectedWritable":writable,"observedWritable":account.writable}));
            }
        } else {
            values.push(serde_json::json!({"code":"missing_required_account","position":position,"role":role}));
        }
    }
    for left in 0..accounts.len() {
        for right in left.saturating_add(1)..accounts.len() {
            if accounts[left].key == accounts[right].key {
                values.push(serde_json::json!({"code":"duplicate_account_key","positions":[left,right],"accountKey":accounts[left].key}));
            }
        }
    }
    let expected = addresses(code, accounts);
    for field in [
        "associatedTokenAccount",
        "ownerAssociatedTokenAccount",
        "nestedAssociatedTokenAccount",
        "walletNestedMintAssociatedTokenAccount",
    ] {
        if expected
            .get(field)
            .and_then(|value| return value.get("valid"))
            .and_then(serde_json::Value::as_bool)
            == std::option::Option::Some(false)
        {
            values.push(serde_json::json!({"code":"associated_address_mismatch","field":field,"observed":expected[field]["observed"],"expected":expected[field]["expected"]}));
        }
    }
    let token_position = if code == "recover_nested" { 6 } else { 5 };
    if accounts.get(token_position).is_some_and(|account| {
        return account.key != kb_program_ids::SPL_TOKEN_PROGRAM_ID
            && account.key != kb_program_ids::SPL_TOKEN_2022_PROGRAM_ID;
    }) {
        values.push(serde_json::json!({"code":"unsupported_token_program","observed":accounts.get(token_position).map(|value| return value.key.clone())}));
    }
    if code != "recover_nested"
        && accounts
            .get(4)
            .is_some_and(|account| return account.key != "11111111111111111111111111111111")
    {
        values
            .push(serde_json::json!({"code":"system_program_mismatch","observed":accounts[4].key}));
    }
    values.truncate(crate::SPL_ASSOCIATED_TOKEN_ACCOUNT_MAX_DIAGNOSTICS);
    return serde_json::Value::Array(values);
}

fn addresses(code: &str, accounts: &[Account]) -> serde_json::Value {
    let token_position = if code == "recover_nested" { 6 } else { 5 };
    let token = accounts.get(token_position).map(|value| return value.key.as_str());
    if code == "recover_nested" {
        return serde_json::json!({
            "tokenProgram": token_kind(token),
            "ownerAssociatedTokenAccount": derived(accounts.get(5), accounts.get(4), token, accounts.get(3)),
            "nestedAssociatedTokenAccount": derived(accounts.get(3), accounts.get(1), token, accounts.first()),
            "walletNestedMintAssociatedTokenAccount": derived(accounts.get(5), accounts.get(1), token, accounts.get(2)),
            "walletOwner": accounts.get(5).map(|value| return value.key.clone()),
            "ownerMint": accounts.get(4).map(|value| return value.key.clone()),
            "nestedMint": accounts.get(1).map(|value| return value.key.clone()),
        });
    }
    return serde_json::json!({
        "tokenProgram": token_kind(token),
        "associatedTokenAccount": derived(accounts.get(2), accounts.get(3), token, accounts.get(1)),
        "fundingAccount": accounts.first().map(|value| return value.key.clone()),
        "walletOwner": accounts.get(2).map(|value| return value.key.clone()),
        "mint": accounts.get(3).map(|value| return value.key.clone()),
    });
}

fn derived(
    wallet: std::option::Option<&Account>,
    mint: std::option::Option<&Account>,
    token: std::option::Option<&str>,
    observed: std::option::Option<&Account>,
) -> serde_json::Value {
    let expected = match (wallet, mint, token) {
        (
            std::option::Option::Some(wallet),
            std::option::Option::Some(mint),
            std::option::Option::Some(token),
        ) => {
            let wallet_key = wallet.key.parse();
            let mint_key = mint.key.parse();
            let token_key = token.parse();
            match (wallet_key, mint_key, token_key) {
                (std::result::Result::Ok(wallet_key), std::result::Result::Ok(mint_key), std::result::Result::Ok(token_key)) => std::option::Option::Some(spl_associated_token_account_interface::address::get_associated_token_address_with_program_id(&wallet_key, &mint_key, &token_key).to_string()),
                _ => std::option::Option::None,
            }
        },
        _ => std::option::Option::None,
    };
    let observed = observed.map(|value| return value.key.clone());
    let valid = expected.is_some() && expected == observed;
    return serde_json::json!({"expected":expected,"observed":observed,"valid":valid});
}

fn token_kind(token: std::option::Option<&str>) -> serde_json::Value {
    return match token {
        std::option::Option::Some(kb_program_ids::SPL_TOKEN_PROGRAM_ID) => {
            serde_json::json!({"programId":kb_program_ids::SPL_TOKEN_PROGRAM_ID,"kind":"spl_token_classic","supported":true})
        },
        std::option::Option::Some(kb_program_ids::SPL_TOKEN_2022_PROGRAM_ID) => {
            serde_json::json!({"programId":kb_program_ids::SPL_TOKEN_2022_PROGRAM_ID,"kind":"token_2022","supported":true})
        },
        std::option::Option::Some(value) => {
            serde_json::json!({"programId":value,"kind":"unsupported","supported":false})
        },
        std::option::Option::None => {
            serde_json::json!({"programId":serde_json::Value::Null,"kind":"missing","supported":false})
        },
    };
}

fn failed(
    entry: std::option::Option<&str>,
    code: &str,
    message: std::string::String,
) -> crate::DecoderExecutionResult {
    return crate::DecoderExecutionResult {
        status: crate::DecoderOutcomeStatus::Failed,
        recognized_entry_code: entry.map(|value| return value.to_string()),
        observations: std::vec::Vec::new(),
        diagnostics: std::vec![crate::DecoderDiagnostic {
            code: code.to_string(),
            message,
            retriable: false
        }],
    };
}

fn hash(bytes: &[u8]) -> std::string::String {
    let digest = sha2::Sha256::digest(bytes);
    let mut output = std::string::String::with_capacity(64);
    for byte in digest {
        output.push_str(format!("{byte:02x}").as_str());
    }
    return output;
}

fn hex(bytes: &[u8]) -> std::string::String {
    return bytes.iter().take(16).map(|value| return format!("{value:02x}")).collect();
}