// file: ks-pipeline/src/spl_token_2022_correlation.rs
// version: 4

//! Deterministic Token-2022 instruction-to-state correlation.

/// Correlation outcome between one committed instruction fact and one final state snapshot.
#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)]
pub enum Token2022CorrelationStatus {
    /// The final snapshot contains the extension expected by the instruction fact.
    Confirmed,
    /// The final snapshot does not contain the extension expected by the instruction fact.
    Contradicted,
    /// The instruction output is outside the supported correlation inventory.
    NotApplicable,
}

/// Deterministic correlation report for one instruction output and one final snapshot.
#[derive(Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
pub struct Token2022StateCorrelationReport {
    /// Stable caller-provided correlation identity.
    pub correlation_key: std::string::String,
    /// Canonical account identity checked against the snapshot.
    pub account: ks_lib::MdPubkey,
    /// Instruction operation or risk fact used for correlation.
    pub fact_code: std::string::String,
    /// Extension expected from the fact when the fact is supported.
    pub expected_extension: std::option::Option<std::string::String>,
    /// Final correlation outcome.
    pub status: Token2022CorrelationStatus,
    /// Snapshot context slot.
    pub snapshot_slot: u64,
    /// Whether the snapshot contains the expected extension.
    pub extension_present: std::option::Option<bool>,
    /// Explicit semantic limitation of this correlation.
    pub fact_only: bool,
}

/// Correlates one materialized instruction fact with one authoritative Token-2022 snapshot.
pub fn correlate_token_2022_instruction_with_snapshot(
    correlation_key: &str,
    account: &ks_lib::MdPubkey,
    instruction_output: &ks_lib::MtApiMaterializedOutput,
    snapshot: &crate::Token2022StatefulSnapshotBundle,
) -> ks_core::Result<Token2022StateCorrelationReport> {
    if correlation_key.trim().is_empty() {
        return std::result::Result::Err(ks_core::Error::config(
            "Token-2022 correlation key must not be empty",
        ));
    }
    if snapshot.account_key != account.0 {
        return std::result::Result::Err(ks_core::Error::new(
            "token_2022_correlation_account_mismatch",
            format!(
                "Token-2022 correlation expected account {}, got snapshot {}",
                account.0, snapshot.account_key
            ),
        ));
    }
    let fact_code = match instruction_output.payload_json.get("riskKind") {
        std::option::Option::Some(value) => value.as_str().unwrap_or("").to_string(),
        std::option::Option::None => instruction_output
            .payload_json
            .get("operation")
            .and_then(serde_json::Value::as_str)
            .unwrap_or("")
            .to_string(),
    };
    if fact_code.is_empty() {
        return std::result::Result::Err(ks_core::Error::new(
            "token_2022_correlation_fact_missing",
            "Token-2022 correlation requires riskKind or operation in the instruction output",
        ));
    }
    let expected_extension = expected_extension(fact_code.as_str());
    let extension_present = expected_extension.map(|extension| {
        return snapshot.extension_names.iter().any(|candidate| return candidate == extension);
    });
    let status = match extension_present {
        std::option::Option::Some(true) => Token2022CorrelationStatus::Confirmed,
        std::option::Option::Some(false) => Token2022CorrelationStatus::Contradicted,
        std::option::Option::None => Token2022CorrelationStatus::NotApplicable,
    };
    return std::result::Result::Ok(Token2022StateCorrelationReport {
        correlation_key: correlation_key.to_string(),
        account: account.clone(),
        fact_code,
        expected_extension: expected_extension.map(str::to_string),
        status,
        snapshot_slot: snapshot.slot,
        extension_present,
        fact_only: true,
    });
}

fn expected_extension(fact_code: &str) -> std::option::Option<&'static str> {
    return match fact_code {
        "default_account_state_configured"
        | "initialize_default_account_state"
        | "update_default_account_state" => std::option::Option::Some("default_account_state"),
        "required_transfer_memos_enabled"
        | "required_transfer_memos_disabled"
        | "enable_required_transfer_memos"
        | "disable_required_transfer_memos" => std::option::Option::Some("memo_transfer"),
        "cpi_guard_enabled" | "cpi_guard_disabled" | "enable_cpi_guard" | "disable_cpi_guard" => {
            std::option::Option::Some("cpi_guard")
        },
        "non_transferable_mint_configured" | "initialize_non_transferable_mint" => {
            std::option::Option::Some("non_transferable")
        },
        "pausable_mint_configured"
        | "mint_activity_paused"
        | "mint_activity_resumed"
        | "initialize_pausable_config"
        | "pause"
        | "resume" => std::option::Option::Some("pausable"),
        "permissioned_burn_configured" | "initialize_permissioned_burn" => {
            std::option::Option::Some("permissioned_burn")
        },
        "confidential_credits_enabled"
        | "confidential_credits_disabled"
        | "non_confidential_credits_enabled"
        | "non_confidential_credits_disabled"
        | "enable_confidential_credits"
        | "disable_confidential_credits"
        | "enable_non_confidential_credits"
        | "disable_non_confidential_credits" => {
            std::option::Option::Some("confidential_transfer_account")
        },
        "initialize_transfer_fee_config"
        | "set_transfer_fee"
        | "withdraw_withheld_tokens_from_mint"
        | "harvest_withheld_tokens_to_mint" => std::option::Option::Some("transfer_fee_config"),
        "withdraw_withheld_tokens_from_accounts" | "transfer_checked_with_fee" => {
            std::option::Option::Some("transfer_fee_amount")
        },
        "initialize_confidential_transfer_fee_config"
        | "withdraw_confidential_withheld_tokens_from_mint"
        | "harvest_confidential_withheld_tokens_to_mint"
        | "enable_confidential_harvest_to_mint"
        | "disable_confidential_harvest_to_mint" => {
            std::option::Option::Some("confidential_transfer_fee_config")
        },
        "withdraw_confidential_withheld_tokens_from_accounts"
        | "transfer_confidential_tokens_with_fee" => {
            std::option::Option::Some("confidential_transfer_fee_amount")
        },
        _ => std::option::Option::None,
    };
}

#[cfg(test)]
mod tests {
    fn snapshot(extension_names: &[&str]) -> crate::Token2022StatefulSnapshotBundle {
        return crate::Token2022StatefulSnapshotBundle {
            account_key: solana_pubkey::Pubkey::new_from_array([7u8; 32]).to_string(),
            slot: 42,
            state_kind: "mint".to_string(),
            extension_names: extension_names
                .iter()
                .map(|value| return (*value).to_string())
                .collect(),
            outputs: std::vec::Vec::new(),
        };
    }

    fn output(field: &str, value: &str) -> ks_lib::MtApiMaterializedOutput {
        let mut payload = serde_json::Map::new();
        payload.insert(field.to_string(), serde_json::Value::String(value.to_string()));
        return ks_lib::MtApiMaterializedOutput {
            output_key: "fact:0".to_string(),
            family: ks_lib::MdMaterializedEventFamily::Risk,
            payload_json: serde_json::Value::Object(payload),
        };
    }

    #[test]
    fn known_fact_is_confirmed_or_contradicted_by_the_final_extension_inventory() {
        let snapshot = snapshot(&["pausable"]);
        let account = ks_lib::MdPubkey(snapshot.account_key.clone());
        let confirmed = super::correlate_token_2022_instruction_with_snapshot(
            "corr:pause",
            &account,
            &output("riskKind", "mint_activity_paused"),
            &snapshot,
        );
        assert!(confirmed.is_ok());
        if let std::result::Result::Ok(confirmed) = confirmed {
            assert_eq!(confirmed.status, super::Token2022CorrelationStatus::Confirmed);
            assert_eq!(confirmed.extension_present, std::option::Option::Some(true));
        }
        let contradicted = super::correlate_token_2022_instruction_with_snapshot(
            "corr:fee",
            &account,
            &output("operation", "initialize_transfer_fee_config"),
            &snapshot,
        );
        assert!(contradicted.is_ok());
        if let std::result::Result::Ok(contradicted) = contradicted {
            assert_eq!(contradicted.status, super::Token2022CorrelationStatus::Contradicted);
            assert_eq!(contradicted.extension_present, std::option::Option::Some(false));
        }
    }

    #[test]
    fn unsupported_fact_is_explicitly_not_applicable_without_inventing_a_conclusion() {
        let snapshot = snapshot(&[]);
        let account = ks_lib::MdPubkey(snapshot.account_key.clone());
        let report = super::correlate_token_2022_instruction_with_snapshot(
            "corr:unknown",
            &account,
            &output("operation", "unknown_future_operation"),
            &snapshot,
        );
        assert!(report.is_ok());
        if let std::result::Result::Ok(report) = report {
            assert_eq!(report.status, super::Token2022CorrelationStatus::NotApplicable);
            assert_eq!(report.extension_present, std::option::Option::None);
            assert!(report.fact_only);
        }
    }

    #[test]
    fn correlation_identity_and_account_mismatch_fail_closed() {
        let snapshot = snapshot(&["memo_transfer"]);
        let account = ks_lib::MdPubkey(snapshot.account_key.clone());
        let empty = super::correlate_token_2022_instruction_with_snapshot(
            "",
            &account,
            &output("riskKind", "required_transfer_memos_enabled"),
            &snapshot,
        );
        assert!(empty.is_err());
        let wrong = ks_lib::MdPubkey(solana_pubkey::Pubkey::new_from_array([8u8; 32]).to_string());
        let mismatch = super::correlate_token_2022_instruction_with_snapshot(
            "corr:mismatch",
            &wrong,
            &output("riskKind", "required_transfer_memos_enabled"),
            &snapshot,
        );
        assert!(mismatch.is_err());
    }
}