ark-core 0.10.1

use crate::boarding_output::BoardingOutput;
use crate::vhtlc::VhtlcOptions;
use crate::vtxo::Vtxo;
use crate::Error;
use bitcoin::absolute;
use bitcoin::key::Secp256k1;
use bitcoin::secp256k1::All;
use bitcoin::taproot::ControlBlock;
use bitcoin::Network;
use bitcoin::ScriptBuf;
use bitcoin::Sequence;
use bitcoin::XOnlyPublicKey;
use serde::Deserialize;
use serde::Serialize;
use std::fmt;

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ContractType(String);

impl ContractType {
    pub fn new(value: impl Into<String>) -> Result<Self, Error> {
        let value = value.into();
        if value.is_empty() {
            return Err(Error::ad_hoc("contract type cannot be empty"));
        }
        Ok(Self(value))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }

    pub fn default_vtxo() -> Self {
        Self("default".to_string())
    }

    pub fn delegate_vtxo() -> Self {
        Self("delegate".to_string())
    }

    pub fn boarding() -> Self {
        Self("boarding".to_string())
    }

    pub fn vhtlc() -> Self {
        Self("vhtlc".to_string())
    }
}

impl fmt::Display for ContractType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl From<&'static str> for ContractType {
    fn from(value: &'static str) -> Self {
        Self(value.to_string())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContractState {
    Active,
    Inactive,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct StoredContract {
    pub contract_type: ContractType,
    pub contract_version: u32,
    pub script_pubkey: ScriptBuf,
    pub state: ContractState,
    pub created_at: u64,
    /// HD derivation index for contracts generated by an index-based key provider.
    ///
    /// Static/custom providers may not expose an index, and imported contracts may not have one.
    pub key_index: Option<u32>,
    pub data: serde_json::Value,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SpendPathKind {
    Forfeit,
    Exit,
    Delegate,
    VhtlcClaim,
    VhtlcRefund,
    VhtlcRefundWithoutReceiver,
    VhtlcUnilateralClaim,
    VhtlcUnilateralRefund,
    VhtlcUnilateralRefundWithoutReceiver,
    Custom(String),
}

impl SpendPathKind {
    pub fn from_vhtlc_name(name: String) -> Self {
        match name.as_str() {
            "claim" => Self::VhtlcClaim,
            "refund" => Self::VhtlcRefund,
            "refund_without_receiver" => Self::VhtlcRefundWithoutReceiver,
            "unilateral_claim" => Self::VhtlcUnilateralClaim,
            "unilateral_refund" => Self::VhtlcUnilateralRefund,
            "unilateral_refund_without_receiver" => Self::VhtlcUnilateralRefundWithoutReceiver,
            _ => Self::Custom(name),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SpendPath {
    pub kind: SpendPathKind,
    pub script: ScriptBuf,
    pub control_block: ControlBlock,
}

impl SpendPath {
    pub fn new(kind: SpendPathKind, script: ScriptBuf, control_block: ControlBlock) -> Self {
        Self {
            kind,
            script,
            control_block,
        }
    }

    pub fn select(self) -> SpendSelection {
        SpendSelection::new(self)
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SpendSelection {
    pub path: SpendPath,
    pub sequence: Option<Sequence>,
    pub locktime: Option<absolute::LockTime>,
    pub extra_witness: Vec<Vec<u8>>,
}

impl SpendSelection {
    pub fn new(path: SpendPath) -> Self {
        Self {
            path,
            sequence: None,
            locktime: None,
            extra_witness: Vec::new(),
        }
    }

    pub fn with_sequence(mut self, sequence: Sequence) -> Self {
        self.sequence = Some(sequence);
        self
    }

    pub fn with_locktime(mut self, locktime: absolute::LockTime) -> Self {
        self.locktime = Some(locktime);
        self
    }

    pub fn with_extra_witness(mut self, extra_witness: Vec<Vec<u8>>) -> Self {
        self.extra_witness = extra_witness;
        self
    }

    pub fn resolved_sequence(&self, default_sequence: Sequence) -> Sequence {
        self.sequence.unwrap_or(default_sequence)
    }

    pub fn resolved_spend_info(
        &self,
        default_sequence: Sequence,
    ) -> (Sequence, (ScriptBuf, ControlBlock)) {
        (self.resolved_sequence(default_sequence), self.spend_info())
    }

    pub fn spend_info(&self) -> (ScriptBuf, ControlBlock) {
        (self.path.script.clone(), self.path.control_block.clone())
    }
}

#[derive(Clone)]
pub struct ContractContext {
    network: Network,
    secp: Secp256k1<All>,
}

impl ContractContext {
    pub fn new(network: Network) -> Self {
        Self {
            network,
            secp: Secp256k1::new(),
        }
    }

    pub fn network(&self) -> Network {
        self.network
    }

    pub fn secp(&self) -> &Secp256k1<All> {
        &self.secp
    }
}

pub trait ContractSpec:
    Clone + Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static
{
    const VERSION: u32;

    fn contract_type() -> ContractType;
    fn script_pubkey(&self, ctx: &ContractContext) -> Result<ScriptBuf, Error>;
    fn spendable_paths(&self, ctx: &ContractContext) -> Result<Vec<SpendPath>, Error>;

    fn spendable_selections(&self, ctx: &ContractContext) -> Result<Vec<SpendSelection>, Error> {
        Ok(self
            .spendable_paths(ctx)?
            .into_iter()
            .map(SpendPath::select)
            .collect())
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DefaultVtxoContract {
    pub server: XOnlyPublicKey,
    pub owner: XOnlyPublicKey,
    pub exit_delay: Sequence,
}

impl ContractSpec for DefaultVtxoContract {
    const VERSION: u32 = 1;

    fn contract_type() -> ContractType {
        ContractType::default_vtxo()
    }

    fn script_pubkey(&self, ctx: &ContractContext) -> Result<ScriptBuf, Error> {
        Ok(self.vtxo(ctx)?.script_pubkey())
    }

    fn spendable_paths(&self, ctx: &ContractContext) -> Result<Vec<SpendPath>, Error> {
        let vtxo = self.vtxo(ctx)?;
        let (forfeit_script, forfeit_control_block) = vtxo.forfeit_spend_info()?;
        let (exit_script, exit_control_block) = vtxo.exit_spend_info()?;
        Ok(vec![
            SpendPath::new(
                SpendPathKind::Forfeit,
                forfeit_script,
                forfeit_control_block,
            ),
            SpendPath::new(SpendPathKind::Exit, exit_script, exit_control_block),
        ])
    }

    fn spendable_selections(&self, ctx: &ContractContext) -> Result<Vec<SpendSelection>, Error> {
        Ok(self
            .spendable_paths(ctx)?
            .into_iter()
            .map(|path| {
                if path.kind == SpendPathKind::Exit {
                    path.select().with_sequence(self.exit_delay)
                } else {
                    path.select()
                }
            })
            .collect())
    }
}

impl DefaultVtxoContract {
    pub fn vtxo(&self, ctx: &ContractContext) -> Result<Vtxo, Error> {
        Vtxo::new_default(
            ctx.secp(),
            self.server,
            self.owner,
            self.exit_delay,
            ctx.network(),
        )
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DelegateVtxoContract {
    pub server: XOnlyPublicKey,
    pub owner: XOnlyPublicKey,
    pub delegator: XOnlyPublicKey,
    pub exit_delay: Sequence,
}

impl ContractSpec for DelegateVtxoContract {
    const VERSION: u32 = 1;

    fn contract_type() -> ContractType {
        ContractType::delegate_vtxo()
    }

    fn script_pubkey(&self, ctx: &ContractContext) -> Result<ScriptBuf, Error> {
        Ok(self.vtxo(ctx)?.script_pubkey())
    }

    fn spendable_paths(&self, ctx: &ContractContext) -> Result<Vec<SpendPath>, Error> {
        let vtxo = self.vtxo(ctx)?;
        let (forfeit_script, forfeit_control_block) = vtxo.forfeit_spend_info()?;
        let (exit_script, exit_control_block) = vtxo.exit_spend_info()?;
        let (delegate_script, delegate_control_block) = vtxo.delegate_spend_info()?;
        Ok(vec![
            SpendPath::new(
                SpendPathKind::Forfeit,
                forfeit_script,
                forfeit_control_block,
            ),
            SpendPath::new(SpendPathKind::Exit, exit_script, exit_control_block),
            SpendPath::new(
                SpendPathKind::Delegate,
                delegate_script,
                delegate_control_block,
            ),
        ])
    }

    fn spendable_selections(&self, ctx: &ContractContext) -> Result<Vec<SpendSelection>, Error> {
        Ok(self
            .spendable_paths(ctx)?
            .into_iter()
            .map(|path| {
                if path.kind == SpendPathKind::Exit {
                    path.select().with_sequence(self.exit_delay)
                } else {
                    path.select()
                }
            })
            .collect())
    }
}

impl DelegateVtxoContract {
    pub fn vtxo(&self, ctx: &ContractContext) -> Result<Vtxo, Error> {
        Vtxo::new_with_delegator(
            ctx.secp(),
            self.server,
            self.owner,
            self.delegator,
            self.exit_delay,
            ctx.network(),
        )
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BoardingContract {
    pub server: XOnlyPublicKey,
    pub owner: XOnlyPublicKey,
    pub exit_delay: Sequence,
}

impl ContractSpec for BoardingContract {
    const VERSION: u32 = 1;

    fn contract_type() -> ContractType {
        ContractType::boarding()
    }

    fn script_pubkey(&self, ctx: &ContractContext) -> Result<ScriptBuf, Error> {
        Ok(self.boarding_output(ctx)?.script_pubkey())
    }

    fn spendable_paths(&self, ctx: &ContractContext) -> Result<Vec<SpendPath>, Error> {
        let boarding_output = self.boarding_output(ctx)?;
        let (forfeit_script, forfeit_control_block) = boarding_output.forfeit_spend_info();
        let (exit_script, exit_control_block) = boarding_output.exit_spend_info();
        Ok(vec![
            SpendPath::new(
                SpendPathKind::Forfeit,
                forfeit_script,
                forfeit_control_block,
            ),
            SpendPath::new(SpendPathKind::Exit, exit_script, exit_control_block),
        ])
    }

    fn spendable_selections(&self, ctx: &ContractContext) -> Result<Vec<SpendSelection>, Error> {
        Ok(self
            .spendable_paths(ctx)?
            .into_iter()
            .map(|path| {
                if path.kind == SpendPathKind::Exit {
                    path.select().with_sequence(self.exit_delay)
                } else {
                    path.select()
                }
            })
            .collect())
    }
}

impl BoardingContract {
    pub fn boarding_output(&self, ctx: &ContractContext) -> Result<BoardingOutput, Error> {
        BoardingOutput::new(
            ctx.secp(),
            self.server,
            self.owner,
            self.exit_delay,
            ctx.network(),
        )
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct VhtlcContract {
    pub options: VhtlcOptions,
}

impl ContractSpec for VhtlcContract {
    const VERSION: u32 = 1;

    fn contract_type() -> ContractType {
        ContractType::vhtlc()
    }

    fn script_pubkey(&self, ctx: &ContractContext) -> Result<ScriptBuf, Error> {
        let script = crate::vhtlc::VhtlcScript::new(self.options.clone(), ctx.network())
            .map_err(|e| Error::ad_hoc(format!("failed to build vhtlc: {e}")))?;
        Ok(script.script_pubkey())
    }

    fn spendable_paths(&self, ctx: &ContractContext) -> Result<Vec<SpendPath>, Error> {
        let script = crate::vhtlc::VhtlcScript::new(self.options.clone(), ctx.network())
            .map_err(|e| Error::ad_hoc(format!("failed to build vhtlc: {e}")))?;
        script
            .get_script_map()
            .into_iter()
            .map(|(name, tapscript)| {
                let control_block = script
                    .taproot_spend_info()
                    .control_block(&(tapscript.clone(), bitcoin::taproot::LeafVersion::TapScript))
                    .ok_or_else(|| Error::ad_hoc("missing vhtlc control block"))?;
                Ok(SpendPath {
                    kind: SpendPathKind::from_vhtlc_name(name),
                    script: tapscript,
                    control_block,
                })
            })
            .collect()
    }

    fn spendable_selections(&self, ctx: &ContractContext) -> Result<Vec<SpendSelection>, Error> {
        Ok(self
            .spendable_paths(ctx)?
            .into_iter()
            .filter_map(|path| match path.kind {
                SpendPathKind::VhtlcClaim | SpendPathKind::VhtlcUnilateralClaim => None,
                SpendPathKind::VhtlcRefundWithoutReceiver => Some(path.select().with_locktime(
                    absolute::LockTime::from_consensus(self.options.refund_locktime),
                )),
                SpendPathKind::VhtlcUnilateralRefund => Some(
                    path.select()
                        .with_sequence(self.options.unilateral_refund_delay),
                ),
                SpendPathKind::VhtlcUnilateralRefundWithoutReceiver => Some(
                    path.select()
                        .with_sequence(self.options.unilateral_refund_without_receiver_delay),
                ),
                SpendPathKind::Forfeit
                | SpendPathKind::Exit
                | SpendPathKind::Delegate
                | SpendPathKind::VhtlcRefund
                | SpendPathKind::Custom(_) => Some(path.select()),
            })
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::hashes::ripemd160;
    use bitcoin::hashes::Hash;
    use std::str::FromStr;

    fn test_key(hex: &str) -> XOnlyPublicKey {
        XOnlyPublicKey::from_str(hex).unwrap()
    }

    fn vhtlc_contract() -> VhtlcContract {
        VhtlcContract {
            options: VhtlcOptions {
                sender: test_key(
                    "f874c4fd782a63a2b078b424f9f4e5edae622880a29fe572bd71bee45ab55ea0",
                ),
                receiver: test_key(
                    "b02b1a956cd0ed91d1a06b65ae976470ffd4b8c9215f446f266bb74fedd153b6",
                ),
                server: test_key(
                    "e35799157be4b37565bb5afe4d04e6a0fa0a4b6a4f4e48b0d904685d253cdbdb",
                ),
                preimage_hash: ripemd160::Hash::hash(b"preimage"),
                refund_locktime: 500,
                unilateral_claim_delay: Sequence::from_height(10),
                unilateral_refund_delay: Sequence::from_height(20),
                unilateral_refund_without_receiver_delay: Sequence::from_height(30),
            },
        }
    }

    #[test]
    fn vhtlc_spendable_selections_include_required_constraints() {
        let ctx = ContractContext::new(Network::Regtest);
        let selections = vhtlc_contract().spendable_selections(&ctx).unwrap();

        assert!(!selections
            .iter()
            .any(|selection| selection.path.kind == SpendPathKind::VhtlcClaim));
        assert!(!selections
            .iter()
            .any(|selection| selection.path.kind == SpendPathKind::VhtlcUnilateralClaim));

        let refund_without_receiver = selections
            .iter()
            .find(|selection| selection.path.kind == SpendPathKind::VhtlcRefundWithoutReceiver)
            .unwrap();
        assert_eq!(
            refund_without_receiver.locktime,
            Some(absolute::LockTime::from_consensus(500))
        );

        let unilateral_refund = selections
            .iter()
            .find(|selection| selection.path.kind == SpendPathKind::VhtlcUnilateralRefund)
            .unwrap();
        assert_eq!(unilateral_refund.sequence, Some(Sequence::from_height(20)));

        let unilateral_refund_without_receiver = selections
            .iter()
            .find(|selection| {
                selection.path.kind == SpendPathKind::VhtlcUnilateralRefundWithoutReceiver
            })
            .unwrap();
        assert_eq!(
            unilateral_refund_without_receiver.sequence,
            Some(Sequence::from_height(30))
        );
    }
}