/*
 copyright: (c) 2013-2018 by Blockstack PBC, a public benefit corporation.

 This file is part of Blockstack.

 Blockstack is free software. You may redistribute or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation, either version 3 of the License or
 (at your option) any later version.

 Blockstack is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY, including without the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with Blockstack. If not, see <http://www.gnu.org/licenses/>.
*/

use std::error;
use std::fmt;

use burnchains::PublicKey;

use deps::bitcoin::blockdata::script::{Script, Instruction, Builder};
use deps::bitcoin::blockdata::opcodes::All as btc_opcodes;

use util::hash::Hash160;

use sha2::Sha256;
use sha2::Digest;

pub mod b58;
pub mod c32;

#[derive(Debug)]
pub enum Error {
    InvalidCrockford32,
    InvalidVersion(u8),
    EmptyData,
    /// Invalid character encountered
    BadByte(u8),
    /// Checksum was not correct (expected, actual)
    BadChecksum(u32, u32),
    /// The length (in bytes) of the object was not correct
    /// Note that if the length is excessively long the provided length may be
    /// an estimate (and the checksum step may be skipped).
    InvalidLength(usize),
    /// Checked data was less than 4 bytes
    TooShort(usize),
    /// Any other error
    Other(String)
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::InvalidCrockford32 => write!(f, "Invalid crockford 32 string"),
            Error::InvalidVersion(ref v) => write!(f, "Invalid version {}", v),
            Error::EmptyData => f.write_str("Empty data"),
            Error::BadByte(b) => write!(f, "invalid base58 character 0x{:x}", b),
            Error::BadChecksum(exp, actual) => write!(f, "base58ck checksum 0x{:x} does not match expected 0x{:x}", actual, exp),
            Error::InvalidLength(ell) => write!(f, "length {} invalid for this base58 type", ell),
            Error::TooShort(_) => write!(f, "base58ck data not even long enough for a checksum"),
            Error::Other(ref s) => f.write_str(s)
        }
    }
}

impl error::Error for Error {
    fn cause(&self) -> Option<&dyn error::Error> { None }
    fn description(&self) -> &'static str {
        match *self {
            Error::InvalidCrockford32 => "Invalid crockford 32 string",
            Error::InvalidVersion(_) => "Invalid version",
            Error::EmptyData => "Empty data",
            Error::BadByte(_) => "invalid b58 character",
            Error::BadChecksum(_, _) => "invalid b58ck checksum",
            Error::InvalidLength(_) => "invalid length for b58 type",
            Error::TooShort(_) => "b58ck data less than 4 bytes",
            Error::Other(_) => "unknown b58 error"
        }
    }
}

#[repr(u8)]
#[derive(Debug, Clone, PartialEq, Eq, Copy, Serialize, Deserialize)]
pub enum AddressHashMode {
    // serialization modes for public keys to addresses.
    // We support four different modes due to legacy compatibility with Stacks v1 addresses:
    SerializeP2PKH = 0x00,          // hash160(public-key), same as bitcoin's p2pkh
    SerializeP2SH = 0x01,           // hash160(multisig-redeem-script), same as bitcoin's multisig p2sh
    SerializeP2WPKH = 0x02,         // hash160(segwit-program-00(p2pkh)), same as bitcoin's p2sh-p2wpkh
    SerializeP2WSH = 0x03,          // hash160(segwit-program-00(public-keys)), same as bitcoin's p2sh-p2wsh
}

/// Given the u8 of an AddressHashMode, deduce the AddressHashNode
pub fn address_serialization_flag(flags: u8) -> AddressHashMode {
    match flags & 0x03 {
        x if x == AddressHashMode::SerializeP2PKH as u8 => AddressHashMode::SerializeP2PKH,
        x if x == AddressHashMode::SerializeP2SH as u8 => AddressHashMode::SerializeP2SH,
        x if x == AddressHashMode::SerializeP2WPKH as u8 => AddressHashMode::SerializeP2WPKH,
        x if x == AddressHashMode::SerializeP2WSH as u8 => AddressHashMode::SerializeP2WSH,
        _ => unreachable!()
    }
}

/// Internally, the Stacks blockchain encodes address the same as Bitcoin
/// single-sig address (p2pkh)
/// Get back the hash of the address
fn to_bits_p2pkh<K: PublicKey>(pubk: &K) -> Hash160 {
    let key_hash = Hash160::from_data(&pubk.to_bytes());
    key_hash
}

/// Internally, the Stacks blockchain encodes address the same as Bitcoin
/// multi-sig address (p2sh)
fn to_bits_p2sh<K: PublicKey>(num_sigs: usize, pubkeys: &Vec<K>) -> Hash160 {
    let mut bldr = Builder::new();
    bldr = bldr.push_int(num_sigs as i64);
    for pubk in pubkeys {
        bldr = bldr.push_slice(&pubk.to_bytes());
    }
    bldr = bldr.push_int(pubkeys.len() as i64);
    bldr = bldr.push_opcode(btc_opcodes::OP_CHECKMULTISIG);
    
    let script = bldr.into_script();
    let script_hash = Hash160::from_data(&script.as_bytes());
    script_hash
}

/// Internally, the Stacks blockchain encodes address the same as Bitcoin
/// single-sig address over p2sh (p2h-p2wpkh)
fn to_bits_p2sh_p2wpkh<K: PublicKey>(pubk: &K) -> Hash160 {
    let key_hash = Hash160::from_data(&pubk.to_bytes());

    let bldr = Builder::new()
        .push_int(0)
        .push_slice(key_hash.as_bytes());

    let script = bldr.into_script();
    let script_hash = Hash160::from_data(&script.as_bytes());
    script_hash
}

/// Internally, the Stacks blockchain encodes address the same as Bitcoin
/// multisig address over p2sh (p2sh-p2wsh)
fn to_bits_p2sh_p2wsh<K: PublicKey>(num_sigs: usize, pubkeys: &Vec<K>) -> Hash160 {
    let mut bldr = Builder::new();
    bldr = bldr.push_int(num_sigs as i64);
    for pubk in pubkeys {
        bldr = bldr.push_slice(&pubk.to_bytes());
    }
    bldr = bldr.push_int(pubkeys.len() as i64);
    bldr = bldr.push_opcode(btc_opcodes::OP_CHECKMULTISIG);

    let mut digest = Sha256::new();
    let mut d = [0u8; 32];

    digest.input(bldr.into_script().as_bytes());
    d.copy_from_slice(digest.result().as_slice());

    let ws = Builder::new().push_int(0).push_slice(&d).into_script();
    let ws_hash = Hash160::from_data(&ws.as_bytes());
    ws_hash
}

/// Convert a number of required signatures and a list of public keys into a byte-vec to hash to an
/// address.  Validity of the hash_flag vis a vis the num_sigs and pubkeys will _NOT_ be checked.
/// This is a low-level method.  Consider using StacksAdress::from_public_keys() if you can.
pub fn public_keys_to_address_hash<K: PublicKey>(hash_flag: &AddressHashMode, num_sigs: usize, pubkeys: &Vec<K>) -> Hash160 {
    match *hash_flag {
        AddressHashMode::SerializeP2PKH => to_bits_p2pkh(&pubkeys[0]),
        AddressHashMode::SerializeP2SH => to_bits_p2sh(num_sigs, pubkeys),
        AddressHashMode::SerializeP2WPKH => to_bits_p2sh_p2wpkh(&pubkeys[0]),
        AddressHashMode::SerializeP2WSH => to_bits_p2sh_p2wsh(num_sigs, pubkeys)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use burnchains::PublicKey;
    use util::log;
    use util::hash::*;
    use util::secp256k1::Secp256k1PublicKey as PubKey;

    struct PubkeyFixture {
        keys: Vec<PubKey>,
        num_required: usize,
        segwit: bool,
        result: Vec<u8>
    }

    #[test]
    fn test_public_keys_to_address_hash() {
        let pubkey_fixtures = vec![
            PubkeyFixture {
                // script pubkey for p2pkh
                keys: vec![
                    PubKey::from_hex("040fadbbcea0ff3b05f03195b41cd991d7a0af8bd38559943aec99cbdaf0b22cc806b9a4f07579934774cc0c155e781d45c989f94336765e88a66d91cfb9f060b0").unwrap(),
                ],
                num_required: 1,
                segwit: false,
                result: hex_bytes("395f3643cea07ec4eec73b4d9a973dcce56b9bf1").unwrap().to_vec()
            },
            PubkeyFixture {
                // script pubkey for multisig p2sh
                keys: vec![
                    PubKey::from_hex("040fadbbcea0ff3b05f03195b41cd991d7a0af8bd38559943aec99cbdaf0b22cc806b9a4f07579934774cc0c155e781d45c989f94336765e88a66d91cfb9f060b0").unwrap(),
                    PubKey::from_hex("04c77f262dda02580d65c9069a8a34c56bd77325bba4110b693b90216f5a3edc0bebc8ce28d61aa86b414aa91ecb29823b11aeed06098fcd97fee4bc73d54b1e96").unwrap(),
                ],
                num_required: 2,
                segwit: false,
                result: hex_bytes("fd3a5e9f5ba311ce6122765f0af8da7488e25d3a").unwrap().to_vec(),
            },
            PubkeyFixture {
                // script pubkey for p2sh-p2wpkh
                keys: vec![
                    PubKey::from_hex("020fadbbcea0ff3b05f03195b41cd991d7a0af8bd38559943aec99cbdaf0b22cc8").unwrap(),
                ],
                num_required: 1,
                segwit: true,
                result: hex_bytes("0ac7ad046fe22c794dd923b3be14b2e668e50c42").unwrap().to_vec(),
            },
            PubkeyFixture {
                // script pubkey for multisig p2sh-p2wsh
                keys: vec![
                    PubKey::from_hex("020fadbbcea0ff3b05f03195b41cd991d7a0af8bd38559943aec99cbdaf0b22cc8").unwrap(),
                    PubKey::from_hex("02c77f262dda02580d65c9069a8a34c56bd77325bba4110b693b90216f5a3edc0b").unwrap(),
                ],
                num_required: 2,
                segwit: true,
                result: hex_bytes("3e02fa83ac2fae11fd6703b91e7c94ad393052e2").unwrap().to_vec(),
            },
        ];

        for pubkey_fixture in pubkey_fixtures {
            let hash_mode = 
                if !pubkey_fixture.segwit {
                    if pubkey_fixture.num_required == 1 {
                        AddressHashMode::SerializeP2PKH
                    }
                    else {
                        AddressHashMode::SerializeP2SH
                    }
                }
                else {
                    if pubkey_fixture.num_required == 1 {
                        AddressHashMode::SerializeP2WPKH
                    }
                    else {
                        AddressHashMode::SerializeP2WSH
                    }
                };

            let result_hash = public_keys_to_address_hash(&hash_mode, pubkey_fixture.num_required, &pubkey_fixture.keys);
            let result = result_hash.as_bytes().to_vec();

            assert_eq!(result, pubkey_fixture.result);
        }
    }
}
