ai-agent-bitcoin-escrow 0.1.0

//! Escrow contract management for AI agents.
//!
//! This module provides the main escrow contract functionality including
//! creation, funding, condition evaluation, and fund release.

use bitcoin::{Address, Network};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::str::FromStr;

use crate::audit::{AuditEventType, AuditLogger};
use crate::conditions::{ConditionEvaluator, StandardEvaluator};
use crate::error::{EscrowError, Result};
use crate::multisig::{MultisigConfig, MultisigWallet};
use crate::oracle::OracleRegistry;
use crate::types::{
    ConditionResult, EscrowConfig, EscrowId, EscrowOutput, EscrowParticipant,
    EscrowStatus, ReleaseCondition, SignedTransaction,
};

/// An escrow contract.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EscrowContract {
    /// Unique identifier.
    pub id: EscrowId,
    /// Configuration.
    pub config: EscrowConfig,
    /// Participants.
    pub participants: Vec<EscrowParticipant>,
    /// Multisig wallet (if created).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub multisig: Option<MultisigWallet>,
    /// Current status.
    pub status: EscrowStatus,
    /// Release conditions.
    pub conditions: Vec<ReleaseCondition>,
    /// Funded outputs (UTXOs).
    pub outputs: Vec<EscrowOutput>,
    /// Total amount in satoshis.
    pub amount_sat: u64,
    /// Recipient address for release (serialized as string).
    #[serde(skip_serializing_if = "Option::is_none", serialize_with = "serialize_address_opt", deserialize_with = "deserialize_address_opt")]
    pub release_address: Option<Address>,
    /// Transaction signed for release.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub release_tx: Option<SignedTransaction>,
    /// Creation timestamp.
    pub created_at: DateTime<Utc>,
    /// Last update timestamp.
    pub updated_at: DateTime<Utc>,
    /// Optional metadata.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<HashMap<String, serde_json::Value>>,
}

fn serialize_address_opt<S>(addr: &Option<Address>, s: S) -> std::result::Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    match addr {
        Some(a) => s.serialize_str(&a.to_string()),
        None => s.serialize_none(),
    }
}

fn deserialize_address_opt<'de, D>(d: D) -> std::result::Result<Option<Address>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    use serde::de::Error;
    let opt: Option<String> = Option::deserialize(d)?;
    match opt {
        Some(s) => {
            let addr = Address::from_str(&s).map_err(|e| D::Error::custom(e.to_string()))?;
            Ok(Some(addr.assume_checked()))
        }
        None => Ok(None),
    }
}

impl EscrowContract {
    /// Create a new escrow contract.
    pub fn new(config: EscrowConfig, description: Option<String>) -> Self {
        Self {
            id: EscrowId::new(),
            config,
            participants: Vec::new(),
            multisig: None,
            status: EscrowStatus::Created,
            conditions: Vec::new(),
            outputs: Vec::new(),
            amount_sat: 0,
            release_address: None,
            release_tx: None,
            created_at: Utc::now(),
            updated_at: Utc::now(),
            metadata: if description.is_some() {
                let mut m = HashMap::new();
                m.insert("description".to_string(), serde_json::json!(description));
                Some(m)
            } else {
                None
            },
        }
    }

    /// Add a participant to the escrow.
    pub fn add_participant(&mut self, participant: EscrowParticipant) -> Result<()> {
        if self.status != EscrowStatus::Created {
            return Err(EscrowError::InvalidState(
                "Cannot add participant after escrow is funded".to_string(),
            ));
        }

        if self.participants.len() >= self.config.total_participants {
            return Err(EscrowError::Contract(
                "Maximum participants reached".to_string(),
            ));
        }

        self.participants.push(participant);
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Initialize the multisig wallet.
    pub fn initialize_multisig(&mut self) -> Result<()> {
        if self.multisig.is_some() {
            return Err(EscrowError::Contract("Multisig already initialized".to_string()));
        }

        let multisig_config = MultisigConfig {
            network: self.config.network,
            threshold: self.config.threshold,
            total: self.config.total_participants,
        };

        let wallet = MultisigWallet::new(multisig_config, &self.participants)?;
        self.multisig = Some(wallet);
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Get the escrow address for funding.
    pub fn escrow_address(&self) -> Result<Address> {
        let multisig = self
            .multisig
            .as_ref()
            .ok_or_else(|| EscrowError::Contract("Multisig not initialized".to_string()))?;

        multisig.address()
    }

    /// Add a funding output.
    pub fn add_output(&mut self, output: EscrowOutput) -> Result<()> {
        if self.multisig.is_none() {
            return Err(EscrowError::Contract("Multisig not initialized".to_string()));
        }

        // Verify the output belongs to this escrow
        let expected_script = self.multisig.as_ref().unwrap().script_pubkey()?;
        if output.script_pubkey != expected_script {
            return Err(EscrowError::Contract(
                "Output does not belong to this escrow".to_string(),
            ));
        }

        self.amount_sat += output.amount;
        self.outputs.push(output);
        self.status = EscrowStatus::Funded;
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Add a release condition.
    pub fn add_condition(&mut self, condition: ReleaseCondition) -> Result<()> {
        self.conditions.push(condition);
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Check if the escrow has expired.
    pub fn is_expired(&self) -> bool {
        if let Some(expires_at) = self.config.expires_at {
            Utc::now() >= expires_at
        } else {
            false
        }
    }

    /// Set the release address.
    pub fn set_release_address(&mut self, address: Address) -> Result<()> {
        self.release_address = Some(address);
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Mark a participant as having signed.
    pub fn mark_signed(&mut self, participant_id: &str) -> Result<()> {
        let participant = self
            .participants
            .iter_mut()
            .find(|p| p.id == participant_id)
            .ok_or_else(|| EscrowError::Contract("Participant not found".to_string()))?;

        participant.signed = true;
        self.updated_at = Utc::now();
        Ok(())
    }

    /// Check if enough participants have signed for release.
    pub fn has_sufficient_signatures(&self) -> bool {
        let signed_count = self.participants.iter().filter(|p| p.signed).count();
        signed_count >= self.config.threshold
    }

    /// Get the number of additional signatures needed.
    pub fn signatures_needed(&self) -> usize {
        let signed_count = self.participants.iter().filter(|p| p.signed).count();
        self.config.threshold.saturating_sub(signed_count)
    }
}

/// Manager for escrow contracts.
pub struct EscrowManager {
    /// Audit logger.
    audit: AuditLogger,
    /// Condition evaluator.
    evaluator: StandardEvaluator,
    /// Oracle registry.
    oracles: OracleRegistry,
    /// Active contracts.
    contracts: HashMap<EscrowId, EscrowContract>,
}

impl EscrowManager {
    /// Create a new escrow manager.
    pub fn new(audit: AuditLogger) -> Self {
        Self {
            audit,
            evaluator: StandardEvaluator::new(),
            oracles: OracleRegistry::new(),
            contracts: HashMap::new(),
        }
    }

    /// Create a new escrow contract.
    pub fn create_contract(
        &mut self,
        config: EscrowConfig,
        description: Option<String>,
        actor: String,
    ) -> Result<EscrowContract> {
        let contract = EscrowContract::new(config, description);

        self.audit.log(
            AuditEventType::ContractCreated,
            contract.id.clone(),
            actor,
            format!("Escrow contract created: {}", contract.id),
            None,
        )?;

        self.contracts.insert(contract.id.clone(), contract.clone());
        Ok(contract)
    }

    /// Get a contract by ID.
    pub fn get_contract(&self, id: &EscrowId) -> Option<&EscrowContract> {
        self.contracts.get(id)
    }

    /// Get a mutable contract by ID.
    pub fn get_contract_mut(&mut self, id: &EscrowId) -> Option<&mut EscrowContract> {
        self.contracts.get_mut(id)
    }

    /// Add a participant to a contract.
    pub fn add_participant(
        &mut self,
        escrow_id: &EscrowId,
        participant: EscrowParticipant,
        actor: String,
    ) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        contract.add_participant(participant.clone())?;

        self.audit.log(
            AuditEventType::ParticipantJoined,
            escrow_id.clone(),
            actor,
            format!(
                "Participant {} joined as {:?}",
                participant.id, participant.role
            ),
            Some({
                let mut m = HashMap::new();
                m.insert(
                    "participant_id".to_string(),
                    serde_json::json!(participant.id),
                );
                m.insert(
                    "role".to_string(),
                    serde_json::json!(participant.role.to_string()),
                );
                m
            }),
        )?;

        Ok(())
    }

    /// Initialize multisig for a contract.
    pub fn initialize_multisig(&mut self, escrow_id: &EscrowId) -> Result<Address> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        contract.initialize_multisig()?;
        contract.escrow_address()
    }

    /// Fund a contract.
    pub fn fund_contract(
        &mut self,
        escrow_id: &EscrowId,
        output: EscrowOutput,
        actor: String,
    ) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        contract.add_output(output.clone())?;

        self.audit.log(
            AuditEventType::FundsDeposited,
            escrow_id.clone(),
            actor,
            format!("Deposited {} sats to escrow", output.amount),
            Some({
                let mut m = HashMap::new();
                m.insert("amount_sat".to_string(), serde_json::json!(output.amount));
                m.insert(
                    "outpoint".to_string(),
                    serde_json::json!(output.outpoint.to_string()),
                );
                m
            }),
        )?;

        Ok(())
    }

    /// Evaluate conditions for a contract.
    pub async fn evaluate_conditions(&mut self, escrow_id: &EscrowId) -> Result<Vec<ConditionResult>> {
        let contract = self
            .contracts
            .get(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        let mut results = Vec::new();
        for condition in &contract.conditions {
            let result = self.evaluator.evaluate(condition).await?;
            results.push(result);
        }

        Ok(results)
    }

    /// Check if all conditions are satisfied.
    pub async fn conditions_satisfied(&mut self, escrow_id: &EscrowId) -> Result<bool> {
        let results = self.evaluate_conditions(escrow_id).await?;
        Ok(results.iter().all(|r| r.satisfied))
    }

    /// Initiate release of funds.
    pub fn initiate_release(
        &mut self,
        escrow_id: &EscrowId,
        release_address: Address,
        actor: String,
    ) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        if contract.status != EscrowStatus::Funded && contract.status != EscrowStatus::Evaluating {
            return Err(EscrowError::InvalidState(format!(
                "Cannot initiate release from status {:?}",
                contract.status
            )));
        }

        contract.set_release_address(release_address)?;
        contract.status = EscrowStatus::PendingRelease;

        self.audit.log(
            AuditEventType::ReleaseInitiated,
            escrow_id.clone(),
            actor,
            "Release initiated".to_string(),
            None,
        )?;

        Ok(())
    }

    /// Record a signature for release.
    pub fn add_signature(
        &mut self,
        escrow_id: &EscrowId,
        participant_id: &str,
        actor: String,
    ) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        contract.mark_signed(participant_id)?;

        self.audit.log(
            AuditEventType::SignatureAdded,
            escrow_id.clone(),
            actor,
            format!("Signature added by {}", participant_id),
            Some({
                let mut m = HashMap::new();
                m.insert("participant_id".to_string(), serde_json::json!(participant_id));
                m
            }),
        )?;

        Ok(())
    }

    /// Release funds (mark as complete).
    pub fn release_funds(
        &mut self,
        escrow_id: &EscrowId,
        tx: SignedTransaction,
        actor: String,
    ) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        if !contract.has_sufficient_signatures() {
            return Err(EscrowError::Signing(
                "Insufficient signatures for release".to_string(),
            ));
        }

        contract.release_tx = Some(tx);
        contract.status = EscrowStatus::Released;
        contract.updated_at = Utc::now();

        self.audit.log(
            AuditEventType::FundsReleased,
            escrow_id.clone(),
            actor,
            format!("Funds released via tx {}", contract.release_tx.as_ref().unwrap().txid),
            Some({
                let mut m = HashMap::new();
                m.insert(
                    "txid".to_string(),
                    serde_json::json!(contract.release_tx.as_ref().unwrap().txid.to_string()),
                );
                m
            }),
        )?;

        Ok(())
    }

    /// Cancel a contract.
    pub fn cancel_contract(&mut self, escrow_id: &EscrowId, reason: String, actor: String) -> Result<()> {
        let contract = self
            .contracts
            .get_mut(escrow_id)
            .ok_or_else(|| EscrowError::Contract("Contract not found".to_string()))?;

        contract.status = EscrowStatus::Cancelled;
        contract.updated_at = Utc::now();

        self.audit.log(
            AuditEventType::ContractCancelled,
            escrow_id.clone(),
            actor,
            format!("Contract cancelled: {}", reason),
            Some({
                let mut m = HashMap::new();
                m.insert("reason".to_string(), serde_json::json!(reason));
                m
            }),
        )?;

        Ok(())
    }

    /// List all contracts.
    pub fn list_contracts(&self) -> Vec<&EscrowContract> {
        self.contracts.values().collect()
    }

    /// Get audit logger.
    pub fn audit(&self) -> &AuditLogger {
        &self.audit
    }

    /// Export audit log.
    pub fn export_audit(&self) -> Result<String> {
        self.audit.export_json()
    }
}

/// Builder for creating escrow contracts.
pub struct EscrowBuilder {
    config: EscrowConfig,
    participants: Vec<EscrowParticipant>,
    conditions: Vec<ReleaseCondition>,
    description: Option<String>,
    metadata: HashMap<String, serde_json::Value>,
}

impl EscrowBuilder {
    /// Create a new builder.
    pub fn new() -> Self {
        Self {
            config: EscrowConfig::default(),
            participants: Vec::new(),
            conditions: Vec::new(),
            description: None,
            metadata: HashMap::new(),
        }
    }

    /// Set the network.
    pub fn network(mut self, network: Network) -> Self {
        self.config.network = network;
        self
    }

    /// Set the threshold.
    pub fn threshold(mut self, threshold: usize) -> Self {
        self.config.threshold = threshold;
        self.config.total_participants = threshold + 1; // Default to threshold+1 participants
        self
    }

    /// Set total participants.
    pub fn total_participants(mut self, total: usize) -> Self {
        self.config.total_participants = total;
        self
    }

    /// Set expiry time.
    pub fn expires_at(mut self, expires_at: DateTime<Utc>) -> Self {
        self.config.expires_at = Some(expires_at);
        self
    }

    /// Set description.
    pub fn description(mut self, description: impl Into<String>) -> Self {
        self.description = Some(description.into());
        self
    }

    /// Set fee rate.
    pub fn fee_rate(mut self, fee_rate: f32) -> Self {
        self.config.fee_rate = fee_rate;
        self
    }

    /// Add a participant.
    pub fn participant(mut self, participant: EscrowParticipant) -> Self {
        self.participants.push(participant);
        self
    }

    /// Add a condition.
    pub fn condition(mut self, condition: ReleaseCondition) -> Self {
        self.conditions.push(condition);
        self
    }

    /// Add metadata.
    pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }

    /// Build the escrow contract.
    pub fn build(self) -> Result<EscrowContract> {
        let mut contract = EscrowContract::new(self.config, self.description);

        for participant in self.participants {
            contract.add_participant(participant)?;
        }

        for condition in self.conditions {
            contract.add_condition(condition)?;
        }

        if !self.metadata.is_empty() {
            contract.metadata = Some(self.metadata);
        }

        Ok(contract)
    }
}

impl Default for EscrowBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::audit::AuditLogger;
    use crate::conditions::ConditionBuilder;
    use crate::multisig::create_participant;
    use crate::types::EscrowRole;

    #[test]
    fn test_escrow_creation() {
        let config = EscrowConfig::default();
        let contract = EscrowContract::new(config, Some("Test escrow".to_string()));
        
        assert_eq!(contract.status, EscrowStatus::Created);
        assert!(contract.participants.is_empty());
    }

    #[test]
    fn test_add_participants() {
        let config = EscrowConfig::default();
        let mut contract = EscrowContract::new(config, None);

        let (buyer, _) = create_participant(EscrowRole::Buyer, "buyer-1".to_string(), Network::Testnet).unwrap();
        contract.add_participant(buyer).unwrap();

        assert_eq!(contract.participants.len(), 1);
    }

    #[test]
    fn test_escrow_builder() {
        let contract = EscrowBuilder::new()
            .network(Network::Testnet)
            .threshold(2)
            .description("Test escrow")
            .build()
            .unwrap();

        assert_eq!(contract.config.network, Network::Testnet);
        assert_eq!(contract.config.threshold, 2);
    }

    #[tokio::test]
    async fn test_escrow_manager() {
        let audit = AuditLogger::in_memory();
        let mut manager = EscrowManager::new(audit);

        let contract = manager
            .create_contract(
                EscrowConfig::default(),
                Some("Test escrow".to_string()),
                "agent-1".to_string(),
            )
            .unwrap();

        assert!(manager.get_contract(&contract.id).is_some());
    }
}