charms 15.0.0

use super::prove_spell_tx::ProveSpellTxImpl;
use crate::tx::bitcoin_tx::from_spell;
use anyhow::{Context, anyhow, bail, ensure};
use bitcoin::Network;
use charms_app_runner::AppRunner;
use charms_client::{
    BeamSource, NormalizedSpell, SignedScrollOutputs,
    cardano_tx::OutputContent,
    ensure_no_orphan_versioned_apps,
    tx::{Chain, Tx, by_txid},
};
use charms_data::{App, AppInput, AppSignature, B32, Data, TxId, util};
use charms_lib::SPELL_VK;
use std::{
    collections::{BTreeMap, BTreeSet},
    str::FromStr,
};

use super::get_charms_fee;

pub fn ensure_exact_app_binaries(
    norm_spell: &NormalizedSpell,
    app_private_inputs: &BTreeMap<App, Data>,
    tx: &charms_data::Transaction,
    binaries: &BTreeMap<B32, Vec<u8>>,
) -> anyhow::Result<()> {
    let required_binary_hashes: BTreeSet<B32> = norm_spell
        .app_public_inputs
        .iter()
        .filter(|(app, data)| {
            !data.is_empty()
                || !app_private_inputs
                    .get(app)
                    .is_none_or(|data| data.is_empty())
                || !charms_data::is_simple_transfer(app, tx)
        })
        .map(|(app, _)| match norm_spell.versioned_apps.get(&app.vk) {
            Some(va) => va.wasm_hash.clone(),
            None => app.vk.clone(),
        })
        .collect();

    let provided_binary_hashes: BTreeSet<B32> = binaries.keys().cloned().collect();

    ensure!(
        required_binary_hashes == provided_binary_hashes,
        "binaries must contain exactly the required app binaries.\n\
         Required binary hashes: {:?}\n\
         Provided binary hashes: {:?}",
        required_binary_hashes,
        provided_binary_hashes
    );

    Ok(())
}

/// `app_signatures` must contain exactly one entry per **non-simple-transfer** versioned
/// app referenced in the spell. Simple-transfer versioned apps are intentionally excluded:
/// the spell-level continuity check pins their `(version, wasm_hash)` to whatever the
/// previous spell already authenticated, so no fresh binary or signature is needed.
///
/// "Simple transfer" here matches [`charms_data::is_simple_transfer`] AND has no public or
/// private input — exactly the condition under which [`AppRunner::run_all`] skips the
/// binary entirely.
pub fn ensure_versioned_apps_have_signatures(
    norm_spell: &NormalizedSpell,
    app_private_inputs: &BTreeMap<App, Data>,
    tx: &charms_data::Transaction,
    app_signatures: &BTreeMap<B32, AppSignature>,
) -> anyhow::Result<()> {
    let required_vks: BTreeSet<&B32> = norm_spell
        .app_public_inputs
        .iter()
        .filter(|(app, _)| norm_spell.versioned_apps.contains_key(&app.vk))
        .filter(|(app, data)| {
            !data.is_empty()
                || !app_private_inputs.get(app).is_none_or(|w| w.is_empty())
                || !charms_data::is_simple_transfer(app, tx)
        })
        .map(|(app, _)| &app.vk)
        .collect();
    let provided_vks: BTreeSet<&B32> = app_signatures.keys().collect();
    ensure!(
        required_vks == provided_vks,
        "app_signatures must contain exactly one entry per non-simple-transfer versioned app \
         referenced in the spell.\n\
         Required app vks: {:?}\n\
         Provided app vks: {:?}",
        required_vks,
        provided_vks
    );
    Ok(())
}

pub fn ensure_unique_spell_inputs(spell: &NormalizedSpell) -> anyhow::Result<()> {
    let spell_ins = spell
        .tx
        .ins
        .as_ref()
        .ok_or_else(|| anyhow!("spell.tx.ins must be present"))?;
    let mut seen = BTreeSet::new();
    for utxo_id in spell_ins {
        ensure!(
            seen.insert(utxo_id),
            "spell.tx.ins contains duplicate UTXO: {}",
            utxo_id
        );
    }
    Ok(())
}

pub fn ensure_all_prev_txs_are_present(
    spell: &NormalizedSpell,
    tx_ins_beamed_source_utxos: &BTreeMap<usize, BeamSource>,
    prev_txs_by_id: &BTreeMap<TxId, Tx>,
) -> anyhow::Result<()> {
    let spell_ins = spell
        .tx
        .ins
        .as_ref()
        .ok_or_else(|| anyhow!("spell.tx.ins must be present"))?;

    ensure!(
        spell_ins
            .iter()
            .all(|utxo_id| prev_txs_by_id.contains_key(&utxo_id.0)),
        "prev_txs MUST contain transactions creating input UTXOs"
    );
    ensure!(
        spell.tx.refs.as_ref().is_none_or(|ins| {
            ins.iter()
                .all(|utxo_id| prev_txs_by_id.contains_key(&utxo_id.0))
        }),
        "prev_txs MUST contain transactions creating ref UTXOs"
    );
    ensure!(
        tx_ins_beamed_source_utxos
            .iter()
            .all(|(&i, beaming_source)| {
                spell_ins.get(i).is_some_and(|utxo_id| {
                    prev_txs_by_id.contains_key(&utxo_id.0)
                        && prev_txs_by_id.contains_key(&(beaming_source.0).0)
                })
            }),
        "prev_txs MUST contain transactions creating beaming source and destination UTXOs"
    );

    // Ensure prev_txs contains ONLY the required transactions (no extras)
    let mut required_txids = BTreeSet::new();

    // Add transaction IDs from spell inputs
    required_txids.extend(spell_ins.iter().map(|utxo_id| &utxo_id.0));

    // Add transaction IDs from spell refs
    if let Some(refs) = spell.tx.refs.as_ref() {
        required_txids.extend(refs.iter().map(|utxo_id| &utxo_id.0));
    }

    // Add transaction IDs from beaming source UTXOs
    required_txids.extend(tx_ins_beamed_source_utxos.values().map(|bs| &(bs.0).0));

    // Check that prev_txs contains exactly the required transactions
    let provided_txids: BTreeSet<_> = prev_txs_by_id.keys().collect();

    ensure!(
        required_txids == provided_txids,
        "prev_txs must contain exactly the transactions producing spell inputs and beaming sources.\n\
         Required: {:?}\n\
         Provided: {:?}",
        required_txids,
        provided_txids
    );

    Ok(())
}

/// Adjust `NativeOutput.content` fields in `norm_spell.tx.coins` according to the target chain.
///
/// - **Cardano**: each `content` is canonicalized through JSON→[`OutputContent`]→[`Data`]
///   round-trip.  If `content` is `None`, the default (empty) [`OutputContent`] is used.
/// - **Bitcoin**: every `content` field **must** be `None`; otherwise an error is returned.
pub fn adjust_coin_contents(norm_spell: &mut NormalizedSpell, chain: Chain) -> anyhow::Result<()> {
    let Some(coins) = norm_spell.tx.coins.as_mut() else {
        bail!("coins must be present");
    };
    ensure!(
        coins.len() == norm_spell.tx.outs.len(),
        "coins length ({}) must match outs length ({})",
        coins.len(),
        norm_spell.tx.outs.len()
    );

    for (i, coin) in coins.iter_mut().enumerate() {
        match chain {
            Chain::Bitcoin => {
                ensure!(
                    coin.content.is_none(),
                    "coins[{i}].content must be None for Bitcoin"
                );
            }
            Chain::Cardano => {
                let output_content: OutputContent = match coin.content.take() {
                    Some(content) => {
                        let json = serde_json::to_value(&content).with_context(|| {
                            format!("coins[{i}].content: failed to serialize to JSON")
                        })?;
                        serde_json::from_value(json).with_context(|| {
                            format!("coins[{i}].content: failed to parse as OutputContent")
                        })?
                    }
                    None => OutputContent::default(),
                };
                coin.content = Some((&output_content).into());
            }
        }
    }

    Ok(())
}

impl ProveSpellTxImpl {
    /// Validate a `ProveRequest`, returning the app-cycle total and the `is_correct`
    /// result. `is_correct` may return `Ok(false)` on the client preflight path when
    /// the spell has unbound Scrolls outputs (the prover server fills those in via
    /// `fill_scroll_outputs`); the caller decides whether to treat that as fatal.
    pub fn validate_prove_request(
        &self,
        prove_request: &mut super::request::ProveRequest,
        scroll_outputs: Option<&SignedScrollOutputs>,
    ) -> anyhow::Result<(u64, bool)> {
        ensure!(
            prove_request.spell.mock == self.mock,
            "cannot prove a mock=={} spell on a mock=={} prover",
            prove_request.spell.mock,
            self.mock
        );

        let prev_txs = &prove_request.prev_txs;
        let prev_txs_by_id = by_txid(prev_txs);

        let norm_spell = &mut prove_request.spell;
        adjust_coin_contents(norm_spell, prove_request.chain)?;
        let app_private_inputs = &prove_request.app_private_inputs;
        let tx_ins_beamed_source_utxos = &prove_request.tx_ins_beamed_source_utxos;

        ensure_all_prev_txs_are_present(&norm_spell, tx_ins_beamed_source_utxos, &prev_txs_by_id)?;
        ensure_unique_spell_inputs(&norm_spell)?;

        let prev_spells = charms_client::prev_spells(prev_txs, &SPELL_VK, &norm_spell)?;

        let tx = charms_client::to_tx(
            &norm_spell,
            &prev_spells,
            &tx_ins_beamed_source_utxos,
            &prev_txs,
        );

        ensure_exact_app_binaries(
            &norm_spell,
            &app_private_inputs,
            &tx,
            &prove_request.binaries,
        )?;

        let app_signatures = prove_request.app_signatures.clone();
        ensure_no_orphan_versioned_apps(&norm_spell)?;
        ensure_versioned_apps_have_signatures(
            &norm_spell,
            &app_private_inputs,
            &tx,
            &app_signatures,
        )?;

        let app_input = match prove_request.binaries.is_empty() {
            true => None,
            false => Some(AppInput {
                app_binaries: prove_request.binaries.clone(),
                app_private_inputs: app_private_inputs.clone(),
                app_signatures: app_signatures.clone(),
            }),
        };

        let verified = charms_client::is_correct(
            &norm_spell,
            &prev_txs,
            app_input.clone(),
            &SPELL_VK,
            &tx_ins_beamed_source_utxos,
            scroll_outputs,
        )?;

        // Calculate cycles for fee estimation
        let total_cycles = if let Some(app_input) = &app_input {
            let version_changed_apps = charms_client::collect_version_changed_apps(
                &norm_spell,
                &prev_spells,
                &tx_ins_beamed_source_utxos,
            );
            let cycles = AppRunner::new(true).run_all(
                &app_input.app_binaries,
                &norm_spell.versioned_apps,
                &app_input.app_signatures,
                &tx,
                &norm_spell.app_public_inputs,
                &app_input.app_private_inputs,
                &version_changed_apps,
            )?;
            cycles.iter().sum()
        } else {
            0
        };

        match prove_request.chain {
            Chain::Bitcoin => {
                let change_address = bitcoin::Address::from_str(&prove_request.change_address)?;

                let network = match &change_address {
                    a if a.is_valid_for_network(Network::Bitcoin) => Network::Bitcoin,
                    a if a.is_valid_for_network(Network::Testnet4) => Network::Testnet4,
                    a if a.is_valid_for_network(Network::Regtest) && self.mock => Network::Regtest,
                    _ => bail!(
                        "Unsupported network of change address: {:?}",
                        change_address
                    ),
                };
                let coin_outs = (norm_spell.tx.coins.as_ref()).expect("coin outputs are expected");

                // Validate that all output addresses are valid for the network.
                // Outputs listed in `tx.scrolls` may legitimately have `dest = []`
                // on the client preflight path -- the prover server fills them in
                // via `fill_scroll_outputs`, and `check_scroll_outputs` later pins
                // each `dest` to the canister-signed scriptPubKey.
                let scroll_indexes: BTreeSet<u32> =
                    norm_spell.tx.scrolls.clone().unwrap_or_default();
                ensure!(
                    coin_outs.iter().enumerate().all(|(i, o)| {
                        if scroll_indexes.contains(&(i as u32)) && o.dest.is_empty() {
                            return true;
                        }
                        bitcoin::Address::from_script(
                            &bitcoin::ScriptBuf::from_bytes(o.dest.clone()),
                            network,
                        )
                        .is_ok()
                    }),
                    "all output addresses must be valid for the network"
                );

                let charms_fee = get_charms_fee(&self.charms_fee_settings, total_cycles).to_sat();

                let spell_ins = norm_spell
                    .tx
                    .ins
                    .as_ref()
                    .expect("spell inputs are expected");
                let total_sats_in: u64 = spell_ins
                    .iter()
                    .map(|utxo_id| {
                        prev_txs_by_id
                            .get(&utxo_id.0)
                            .and_then(|prev_tx| {
                                if let Tx::Bitcoin(bitcoin_tx) = prev_tx {
                                    bitcoin_tx
                                        .inner()
                                        .output
                                        .get(utxo_id.1 as usize)
                                        .map(|o| o.value.to_sat())
                                } else {
                                    None
                                }
                            })
                            .ok_or(anyhow!("utxo not found in prev_txs: {}", utxo_id))
                    })
                    .try_fold(0u64, |acc, v| {
                        acc.checked_add(v?)
                            .ok_or_else(|| anyhow!("total input sats overflow u64"))
                    })?;
                let total_sats_out: u64 = coin_outs.iter().try_fold(0u64, |acc, o| {
                    acc.checked_add(o.amount)
                        .ok_or_else(|| anyhow!("total output sats overflow u64"))
                })?;

                let bitcoin_tx = from_spell(&norm_spell)?;
                let tx_size = bitcoin_tx.inner().vsize();
                let mut norm_spell_for_size = norm_spell.clone();
                norm_spell_for_size.tx.ins = None;
                let proof_dummy: Vec<u8> = vec![0xff; 128];
                let spell_cbor = util::write(&(norm_spell_for_size, proof_dummy))?;
                let num_inputs = bitcoin_tx.inner().input.len();
                let estimated_bitcoin_fee: u64 = (111
                    + (spell_cbor.len() as u64 + 372) / 4
                    + tx_size as u64
                    + 28 * num_inputs as u64)
                    * prove_request.fee_rate as u64;

                tracing::info!(
                    total_sats_in,
                    total_sats_out,
                    charms_fee,
                    estimated_bitcoin_fee
                );

                let total_sats_required = total_sats_out
                    .checked_add(charms_fee)
                    .and_then(|s| s.checked_add(estimated_bitcoin_fee))
                    .ok_or_else(|| anyhow!("total required sats (outputs + fees) overflow u64"))?;

                ensure!(
                    total_sats_in > total_sats_required,
                    "spell inputs must have sufficient value to cover outputs and fees"
                );
                Ok((total_cycles, verified))
            }
            Chain::Cardano => {
                // TODO
                tracing::warn!("spell validation for cardano is not yet implemented");
                Ok((total_cycles, verified))
            }
        }
    }
}