// file: kb-pipeline/src/solana_token_2022_preflight.rs
// version: 3

//! Bounded Token-2022 stateful preflight orchestration.

/// Maximum distinct Token-2022 accounts accepted by one preflight inspection.
pub const MAX_TOKEN_2022_PREFLIGHT_ACCOUNTS: usize = 16;
/// Maximum aggregate account-data budget accepted by one preflight inspection.
pub const MAX_TOKEN_2022_PREFLIGHT_TOTAL_BYTES: usize = 262_144;

/// One exact Token-2022 account requirement for a stateful preflight.
#[derive(Clone, Debug, PartialEq)]
pub struct Token2022PreflightRequirement {
    /// Stable semantic role used in diagnostics.
    pub role: std::string::String,
    /// Canonical account address.
    pub account: kb_lib::MdPubkey,
    /// Expected Token-2022 state category.
    pub kind: kb_lib::DcToken2022StateKind,
    /// Maximum complete account bytes accepted for this requirement.
    pub max_data_bytes: usize,
    /// Optional expected mint for a Token Account.
    pub expected_mint: std::option::Option<kb_lib::MdPubkey>,
    /// Optional expected owner for a Token Account.
    pub expected_owner: std::option::Option<kb_lib::MdPubkey>,
    /// Optional exact decimals expected for a Mint.
    pub expected_decimals: std::option::Option<u8>,
    /// Published extension names required on the account.
    pub required_extensions: std::vec::Vec<std::string::String>,
    /// Optional external identities needed by cross-account extensions.
    pub context: crate::Token2022StatefulContext,
}

/// One bounded Token-2022 preflight request.
#[derive(Clone, Debug, PartialEq)]
pub struct Token2022PreflightRequest {
    /// Endpoint role used by all RPC reads.
    pub query_role: std::string::String,
    /// Optional minimum context slot shared by all reads.
    pub min_context_slot: std::option::Option<u64>,
    /// Maximum distinct accounts accepted after deduplication.
    pub max_accounts: usize,
    /// Maximum aggregate requested account bytes.
    pub max_total_data_bytes: usize,
    /// Ordered account requirements.
    pub requirements: std::vec::Vec<Token2022PreflightRequirement>,
    /// Optional ElGamal registry read required by the operation.
    pub elgamal_registry: std::option::Option<crate::ElGamalRegistryStatefulReadRequest>,
}

/// One validated account result in a Token-2022 preflight report.
#[derive(Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct Token2022PreflightAccountReport {
    /// Stable semantic role.
    pub role: std::string::String,
    /// Canonical account address.
    pub account: kb_lib::MdPubkey,
    /// State category observed after parsing.
    pub state_kind: std::string::String,
    /// Context slot returned by the endpoint.
    pub context_slot: u64,
    /// Ordered extension names observed on the account.
    pub extension_names: std::vec::Vec<std::string::String>,
    /// Ordered successful semantic checks.
    pub checks: std::vec::Vec<std::string::String>,
}

/// Complete bounded Token-2022 stateful preflight report.
#[derive(Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct Token2022PreflightReport {
    /// Commitment used for all account reads.
    pub commitment: std::string::String,
    /// Highest context slot observed across all reads.
    pub context_slot: u64,
    /// Aggregate requested account-data budget.
    pub requested_data_bytes: usize,
    /// Ordered distinct Token-2022 account reports.
    pub accounts: std::vec::Vec<Token2022PreflightAccountReport>,
    /// Whether an ElGamal registry was required and validated.
    pub elgamal_registry_validated: bool,
}

/// Inspects all bounded Token-2022 state required before simulation.
pub async fn inspect_token_2022_preflight(
    pool: &kb_onchain_transport::HttpEndpointPool,
    request: &crate::Token2022PreflightRequest,
) -> kb_core::Result<crate::Token2022PreflightReport> {
    let requirements = match validate_and_deduplicate_requirements(request) {
        std::result::Result::Ok(value) => value,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    let requested_data_bytes = requirements
        .iter()
        .fold(0usize, |total, requirement| {
            return total.saturating_add(requirement.max_data_bytes);
        })
        .saturating_add(if request.elgamal_registry.is_some() {
            crate::ELGAMAL_REGISTRY_STATEFUL_ACCOUNT_BYTES
        } else {
            0
        });
    let mut context_slot = request.min_context_slot.unwrap_or(0);
    let mut accounts = std::vec::Vec::with_capacity(requirements.len());
    for requirement in requirements {
        let read_request = crate::Token2022StatefulReadRequest {
            query_role: request.query_role.clone(),
            account: requirement.account.clone(),
            kind: requirement.kind,
            min_context_slot: request.min_context_slot,
            max_data_bytes: requirement.max_data_bytes,
            context: requirement.context.clone(),
        };
        let read_result = match crate::read_token_2022_stateful_snapshot(pool, &read_request).await
        {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        let account_report = match validate_snapshot_requirement(&requirement, &read_result) {
            std::result::Result::Ok(value) => value,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        context_slot = context_slot.max(account_report.context_slot);
        accounts.push(account_report);
    }
    let mut elgamal_registry_validated = false;
    if let std::option::Option::Some(registry_request) = request.elgamal_registry.as_ref() {
        if registry_request.query_role != request.query_role {
            return std::result::Result::Err(kb_core::Error::config(
                "Token-2022 preflight ElGamal registry query role must match the shared query role",
            ));
        }
        let registry_result =
            match crate::read_elgamal_registry_stateful_snapshot(pool, registry_request).await {
                std::result::Result::Ok(value) => value,
                std::result::Result::Err(error) => return std::result::Result::Err(error),
            };
        context_slot = context_slot.max(registry_result.context_slot);
        elgamal_registry_validated = true;
    }
    return std::result::Result::Ok(crate::Token2022PreflightReport {
        commitment: "confirmed".to_string(),
        context_slot,
        requested_data_bytes,
        accounts,
        elgamal_registry_validated,
    });
}

fn validate_and_deduplicate_requirements(
    request: &crate::Token2022PreflightRequest,
) -> kb_core::Result<std::vec::Vec<crate::Token2022PreflightRequirement>> {
    if request.query_role.trim().is_empty() {
        return std::result::Result::Err(kb_core::Error::config(
            "Token-2022 preflight query_role must not be empty",
        ));
    }
    if request.max_accounts == 0 || request.max_accounts > crate::MAX_TOKEN_2022_PREFLIGHT_ACCOUNTS
    {
        return std::result::Result::Err(kb_core::Error::config(format!(
            "Token-2022 preflight max_accounts must be between 1 and {}",
            crate::MAX_TOKEN_2022_PREFLIGHT_ACCOUNTS
        )));
    }
    if request.max_total_data_bytes == 0
        || request.max_total_data_bytes > crate::MAX_TOKEN_2022_PREFLIGHT_TOTAL_BYTES
    {
        return std::result::Result::Err(kb_core::Error::config(format!(
            "Token-2022 preflight max_total_data_bytes must be between 1 and {}",
            crate::MAX_TOKEN_2022_PREFLIGHT_TOTAL_BYTES
        )));
    }
    if request.requirements.is_empty() {
        return std::result::Result::Err(kb_core::Error::config(
            "Token-2022 preflight requires at least one account",
        ));
    }
    let mut indexes = std::collections::BTreeMap::<std::string::String, usize>::new();
    let mut unique = std::vec::Vec::<crate::Token2022PreflightRequirement>::new();
    for requirement in &request.requirements {
        if requirement.role.trim().is_empty() {
            return std::result::Result::Err(kb_core::Error::config(
                "Token-2022 preflight requirement role must not be empty",
            ));
        }
        if requirement.max_data_bytes == 0
            || requirement.max_data_bytes > crate::MAX_TOKEN_2022_STATEFUL_ACCOUNT_BYTES
        {
            return std::result::Result::Err(kb_core::Error::config(format!(
                "Token-2022 preflight account {} has invalid max_data_bytes {}",
                requirement.account.0, requirement.max_data_bytes
            )));
        }
        if let std::option::Option::Some(index) = indexes.get(requirement.account.0.as_str()) {
            if &unique[*index] != requirement {
                return std::result::Result::Err(kb_core::Error::new(
                    "token_2022_preflight_conflicting_duplicate",
                    format!(
                        "Token-2022 preflight account {} has conflicting requirements",
                        requirement.account.0
                    ),
                ));
            }
            continue;
        }
        indexes.insert(requirement.account.0.clone(), unique.len());
        unique.push(requirement.clone());
    }
    let total_accounts =
        unique
            .len()
            .saturating_add(if request.elgamal_registry.is_some() { 1 } else { 0 });
    if total_accounts > request.max_accounts {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_preflight_account_limit_exceeded",
            format!(
                "Token-2022 preflight requires {total_accounts} distinct accounts above limit {}",
                request.max_accounts
            ),
        ));
    }
    let requested_data_bytes = unique
        .iter()
        .fold(0usize, |total, requirement| {
            return total.saturating_add(requirement.max_data_bytes);
        })
        .saturating_add(if request.elgamal_registry.is_some() {
            crate::ELGAMAL_REGISTRY_STATEFUL_ACCOUNT_BYTES
        } else {
            0
        });
    if requested_data_bytes > request.max_total_data_bytes {
        return std::result::Result::Err(kb_core::Error::new(
            "token_2022_preflight_data_budget_exceeded",
            format!(
                "Token-2022 preflight requests {requested_data_bytes} bytes above aggregate limit {}",
                request.max_total_data_bytes
            ),
        ));
    }
    return std::result::Result::Ok(unique);
}

fn validate_snapshot_requirement(
    requirement: &crate::Token2022PreflightRequirement,
    result: &crate::Token2022StatefulReadResult,
) -> kb_core::Result<crate::Token2022PreflightAccountReport> {
    let token_output = result
        .snapshot
        .outputs
        .iter()
        .find(|output| return output.payload_json["domain"] == "spl_token_2022_account_state");
    let token_output = match token_output {
        std::option::Option::Some(value) => value,
        std::option::Option::None => {
            return std::result::Result::Err(kb_core::Error::invalid_state(
                "Token-2022 preflight snapshot is missing its token account owner projection",
            ));
        },
    };
    let base_fields = &token_output.payload_json["baseFields"];
    let mut checks = std::vec![
        "owner_program_id".to_string(),
        "complete_account_data".to_string(),
        "state_kind".to_string()
    ];
    if let std::option::Option::Some(expected_mint) = requirement.expected_mint.as_ref() {
        if base_fields["mint"].as_str() != std::option::Option::Some(expected_mint.0.as_str()) {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_preflight_mint_mismatch",
                format!(
                    "Token-2022 account {} mint does not match expected {}",
                    requirement.account.0, expected_mint.0
                ),
            ));
        }
        checks.push("mint_identity".to_string());
    }
    if let std::option::Option::Some(expected_owner) = requirement.expected_owner.as_ref() {
        if base_fields["owner"].as_str() != std::option::Option::Some(expected_owner.0.as_str()) {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_preflight_owner_mismatch",
                format!(
                    "Token-2022 account {} owner does not match expected {}",
                    requirement.account.0, expected_owner.0
                ),
            ));
        }
        checks.push("token_account_owner".to_string());
    }
    if let std::option::Option::Some(expected_decimals) = requirement.expected_decimals {
        if base_fields["decimals"].as_u64() != std::option::Option::Some(expected_decimals as u64) {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_preflight_decimals_mismatch",
                format!(
                    "Token-2022 mint {} decimals do not match expected {expected_decimals}",
                    requirement.account.0
                ),
            ));
        }
        checks.push("mint_decimals".to_string());
    }
    for required_extension in &requirement.required_extensions {
        if !result
            .snapshot
            .extension_names
            .iter()
            .any(|value| return value == required_extension)
        {
            return std::result::Result::Err(kb_core::Error::new(
                "token_2022_preflight_extension_missing",
                format!(
                    "Token-2022 account {} is missing required extension {required_extension}",
                    requirement.account.0
                ),
            ));
        }
    }
    if !requirement.required_extensions.is_empty() {
        checks.push("required_extensions".to_string());
    }
    return std::result::Result::Ok(crate::Token2022PreflightAccountReport {
        role: requirement.role.clone(),
        account: requirement.account.clone(),
        state_kind: result.snapshot.state_kind.clone(),
        context_slot: result.context_slot,
        extension_names: result.snapshot.extension_names.clone(),
        checks,
    });
}

#[cfg(test)]
mod tests {
    fn pubkey(byte: u8) -> kb_lib::MdPubkey {
        return kb_lib::MdPubkey(bs58::encode([byte; 32]).into_string());
    }

    fn requirement(account: kb_lib::MdPubkey) -> crate::Token2022PreflightRequirement {
        return crate::Token2022PreflightRequirement {
            role: "source".to_string(),
            account,
            kind: kb_lib::DcToken2022StateKind::Account,
            max_data_bytes: 512,
            expected_mint: std::option::Option::None,
            expected_owner: std::option::Option::None,
            expected_decimals: std::option::Option::None,
            required_extensions: std::vec::Vec::new(),
            context: crate::Token2022StatefulContext::default(),
        };
    }

    #[test]
    fn exact_duplicate_requirements_are_deduplicated_and_bounded() {
        let item = requirement(pubkey(1));
        let request = crate::Token2022PreflightRequest {
            query_role: "query".to_string(),
            min_context_slot: std::option::Option::Some(7),
            max_accounts: 1,
            max_total_data_bytes: 512,
            requirements: std::vec![item.clone(), item],
            elgamal_registry: std::option::Option::None,
        };
        let result = super::validate_and_deduplicate_requirements(&request);
        assert_eq!(result.as_ref().map(std::vec::Vec::len), std::result::Result::Ok(1));
    }

    #[test]
    fn conflicting_duplicates_and_aggregate_budget_fail_closed() {
        let first = requirement(pubkey(2));
        let mut conflicting = first.clone();
        conflicting.role = "destination".to_string();
        let conflict_request = crate::Token2022PreflightRequest {
            query_role: "query".to_string(),
            min_context_slot: std::option::Option::None,
            max_accounts: 2,
            max_total_data_bytes: 1024,
            requirements: std::vec![first, conflicting],
            elgamal_registry: std::option::Option::None,
        };
        assert!(super::validate_and_deduplicate_requirements(&conflict_request).is_err());
        let budget_request = crate::Token2022PreflightRequest {
            query_role: "query".to_string(),
            min_context_slot: std::option::Option::None,
            max_accounts: 2,
            max_total_data_bytes: 700,
            requirements: std::vec![requirement(pubkey(3)), requirement(pubkey(4))],
            elgamal_registry: std::option::Option::None,
        };
        assert!(super::validate_and_deduplicate_requirements(&budget_request).is_err());
    }

    #[test]
    fn requirement_validation_checks_mint_owner_decimals_and_extensions() {
        let account = pubkey(5);
        let mint = pubkey(6);
        let owner = pubkey(7);
        let requirement = crate::Token2022PreflightRequirement {
            role: "source".to_string(),
            account: account.clone(),
            kind: kb_lib::DcToken2022StateKind::Account,
            max_data_bytes: 512,
            expected_mint: std::option::Option::Some(mint.clone()),
            expected_owner: std::option::Option::Some(owner.clone()),
            expected_decimals: std::option::Option::None,
            required_extensions: std::vec!["memo_transfer".to_string()],
            context: crate::Token2022StatefulContext::default(),
        };
        let result = crate::Token2022StatefulReadResult {
            commitment: "confirmed".to_string(),
            context_slot: 42,
            snapshot: crate::Token2022StatefulSnapshotBundle {
                account_key: account.0.clone(),
                slot: 42,
                state_kind: "account".to_string(),
                extension_names: std::vec!["memo_transfer".to_string()],
                outputs: std::vec![kb_lib::MtApiMaterializedOutput {
                    output_key: "state".to_string(),
                    family: kb_lib::MdMaterializedEventFamily::TokenAccount,
                    payload_json: serde_json::json!({
                        "domain":"spl_token_2022_account_state",
                        "baseFields":{"mint":mint.0,"owner":owner.0}
                    }),
                }],
            },
        };
        let report = super::validate_snapshot_requirement(&requirement, &result);
        assert_eq!(
            report.as_ref().map(|value| return value.checks.len()),
            std::result::Result::Ok(6)
        );
    }
}