pub mod errors;
use bitcoincash::Network;
pub use errors::{DecodingError, EncodingError};

// Prefixes
const MAINNET_PREFIX: &str = "nexa";
const TESTNET_PREFIX: &str = "nexatest";
const REGNET_PREFIX: &str = "nexareg";

// The cashaddr character set for encoding
const CHARSET: &[u8; 32] = b"qpzry9x8gf2tvdw0s3jn54khce6mua7l";

// The cashaddr character set for decoding
#[rustfmt::skip]
const CHARSET_REV: [Option<u8>; 128] = [
    None,     None,     None,     None,     None,     None,     None,     None,
    None,     None,     None,     None,     None,     None,     None,     None,
    None,     None,     None,     None,     None,     None,     None,     None,
    None,     None,     None,     None,     None,     None,     None,     None,
    None,     None,     None,     None,     None,     None,     None,     None,
    None,     None,     None,     None,     None,     None,     None,     None,
    Some(15), None,     Some(10), Some(17), Some(21), Some(20), Some(26), Some(30),
    Some(7),  Some(5),  None,     None,     None,     None,     None,     None,
    None,     Some(29), None,     Some(24), Some(13), Some(25), Some(9),  Some(8),
    Some(23), None,     Some(18), Some(22), Some(31), Some(27), Some(19), None,
    Some(1),  Some(0),  Some(3),  Some(16), Some(11), Some(28), Some(12), Some(14),
    Some(6),  Some(4),  Some(2),  None,     None,     None,     None,     None,
    None,     Some(29),  None,    Some(24), Some(13), Some(25), Some(9),  Some(8),
    Some(23), None,     Some(18), Some(22), Some(31), Some(27), Some(19), None,
    Some(1),  Some(0),  Some(3),  Some(16), Some(11), Some(28), Some(12), Some(14),
    Some(6),  Some(4),  Some(2),  None,     None,     None,     None,     None,
];

// Version byte flags
#[allow(dead_code)]
#[rustfmt::skip]
pub mod version_byte_flags {
    pub const TYPE_MASK: u8            = 0b01111000;
    pub const SIZE_MASK: u8            = 0b00000111;

    pub const TYPE_P2PKH: u8           = 0b00000;
    pub const TYPE_P2SH: u8            = 0b01000;

    // Not masked, use full byte.
    pub const TYPE_GROUP: u8           = 88;
    pub const TYPE_SCRIPT_TEMPLATE: u8 = 152;

    pub const SIZE_160: u8 = 0x00;
    pub const SIZE_192: u8 = 0x01;
    pub const SIZE_224: u8 = 0x02;
    pub const SIZE_256: u8 = 0x03;
    pub const SIZE_320: u8 = 0x04;
    pub const SIZE_384: u8 = 0x05;
    pub const SIZE_448: u8 = 0x06;
    pub const SIZE_512: u8 = 0x07;
}

// https://github.com/Bitcoin-ABC/bitcoin-abc/blob/2804a49bfc0764ba02ce2999809c52b3b9bb501e/src/cashaddr.cpp#L42
fn polymod(v: &[u8]) -> u64 {
    let mut c: u64 = 1;
    for d in v.iter() {
        let c0: u8 = (c >> 35) as u8;
        c = ((c & 0x0007_ffff_ffff) << 5) ^ u64::from(*d);
        if c0 & 0x01 != 0 {
            c ^= 0x0098_f2bc_8e61;
        }
        if c0 & 0x02 != 0 {
            c ^= 0x0079_b76d_99e2;
        }
        if c0 & 0x04 != 0 {
            c ^= 0x00f3_3e5f_b3c4;
        }
        if c0 & 0x08 != 0 {
            c ^= 0x00ae_2eab_e2a8;
        }
        if c0 & 0x10 != 0 {
            c ^= 0x001e_4f43_e470;
        }
    }
    c ^ 1
}

// Expand the address prefix for the checksum operation.
fn expand_prefix(prefix: &str) -> Vec<u8> {
    let mut ret: Vec<u8> = prefix.chars().map(|c| (c as u8) & 0x1f).collect();
    ret.push(0);
    ret
}

fn convert_bits(data: &[u8], inbits: u8, outbits: u8, pad: bool) -> Vec<u8> {
    assert!(inbits <= 8 && outbits <= 8);
    let num_bytes = (data.len() * inbits as usize).div_ceil(outbits as usize);
    let mut ret = Vec::with_capacity(num_bytes);
    let mut acc: u16 = 0; // accumulator of bits
    let mut num: u8 = 0; // num bits in acc
    let groupmask = (1 << outbits) - 1;
    for d in data.iter() {
        // We push each input chunk into a 16-bit accumulator
        acc = (acc << inbits) | u16::from(*d);
        num += inbits;
        // Then we extract all the output groups we can
        while num > outbits {
            ret.push((acc >> (num - outbits)) as u8);
            acc &= !(groupmask << (num - outbits));
            num -= outbits;
        }
    }
    if pad {
        // If there's some bits left, pad and add it
        if num > 0 {
            ret.push((acc << (outbits - num)) as u8);
        }
    } else {
        // If there's some bits left, figure out if we need to remove padding and add it
        let padding = (data.len() * inbits as usize) % outbits as usize;
        if num as usize > padding {
            ret.push((acc >> padding) as u8);
        }
    }
    ret
}

fn get_version_type(version_byte: u8) -> u8 {
    match version_byte {
        version_byte_flags::TYPE_GROUP | version_byte_flags::TYPE_SCRIPT_TEMPLATE => version_byte,
        _ => version_byte & version_byte_flags::TYPE_MASK,
    }
}

pub fn encode(
    raw: &[u8],
    hash_flag: u8, // see version_byte_flags
    network: Network,
) -> Result<String, EncodingError> {
    // Calculate version byte
    let length = raw.len();
    let version_byte: u8 = if hash_flag != version_byte_flags::TYPE_GROUP
        && hash_flag != version_byte_flags::TYPE_SCRIPT_TEMPLATE
    {
        hash_flag
            | match length {
                20 => version_byte_flags::SIZE_160,
                24 => version_byte_flags::SIZE_192,
                28 => version_byte_flags::SIZE_224,
                32 => version_byte_flags::SIZE_256,
                40 => version_byte_flags::SIZE_320,
                48 => version_byte_flags::SIZE_384,
                56 => version_byte_flags::SIZE_448,
                64 => version_byte_flags::SIZE_512,
                _ => return Err(EncodingError(length)),
            }
    } else {
        hash_flag
    };

    // Get prefix
    let prefix = match network {
        Network::Bitcoin => MAINNET_PREFIX,
        Network::Testnet => TESTNET_PREFIX,
        Network::Regtest => REGNET_PREFIX,
        _ => {
            // Network not supported, default to testnet.
            TESTNET_PREFIX
        }
    };

    let raw = if version_byte == version_byte_flags::TYPE_SCRIPT_TEMPLATE {
        // Pay to script template wraps the payload network serialization
        bitcoincash::consensus::serialize::<Vec<u8>>(&raw.to_vec())
    } else {
        raw.to_vec()
    };

    // Convert payload to 5 bit array
    let mut payload = Vec::with_capacity(1 + raw.len());
    payload.push(version_byte);
    payload.extend(raw);
    let payload_5_bits = convert_bits(&payload, 8, 5, true);

    // Construct payload string using CHARSET
    let payload_str: String = payload_5_bits
        .iter()
        .map(|b| CHARSET[*b as usize] as char)
        .collect();

    // Create checksum
    let expanded_prefix = expand_prefix(prefix);
    let checksum_input = [&expanded_prefix[..], &payload_5_bits, &[0; 8][..]].concat();
    let checksum = polymod(&checksum_input);

    // Convert checksum to string
    let checksum_str: String = (0..8)
        .rev()
        .map(|i| CHARSET[((checksum >> (i * 5)) & 31) as usize] as char)
        .collect();

    // Concatentate all parts
    let cashaddr = [prefix, ":", &payload_str, &checksum_str].concat();
    Ok(cashaddr)
}

pub fn decode(addr_str: &str) -> Result<(Vec<u8>, u8, Network), DecodingError> {
    // Delimit and extract prefix
    let parts: Vec<&str> = addr_str.split(':').collect();
    let (prefix, payload_str) = if parts.len() == 2 {
        (parts[0], parts[1])
    } else {
        (MAINNET_PREFIX, parts[0])
    };

    // Match network
    let network = match prefix {
        MAINNET_PREFIX => Network::Bitcoin,
        TESTNET_PREFIX => Network::Testnet,
        REGNET_PREFIX => Network::Regtest,
        _ => return Err(DecodingError::InvalidPrefix(prefix.to_string())),
    };

    // Do some sanity checks on the string
    let mut payload_chars = payload_str.chars();
    if let Some(first_char) = payload_chars.next() {
        if first_char.is_lowercase() {
            if payload_chars.any(|c| c.is_uppercase()) {
                return Err(DecodingError::MixedCase);
            }
        } else if payload_chars.any(|c| c.is_lowercase()) {
            return Err(DecodingError::MixedCase);
        }
    } else {
        return Err(DecodingError::InvalidLength(0));
    }

    // Decode payload to 5 bit array
    let payload_chars = payload_str.chars(); // Reintialize iterator here
    let payload_5_bits: Result<Vec<u8>, DecodingError> = payload_chars
        .map(|c| {
            let i = c as usize;
            if let Some(Some(d)) = CHARSET_REV.get(i) {
                Ok(*d)
            } else {
                Err(DecodingError::InvalidChar(c))
            }
        })
        .collect();
    let payload_5_bits = payload_5_bits?;

    // Verify the checksum
    let checksum = polymod(&[&expand_prefix(prefix), &payload_5_bits[..]].concat());
    if checksum != 0 {
        return Err(DecodingError::ChecksumFailed(checksum));
    }

    // Convert from 5 bit array to byte array
    let len_5_bit = payload_5_bits.len();
    let payload = convert_bits(&payload_5_bits[..(len_5_bit - 8)], 5, 8, false);

    // Verify the version byte
    let version = payload[0];

    // Check length
    let body = &payload[1..];
    let body_len = body.len();

    // Extract the hash type and return
    let version_type = get_version_type(version);

    if version_type == version_byte_flags::TYPE_P2PKH
        || version_type == version_byte_flags::TYPE_P2SH
    {
        let payload_size = version & version_byte_flags::SIZE_MASK;
        if (payload_size == version_byte_flags::SIZE_160 && body_len != 20)
            || (payload_size == version_byte_flags::SIZE_192 && body_len != 24)
            || (payload_size == version_byte_flags::SIZE_224 && body_len != 28)
            || (payload_size == version_byte_flags::SIZE_256 && body_len != 32)
            || (payload_size == version_byte_flags::SIZE_320 && body_len != 40)
            || (payload_size == version_byte_flags::SIZE_384 && body_len != 48)
            || (payload_size == version_byte_flags::SIZE_448 && body_len != 56)
            || (payload_size == version_byte_flags::SIZE_512 && body_len != 64)
        {
            return Err(DecodingError::InvalidLength(body_len));
        }
    }

    if version_type != version_byte_flags::TYPE_P2PKH
        && version_type != version_byte_flags::TYPE_P2SH
        && version_type != version_byte_flags::TYPE_GROUP
        && version_type != version_byte_flags::TYPE_SCRIPT_TEMPLATE
    {
        return Err(DecodingError::InvalidVersion(version));
    }

    if version_type == version_byte_flags::TYPE_SCRIPT_TEMPLATE {
        // The payload is wrapped in network serialization

        let inner: Vec<u8> = match bitcoincash::consensus::deserialize::<Vec<u8>>(body) {
            Ok(i) => i,
            Err(e) => return Err(DecodingError::Other(format!("Decode P2PKT failed {}", e))),
        };
        return Ok((inner, version_type, network));
    }

    Ok((body.to_vec(), version_type, network))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn p2pkt() {
        // Pay-to-public-key-template (P2PKT)
        // https://gitlab.com/nexa/nexa/-/blob/b9fd30c28081b5ae9df99468fee3c85cd52bbbd0/src/test/dstencode_tests.cpp#L120

        let address = "nexa:nqtsq5g5w6syq5aa5z5ghkj3w7ux59wrk204txrn64e2gs92";
        let (payload, address_type, network) = decode(address).unwrap();
        #[rustfmt::skip]
        assert_eq!(
            vec![
                0,  // No group (OP_0)
                81, // Well known template 1 (1 encoded as OP_TRUE)
                20, // Push 20 bytes
                // 20 byte hash
                118, 160, 64, 83, 189, 160, 168, 139, 218, 81, 119, 184, 106, 21, 195, 178, 159,
                85, 152, 115
            ],
            payload
        );
        assert_eq!(address_type, version_byte_flags::TYPE_SCRIPT_TEMPLATE,);
        assert_eq!(network, Network::Bitcoin);
        assert_eq!(address, encode(&payload, address_type, network).unwrap());
    }

    #[test]
    fn gp2pkt() {
        // Grouped, unspecified token amount, pay-to-public-key-template (GP2PKT)
        // https://gitlab.com/nexa/nexa/-/blob/b9fd30c28081b5ae9df99468fee3c85cd52bbbd0/src/test/dstencode_tests.cpp#L133

        let address = "nexa:nqazqy3uqsp4j0zyyufqzy65qc2u9vvm2jthyqgzqvzq2ps8pqys5zcvqgqqq5g5w6syq5aa5z5ghkj3w7ux59wrk204txrn92xzqzsu";
        let (payload, address_type, network) = decode(address).unwrap();

        #[rustfmt::skip]
        assert_eq!(
            payload,
            vec![
                32, // Push group ID (32 bytes)
                // Group ID
                18, 60, 4, 3, 89, 60, 68, 39, 18, 1, 19, 84, 6, 21, 194, 177, 155, 84, 151, 114, 1,
                2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
                2, 0, 0, // Unspecified token amount (Push 2, OP_0, OP_0)
                81, // Well known template 1 (1 encoded as OP_TRUE)
                20, // Push 20 bytes (PKH)
                // PKH
                118, 160, 64, 83, 189, 160, 168, 139, 218, 81, 119, 184, 106, 21, 195, 178, 159,
                85, 152, 115
            ],
        );

        assert_eq!(address_type, version_byte_flags::TYPE_SCRIPT_TEMPLATE,);
        assert_eq!(network, Network::Bitcoin);
        assert_eq!(address, encode(&payload, address_type, network).unwrap());
    }

    #[test]
    fn p2cat() {
        // Ungrouped Pay-To-Contract-Args-Template (P2CAT) with no args
        // https://gitlab.com/nexa/nexa/-/blob/b9fd30c28081b5ae9df99468fee3c85cd52bbbd0/src/test/dstencode_tests.cpp#L172
        let address = "nexa:nqtsq9rk5pq980dq4z9a55thhp4ptsajna2esucqqj42vk56";
        let (payload, address_type, network) = decode(address).unwrap();

        #[rustfmt::skip]
        assert_eq!(
            payload,
            vec![
                0, // No group (OP_0)
                20, // Push 20 bytes (PKH)
                // PKH
                118, 160, 64, 83, 189, 160, 168, 139, 218, 81, 119, 184, 106, 21, 195, 178, 159,
                85, 152, 115,
                // ???
                0
            ]
        );
        assert_eq!(address_type, version_byte_flags::TYPE_SCRIPT_TEMPLATE,);
        assert_eq!(network, Network::Bitcoin);
        assert_eq!(address, encode(&payload, address_type, network).unwrap());
    }

    #[test]
    fn gp2cat() {
        // Grouped, Pay-To-Contract-Args-Template (GP2CAT) with args
        let address = "nexa:np9sq9rk5pq980dq4z9a55thhp4ptsajna2esuc5zg7qgq6e83zzwyspzd2qv9wzkxd4f9mjzg7qgq6e83zzwyspzd2qv9wzkxd4f9mjqypqxpq9qcrsszg2pvxq3p2q9kaa";
        let (payload, address_type, network) = decode(address).unwrap();
        #[rustfmt::skip]
        assert_eq!(
            payload,
            vec![
                0, // ???
                20, // Push 20 bytes
                // Group hash
                118, 160, 64, 83, 189, 160, 168, 139, 218, 81,
                119, 184, 106, 21, 195, 178, 159, 85, 152, 115,
                20, // Push 20 bytes (argument hash)
                // Arguments hash
                18, 60, 4, 3, 89, 60, 68, 39, 18, 1, 19, 84, 6, 21, 194, 177, 155, 84, 151, 114,
                18, // Push 18 bytes
                // Push visible argument
                60, 4, 3, 89, 60, 68, 39, 18, 1, 19, 84, 6, 21, 194, 177, 155, 84, 151,
                // ????
                114, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12
            ]
        );
        assert_eq!(address_type, version_byte_flags::TYPE_SCRIPT_TEMPLATE,);
        assert_eq!(network, Network::Bitcoin);
        assert_eq!(address, encode(&payload, address_type, network).unwrap());
    }

    #[test]
    fn issue211() {
        // issue was that rostrum used script encoding rather than network encoding
        // on the outer wrapper of p2st.
        // this was first triggered on a large payload.

        let addr = "nexa:npnsq99lyhlv4tk65kvyr60yu9vy503uxm5r4wcqyypmppfe5e8h3j4j28eldtqgqtr3v4t5eazvj96akzl7x3635kvgq9gpvs2p053mapa874z8jvjt5xcxwt9desksjm33fvsu99xrkstkqwc6u7kx9s884za33nhfqksqrpjqjvv5";
        let (payload, address_type, network) = decode(addr).unwrap();

        assert_eq!(
            "0014bf25fecaaedaa59841e9e4e1584a3e3c36e83abb002103b08539a64f78cab251f3f6ac0802c7165574cf44c9175db0bfe34751a598801501641417d23be87a7f54479324ba1b0672cadcc2d096e314b21c294c3b417603b1ae7ac62c0e7a8bb18cee905a00",
            &hex::encode(&payload));
        assert_eq!(address_type, version_byte_flags::TYPE_SCRIPT_TEMPLATE);
        assert_eq!(Network::Bitcoin, network);
        assert_eq!(addr, &encode(&payload, address_type, network).unwrap());
    }
}
