// file: kb-pipeline/src/solana_token2022_proof_orchestration.rs
// version: 2

//! Deterministic orchestration contract for mixed Token-2022 proof locations.

/// Maximum proof references accepted by one confidential Token-2022 operation.
pub const MAX_TOKEN2022_OPERATION_PROOFS: usize = 5;

/// One bounded proof-orchestration request prepared before transaction assembly.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Token2022ProofOrchestrationRequest {
    /// Stable executor operation code.
    pub operation_code: std::string::String,
    /// Exact ordered proof references required by the operation.
    pub proofs: std::vec::Vec<kb_lib::ExSplTokenConfidentialProofReference>,
    /// Optional authority expected on every context-state account.
    pub expected_context_authority: std::option::Option<kb_lib::MdPubkey>,
    /// Maximum compute-unit limit accepted for the future transaction.
    pub compute_unit_limit: u32,
    /// Maximum total fee accepted for the future transaction.
    pub max_fee_lamports: u64,
    /// Whether a future submission was explicitly requested.
    pub submit: bool,
    /// Whether the operator explicitly confirmed the future submission.
    pub operator_confirmed: bool,
}

/// One deterministic mixed-proof orchestration result.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Token2022ProofOrchestrationReport {
    /// Stable executor operation code.
    pub operation_code: std::string::String,
    /// Whether the instructions sysvar must be present.
    pub instructions_sysvar_required: bool,
    /// Ordered non-zero inline offsets.
    pub inline_offsets: std::vec::Vec<i8>,
    /// Ordered validated context-state requirements.
    pub context_requirements: std::vec::Vec<crate::Token2022ProofContextRequirement>,
    /// Simulation is always mandatory.
    pub simulation_required: bool,
    /// Whether the future send path is authorized by request-local policy.
    pub send_authorized: bool,
    /// Ordered successful checks.
    pub checks: std::vec::Vec<std::string::String>,
}

/// Validates mixed inline/context-state proof orchestration before transaction assembly.
pub fn orchestrate_token2022_proofs(
    request: &crate::Token2022ProofOrchestrationRequest,
    cryptographic_preflight: &crate::Token2022CryptographicPreflightReport,
) -> kb_core::Result<crate::Token2022ProofOrchestrationReport> {
    if request.operation_code.trim().is_empty() {
        return std::result::Result::Err(kb_core::Error::config(
            "Token-2022 proof orchestration operation_code must not be empty",
        ));
    }
    if request.proofs.len() > crate::MAX_TOKEN2022_OPERATION_PROOFS {
        return std::result::Result::Err(kb_core::Error::new(
            "token2022_operation_proof_limit_exceeded",
            format!(
                "Token-2022 proof orchestration accepts at most {} proofs",
                crate::MAX_TOKEN2022_OPERATION_PROOFS
            ),
        ));
    }
    if request.compute_unit_limit == 0 || request.max_fee_lamports == 0 {
        return std::result::Result::Err(kb_core::Error::config(
            "Token-2022 proof orchestration requires non-zero compute and fee ceilings",
        ));
    }
    if request.submit && !request.operator_confirmed {
        return std::result::Result::Err(kb_core::Error::new(
            "token2022_submission_confirmation_required",
            "Token-2022 submission requires explicit operator confirmation",
        ));
    }
    let mut kinds = std::collections::BTreeSet::<std::string::String>::new();
    let mut offsets = std::collections::BTreeSet::<i8>::new();
    let mut context_accounts = std::collections::BTreeSet::<std::string::String>::new();
    let mut inline_offsets = std::vec::Vec::<i8>::new();
    let mut context_requirements = std::vec::Vec::<crate::Token2022ProofContextRequirement>::new();
    for proof in &request.proofs {
        let kind_code = proof_kind_code(proof.kind).to_string();
        if !kinds.insert(kind_code.clone()) {
            return std::result::Result::Err(kb_core::Error::new(
                "token2022_duplicate_proof_kind",
                format!(
                    "Token-2022 operation {} contains duplicate proof kind {kind_code}",
                    request.operation_code
                ),
            ));
        }
        if let std::result::Result::Err(error) = proof.location.validate() {
            return std::result::Result::Err(kb_core::Error::invalid_state(error));
        }
        match &proof.location {
            kb_lib::ExSplTokenConfidentialProofLocation::InstructionOffset { offset } => {
                if !offsets.insert(*offset) {
                    return std::result::Result::Err(kb_core::Error::new(
                        "token2022_duplicate_inline_proof_offset",
                        format!(
                            "Token-2022 operation {} reuses inline proof offset {offset}",
                            request.operation_code
                        ),
                    ));
                }
                inline_offsets.push(*offset);
            },
            kb_lib::ExSplTokenConfidentialProofLocation::ContextStateAccount { account } => {
                if !context_accounts.insert(account.0.clone()) {
                    return std::result::Result::Err(kb_core::Error::new(
                        "token2022_duplicate_context_state_account",
                        format!(
                            "Token-2022 operation {} reuses context-state account {}",
                            request.operation_code, account.0
                        ),
                    ));
                }
                context_requirements.push(crate::Token2022ProofContextRequirement {
                    role: kind_code,
                    account: account.clone(),
                    proof_type: proof_kind_to_zk_type(proof.kind),
                    expected_authority: request.expected_context_authority.clone(),
                });
            },
        }
    }
    if let std::result::Result::Err(error) =
        validate_context_reports(context_requirements.as_slice(), cryptographic_preflight)
    {
        return std::result::Result::Err(error);
    }
    return std::result::Result::Ok(crate::Token2022ProofOrchestrationReport {
        operation_code: request.operation_code.clone(),
        instructions_sysvar_required: !inline_offsets.is_empty(),
        inline_offsets,
        context_requirements,
        simulation_required: true,
        send_authorized: request.submit && request.operator_confirmed,
        checks: vec![
            "ordered_unique_proof_kinds".to_string(),
            "non_zero_unique_inline_offsets".to_string(),
            "unique_context_state_accounts".to_string(),
            "context_preflight_matches_operation".to_string(),
            "simulation_required".to_string(),
            "bounded_compute_and_fee".to_string(),
        ],
    });
}

fn validate_context_reports(
    requirements: &[crate::Token2022ProofContextRequirement],
    report: &crate::Token2022CryptographicPreflightReport,
) -> kb_core::Result<()> {
    if requirements.len() != report.proof_contexts.len() {
        return std::result::Result::Err(kb_core::Error::new(
            "token2022_context_preflight_count_mismatch",
            format!(
                "Token-2022 operation requires {} context states, preflight contains {}",
                requirements.len(),
                report.proof_contexts.len()
            ),
        ));
    }
    for (requirement, observed) in requirements.iter().zip(report.proof_contexts.iter()) {
        if requirement.account != observed.account
            || requirement.proof_type.discriminator() != observed.proof_discriminator
        {
            return std::result::Result::Err(kb_core::Error::new(
                "token2022_context_preflight_mismatch",
                format!(
                    "Token-2022 proof context {} does not match the ordered operation requirement",
                    observed.account.0
                ),
            ));
        }
    }
    return std::result::Result::Ok(());
}

fn proof_kind_code(kind: kb_lib::ExSplTokenConfidentialProofKind) -> &'static str {
    return match kind {
        kb_lib::ExSplTokenConfidentialProofKind::PubkeyValidity => "pubkey_validity",
        kb_lib::ExSplTokenConfidentialProofKind::ZeroCiphertext => "zero_ciphertext",
        kb_lib::ExSplTokenConfidentialProofKind::CiphertextCommitmentEquality => {
            "ciphertext_commitment_equality"
        },
        kb_lib::ExSplTokenConfidentialProofKind::CiphertextCiphertextEquality => {
            "ciphertext_ciphertext_equality"
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedGroupedCiphertext3HandlesValidity => {
            "batched_grouped_ciphertext_3_handles_validity"
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedGroupedCiphertext2HandlesValidity => {
            "batched_grouped_ciphertext_2_handles_validity"
        },
        kb_lib::ExSplTokenConfidentialProofKind::PercentageWithFee => "percentage_with_fee",
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU64 => "batched_range_proof_u64",
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU128 => {
            "batched_range_proof_u128"
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU256 => {
            "batched_range_proof_u256"
        },
    };
}

fn proof_kind_to_zk_type(
    kind: kb_lib::ExSplTokenConfidentialProofKind,
) -> kb_lib::ExSolanaCoreZkElGamalProofType {
    return match kind {
        kb_lib::ExSplTokenConfidentialProofKind::PubkeyValidity => {
            kb_lib::ExSolanaCoreZkElGamalProofType::PubkeyValidity
        },
        kb_lib::ExSplTokenConfidentialProofKind::ZeroCiphertext => {
            kb_lib::ExSolanaCoreZkElGamalProofType::ZeroCiphertext
        },
        kb_lib::ExSplTokenConfidentialProofKind::CiphertextCommitmentEquality => {
            kb_lib::ExSolanaCoreZkElGamalProofType::CiphertextCommitmentEquality
        },
        kb_lib::ExSplTokenConfidentialProofKind::CiphertextCiphertextEquality => {
            kb_lib::ExSolanaCoreZkElGamalProofType::CiphertextCiphertextEquality
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedGroupedCiphertext3HandlesValidity => {
            kb_lib::ExSolanaCoreZkElGamalProofType::BatchedGroupedCiphertext3HandlesValidity
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedGroupedCiphertext2HandlesValidity => {
            kb_lib::ExSolanaCoreZkElGamalProofType::BatchedGroupedCiphertext2HandlesValidity
        },
        kb_lib::ExSplTokenConfidentialProofKind::PercentageWithFee => {
            kb_lib::ExSolanaCoreZkElGamalProofType::PercentageWithCap
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU64 => {
            kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU64
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU128 => {
            kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU128
        },
        kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU256 => {
            kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU256
        },
    };
}

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

    fn context_report(
        account: kb_lib::MdPubkey,
        proof_type: kb_lib::ExSolanaCoreZkElGamalProofType,
    ) -> crate::Token2022ProofContextReport {
        return crate::Token2022ProofContextReport {
            role: "proof".to_string(),
            account,
            proof_discriminator: proof_type.discriminator(),
            context_state_bytes: proof_type.context_state_size(),
            context_slot: 100,
            checks: vec!["validated".to_string()],
        };
    }

    #[test]
    fn mixed_proofs_require_sysvar_and_preserve_ordered_contexts() {
        let context = pubkey(1);
        let request = crate::Token2022ProofOrchestrationRequest {
            operation_code: "spl_token2022.confidential_transfer".to_string(),
            proofs: vec![
                kb_lib::ExSplTokenConfidentialProofReference {
                    kind: kb_lib::ExSplTokenConfidentialProofKind::CiphertextCommitmentEquality,
                    location: kb_lib::ExSplTokenConfidentialProofLocation::InstructionOffset {
                        offset: 1,
                    },
                },
                kb_lib::ExSplTokenConfidentialProofReference {
                    kind: kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU128,
                    location: kb_lib::ExSplTokenConfidentialProofLocation::ContextStateAccount {
                        account: context.clone(),
                    },
                },
            ],
            expected_context_authority: std::option::Option::None,
            compute_unit_limit: 400_000,
            max_fee_lamports: 50_000,
            submit: false,
            operator_confirmed: false,
        };
        let preflight = crate::Token2022CryptographicPreflightReport {
            commitment: "confirmed".to_string(),
            context_slot: 100,
            proof_contexts: vec![context_report(
                context,
                kb_lib::ExSolanaCoreZkElGamalProofType::BatchedRangeProofU128,
            )],
        };
        let report = crate::orchestrate_token2022_proofs(&request, &preflight)
            .expect("mixed proof orchestration must succeed");
        assert!(report.instructions_sysvar_required);
        assert_eq!(report.inline_offsets, vec![1]);
        assert_eq!(report.context_requirements.len(), 1);
        assert!(report.simulation_required);
        assert!(!report.send_authorized);
    }

    #[test]
    fn context_only_proofs_do_not_require_instructions_sysvar() {
        let context = pubkey(2);
        let request = crate::Token2022ProofOrchestrationRequest {
            operation_code: "spl_token2022.empty_confidential_account".to_string(),
            proofs: vec![kb_lib::ExSplTokenConfidentialProofReference {
                kind: kb_lib::ExSplTokenConfidentialProofKind::ZeroCiphertext,
                location: kb_lib::ExSplTokenConfidentialProofLocation::ContextStateAccount {
                    account: context.clone(),
                },
            }],
            expected_context_authority: std::option::Option::None,
            compute_unit_limit: 200_000,
            max_fee_lamports: 20_000,
            submit: true,
            operator_confirmed: true,
        };
        let preflight = crate::Token2022CryptographicPreflightReport {
            commitment: "confirmed".to_string(),
            context_slot: 100,
            proof_contexts: vec![context_report(
                context,
                kb_lib::ExSolanaCoreZkElGamalProofType::ZeroCiphertext,
            )],
        };
        let report = crate::orchestrate_token2022_proofs(&request, &preflight)
            .expect("context-only orchestration must succeed");
        assert!(!report.instructions_sysvar_required);
        assert!(report.send_authorized);
    }

    #[test]
    fn duplicate_offsets_contexts_and_unconfirmed_submission_fail_closed() {
        let duplicate_offset = crate::Token2022ProofOrchestrationRequest {
            operation_code: "operation".to_string(),
            proofs: vec![
                kb_lib::ExSplTokenConfidentialProofReference {
                    kind: kb_lib::ExSplTokenConfidentialProofKind::ZeroCiphertext,
                    location: kb_lib::ExSplTokenConfidentialProofLocation::InstructionOffset {
                        offset: 1,
                    },
                },
                kb_lib::ExSplTokenConfidentialProofReference {
                    kind: kb_lib::ExSplTokenConfidentialProofKind::BatchedRangeProofU64,
                    location: kb_lib::ExSplTokenConfidentialProofLocation::InstructionOffset {
                        offset: 1,
                    },
                },
            ],
            expected_context_authority: std::option::Option::None,
            compute_unit_limit: 1,
            max_fee_lamports: 1,
            submit: false,
            operator_confirmed: false,
        };
        let empty = crate::Token2022CryptographicPreflightReport {
            commitment: "confirmed".to_string(),
            context_slot: 0,
            proof_contexts: vec![],
        };
        assert!(crate::orchestrate_token2022_proofs(&duplicate_offset, &empty).is_err());
        let mut unconfirmed = duplicate_offset;
        unconfirmed.proofs.clear();
        unconfirmed.submit = true;
        assert!(crate::orchestrate_token2022_proofs(&unconfirmed, &empty).is_err());
    }
}