nil-zonefile 0.3.0

use serde::{Deserialize, Serialize};
use nil_slip44::{Coin, Symbol};
#[allow(unused)]
use log;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum WalletType {
    // For on-chain addresses
    OnChain(Coin),
    // For Lightning Network addresses
    Lightning,
}

// Custom serialization to use coin IDs
impl Serialize for WalletType {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            WalletType::OnChain(coin) => serializer.serialize_u32(coin.id()),
            WalletType::Lightning => serializer.serialize_str("lightning"),
        }
    }
}

// Custom deserialization to handle both coin IDs and "lightning"
impl<'de> Deserialize<'de> for WalletType {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;
        
        // Try to deserialize as either a number or string
        #[derive(Deserialize)]
        #[serde(untagged)]
        enum IdOrString {
            Id(u32),
            String(String),
        }

        match IdOrString::deserialize(deserializer)? {
            IdOrString::Id(id) => {
                // Convert coin ID back to Coin
                if let Ok(coin) = Coin::try_from(id) {
                    Ok(WalletType::OnChain(coin))
                } else {
                    Err(Error::custom(format!("Invalid coin ID: {}", id)))
                }
            }
            IdOrString::String(s) => {
                // Check for Lightning first
                if s.to_lowercase() == "lightning" {
                    return Ok(WalletType::Lightning);
                }
                
                // Try parsing string as a number
                if let Ok(id) = s.parse::<u32>() {
                    if let Ok(coin) = Coin::try_from(id) {
                        return Ok(WalletType::OnChain(coin));
                    }
                }
                
                // Try parsing as Symbol
                if let Ok(symbol) = s.parse::<Symbol>() {
                    return Ok(WalletType::OnChain(Coin::from(symbol)));
                }
                
                Err(Error::custom(format!("Invalid wallet type string: {}", s)))
            }
        }
    }
}

impl WalletType {
    /// Create a Bitcoin wallet type
    pub fn bitcoin() -> Self {
        WalletType::OnChain(Coin::Bitcoin)
    }

    /// Create a Stacks wallet type
    pub fn stacks() -> Self {
        WalletType::OnChain(Coin::Stacks)
    }

    /// Create a Solana wallet type
    pub fn solana() -> Self {
        WalletType::OnChain(Coin::Solana)
    }

    /// Create an Ethereum wallet type
    pub fn ethereum() -> Self {
        WalletType::OnChain(Coin::Ethereum)
    }

    /// Create a Lightning Network wallet type
    pub fn lightning() -> Self {
        WalletType::Lightning
    }

    pub fn from_str(s: &str) -> Option<Self> {
        // Try parsing as SLIP44 symbol first for on-chain addresses
        if let Ok(symbol) = s.parse::<Symbol>() {
            return Some(WalletType::OnChain(Coin::from(symbol)));
        }

        // Check for Lightning
        match s.to_lowercase().as_str() {
            "lightning" => Some(WalletType::Lightning),
            _ => None,
        }
    }

    pub fn to_string(&self) -> String {
        match self {
            WalletType::OnChain(coin) => coin.to_string(),
            WalletType::Lightning => "lightning".to_string(),
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Wallet {
    pub address: String,
    pub coin_type: WalletType,
}

#[cfg(test)]
mod tests {
    use serde_json::json;
    use super::*;
    use log::{debug};

    #[test]
    fn test_wallet_type_from_str() {
        // Test on-chain addresses
        assert_eq!(WalletType::from_str("BTC"), Some(WalletType::OnChain(Coin::Bitcoin)));
        assert_eq!(WalletType::from_str("STX"), Some(WalletType::OnChain(Coin::Stacks)));
        assert_eq!(WalletType::from_str("SOL"), Some(WalletType::OnChain(Coin::Solana)));
        assert_eq!(WalletType::from_str("ETH"), Some(WalletType::OnChain(Coin::Ethereum)));
        
        // Test Lightning Network
        assert_eq!(WalletType::from_str("lightning"), Some(WalletType::Lightning));
        assert_eq!(WalletType::from_str("LIGHTNING"), Some(WalletType::Lightning));
        
        // Test invalid input
        assert_eq!(WalletType::from_str("INVALID"), None);
    }

    #[test]
    fn test_wallet_type_to_string() {
        assert_eq!(WalletType::OnChain(Coin::Bitcoin).to_string(), "Bitcoin");
        assert_eq!(WalletType::OnChain(Coin::Solana).to_string(), "Solana");
        assert_eq!(WalletType::Lightning.to_string(), "lightning");
    }

    #[test]
    fn test_wallet_serialization() {
        let wallet = Wallet {
            address: "bc1qxxx...".to_string(),
            coin_type: WalletType::OnChain(Coin::Bitcoin),
        };
        
        let json = serde_json::to_string(&wallet).unwrap();
        debug!("Bitcoin wallet JSON: {}", json);
        let deserialized: Wallet = serde_json::from_str(&json).unwrap();
        
        assert_eq!(wallet.address, deserialized.address);
        assert_eq!(wallet.coin_type, deserialized.coin_type);

        let lightning_wallet = Wallet {
            address: "username@domain.com".to_string(),
            coin_type: WalletType::Lightning,
        };
        
        let json = serde_json::to_string(&lightning_wallet).unwrap();
        debug!("Lightning wallet JSON: {}", json);
        let deserialized: Wallet = serde_json::from_str(&json).unwrap();
        
        assert_eq!(lightning_wallet.address, deserialized.address);
        assert_eq!(lightning_wallet.coin_type, deserialized.coin_type);
    }

    #[test]
    fn test_coin_id_serialization() {
        // Test that Bitcoin serializes to ID 0
        let btc_wallet = WalletType::OnChain(Coin::Bitcoin);
        let json = serde_json::to_string(&btc_wallet).unwrap();
        assert_eq!(json, "0");

        // Test that Stacks serializes to ID 5757
        let stx_wallet = WalletType::OnChain(Coin::Stacks);
        let json = serde_json::to_string(&stx_wallet).unwrap();
        assert_eq!(json, "5757");

        // Test that Solana serializes to ID 501
        let sol_wallet = WalletType::OnChain(Coin::Solana);
        let json = serde_json::to_string(&sol_wallet).unwrap();
        assert_eq!(json, "501");

        // Test that Ethereum serializes to ID 60
        let eth_wallet = WalletType::OnChain(Coin::Ethereum);
        let json = serde_json::to_string(&eth_wallet).unwrap();
        assert_eq!(json, "60");

        // Test deserialization from IDs
        let deserialized: WalletType = serde_json::from_str("0").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Bitcoin));

        let deserialized: WalletType = serde_json::from_str("5757").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Stacks));

        let deserialized: WalletType = serde_json::from_str("501").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Solana));

        let deserialized: WalletType = serde_json::from_str("60").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Ethereum));

        // Test deserialization from string IDs
        let deserialized: WalletType = serde_json::from_str("\"0\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Bitcoin));

        let deserialized: WalletType = serde_json::from_str("\"5757\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Stacks));

        let deserialized: WalletType = serde_json::from_str("\"501\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Solana));

        let deserialized: WalletType = serde_json::from_str("\"60\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Ethereum));

        // Test deserialization from symbols
        let deserialized: WalletType = serde_json::from_str("\"BTC\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Bitcoin));

        let deserialized: WalletType = serde_json::from_str("\"STX\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Stacks));

        let deserialized: WalletType = serde_json::from_str("\"SOL\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Solana));

        let deserialized: WalletType = serde_json::from_str("\"ETH\"").unwrap();
        assert_eq!(deserialized, WalletType::OnChain(Coin::Ethereum));
    }

    #[test]
    fn test_onchain_numeric_serialization() {
        // Test that OnChain types serialize as numbers, not strings
        let btc_wallet = WalletType::OnChain(Coin::Bitcoin);
        let stx_wallet = WalletType::OnChain(Coin::Stacks);
        let sol_wallet = WalletType::OnChain(Coin::Solana);

        // Verify serialization produces numbers
        assert_eq!(serde_json::to_value(&btc_wallet).unwrap(), json!(0));
        assert_eq!(serde_json::to_value(&stx_wallet).unwrap(), json!(5757));
        assert_eq!(serde_json::to_value(&sol_wallet).unwrap(), json!(501));

        // Verify serialization does NOT produce strings
        assert_ne!(serde_json::to_value(&btc_wallet).unwrap(), json!("0"));
        assert_ne!(serde_json::to_value(&stx_wallet).unwrap(), json!("5757"));
        assert_ne!(serde_json::to_value(&sol_wallet).unwrap(), json!("501"));

        // Verify the actual JSON string output
        assert_eq!(serde_json::to_string(&btc_wallet).unwrap(), "0");
        assert_eq!(serde_json::to_string(&stx_wallet).unwrap(), "5757");
        assert_eq!(serde_json::to_string(&sol_wallet).unwrap(), "501");
        assert_ne!(serde_json::to_string(&btc_wallet).unwrap(), "\"0\"");
        assert_ne!(serde_json::to_string(&stx_wallet).unwrap(), "\"5757\"");
        assert_ne!(serde_json::to_string(&sol_wallet).unwrap(), "\"501\"");
    }

    #[test]
    fn test_convenience_constructors() {
        // Test that convenience constructors create the correct types
        assert_eq!(WalletType::bitcoin(), WalletType::OnChain(Coin::Bitcoin));
        assert_eq!(WalletType::stacks(), WalletType::OnChain(Coin::Stacks));
        assert_eq!(WalletType::solana(), WalletType::OnChain(Coin::Solana));
        assert_eq!(WalletType::ethereum(), WalletType::OnChain(Coin::Ethereum));
        assert_eq!(WalletType::lightning(), WalletType::Lightning);

        // Test that they serialize to the correct IDs
        assert_eq!(serde_json::to_string(&WalletType::bitcoin()).unwrap(), "0");
        assert_eq!(serde_json::to_string(&WalletType::stacks()).unwrap(), "5757");
        assert_eq!(serde_json::to_string(&WalletType::solana()).unwrap(), "501");
        assert_eq!(serde_json::to_string(&WalletType::ethereum()).unwrap(), "60");
        assert_eq!(serde_json::to_string(&WalletType::lightning()).unwrap(), "\"lightning\"");
    }
}