use ghost_traits::hashing::GhostHasher;
use sp_std::vec::Vec;

const MAX_INPUTS: usize = 3;
const MAX_OUTPUTS: usize = 4;
const MAX_SCRIPT_LEN: usize = 34;
const MAX_WITNESS_STACK_DEPTH: usize = 3;
const MAX_WITNESS_ITEMS_LEN: usize = 73;

pub const MAX_TX_LEN: usize = 512;

pub struct TxIn<H: GhostHasher> {
    pub prev_txid: H::Hash,
    pub vout: u32,
    pub script_sig: Vec<u8>,
    pub sequence: u32,
}

pub struct TxOut {
    pub value: u64,
    pub script_pubkey: Vec<u8>,
}

pub struct Witness {
    pub stack: Vec<Vec<u8>>,
}

pub enum Transaction<H: GhostHasher> {
    Legacy {
        version: i32,
        inputs: Vec<TxIn<H>>,
        outputs: Vec<TxOut>,
        lock_time: u32,
    },
    SegWit {
        version: i32,
        inputs: Vec<TxIn<H>>,
        outputs: Vec<TxOut>,
        witnesses: Vec<Witness>,
        lock_time: u32,
    },
}

pub struct Header<H: GhostHasher> {
    pub merkle_root: H::Hash,
    pub tx_count: u32,
    pub tree_hashes: Vec<H::Hash>,
    pub flags: Vec<u8>,
}

impl<H: GhostHasher> Header<H> {
    fn calc_tree_width(&self, h: u32) -> u32 {
        (self.tx_count + (1 << h) - 1) >> h
    }

    fn traverse(
        &self,
        height: u32,
        pos: u32,
        bits_used: &mut u32,
        hash_used: &mut u32,
        matched_index: &mut Option<usize>,
        target_txid: H::Hash,
    ) -> Result<H::Hash, ParseError> {
        let byte_idx = (*bits_used / 8) as usize;
        let bit_idx = (*bits_used % 8) as u8;

        if byte_idx >= self.flags.len() {
            return Err(ParseError::InvalidVarint(VarintError::BufferTooShort));
        }

        let parent_of_match = ((self.flags[byte_idx] >> bit_idx) & 1) == 1;
        *bits_used += 1;

        if height == 0 || !parent_of_match {
            if *hash_used as usize >= self.tree_hashes.len() {
                return Err(ParseError::InvalidLength);
            }

            let curr_hash = self.tree_hashes[*hash_used as usize];
            *hash_used += 1;

            if height == 0 && parent_of_match && curr_hash == target_txid {
                *matched_index = Some(pos as usize);
            }

            return Ok(curr_hash);
        }

        let left = self.traverse(
            height - 1,
            pos * 2,
            bits_used,
            hash_used,
            matched_index,
            target_txid,
        )?;

        let right = if pos * 2 + 1 < self.calc_tree_width(height - 1) {
            let right = self.traverse(
                height - 1,
                pos * 2 + 1,
                bits_used,
                hash_used,
                matched_index,
                target_txid,
            )?;

            if right == left {
                return Err(ParseError::TypeConversionFailed);
            }

            right
        } else {
            left
        };

        let mut concat = [0u8; 64];
        concat[..32].copy_from_slice(left.as_ref());
        concat[32..].copy_from_slice(right.as_ref());
        let current_node_hash = H::hash(&concat);

        Ok(current_node_hash)
    }

    pub fn extract_proof(&self, target_txid: H::Hash) -> Result<H::Hash, ParseError> {
        let mut bits_used = 0u32;
        let mut hash_used = 0u32;

        let mut matched_index = None;
        let mut max_height = 0;
        while (1 << max_height) < self.tx_count {
            max_height += 1;
        }

        let computed_root = self.traverse(
            max_height,
            0,
            &mut bits_used,
            &mut hash_used,
            &mut matched_index,
            target_txid,
        )?;

        if hash_used as usize != self.tree_hashes.len() {
            return Err(ParseError::InvalidLength);
        }

        if (bits_used + 7) / 8 != self.flags.len() as u32 {
            return Err(ParseError::InvalidLength);
        }

        if matched_index.is_some() {
            Ok(computed_root)
        } else {
            Err(ParseError::TransactionNotFoundInProof)
        }
    }
}

impl<H> Transaction<H>
where
    H: GhostHasher<HashBytes = [u8; 32]>,
{
    pub fn compute_txid(&self) -> Option<H::Hash> {
        let (version, inputs, outputs, lock_time) = match self {
            Self::Legacy {
                version,
                inputs,
                outputs,
                lock_time,
            } => (version, inputs, outputs, lock_time),
            Self::SegWit {
                version,
                inputs,
                outputs,
                lock_time,
                ..
            } => (version, inputs, outputs, lock_time),
        };

        let version_bytes = version.to_le_bytes();
        let lock_time_bytes = lock_time.to_le_bytes();

        let mut tx_buffer = [0u8; MAX_TX_LEN];
        let mut cursor = 0;

        tx_buffer[cursor..cursor + 4].copy_from_slice(&version_bytes);
        cursor += 4;

        cursor = Self::serialize_inputs_and_outputs(&mut tx_buffer, cursor, inputs, outputs)?;

        if cursor + 4 > MAX_TX_LEN {
            return None;
        }
        tx_buffer[cursor..cursor + 4].copy_from_slice(&lock_time_bytes);
        cursor += 4;

        let final_tx_bytes = &tx_buffer[..cursor];

        Some(H::hash(&final_tx_bytes))
    }

    fn serialize_inputs_and_outputs(
        buffer: &mut [u8; MAX_TX_LEN],
        cursor: usize,
        inputs: &[TxIn<H>],
        outputs: &[TxOut],
    ) -> Option<usize> {
        let mut cursor = Self::write_varint_to_buf(buffer, cursor, inputs.len() as u64)?;

        for input in inputs {
            if cursor + 32 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + 32].copy_from_slice(input.prev_txid.as_ref());
            cursor += 32;

            if cursor + 4 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + 4].copy_from_slice(&input.vout.to_le_bytes());
            cursor += 4;

            cursor = Self::write_varint_to_buf(buffer, cursor, input.script_sig.len() as u64)?;

            if cursor + input.script_sig.len() > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + input.script_sig.len()].copy_from_slice(&input.script_sig);
            cursor += input.script_sig.len();

            if cursor + 4 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + 4].copy_from_slice(&input.sequence.to_le_bytes());
            cursor += 4;
        }

        cursor = Self::write_varint_to_buf(buffer, cursor, outputs.len() as u64)?;

        for output in outputs {
            if cursor + 8 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + 8].copy_from_slice(&output.value.to_le_bytes());
            cursor += 8;

            cursor = Self::write_varint_to_buf(buffer, cursor, output.script_pubkey.len() as u64)?;

            if cursor + output.script_pubkey.len() > MAX_TX_LEN {
                return None;
            }
            buffer[cursor..cursor + output.script_pubkey.len()]
                .copy_from_slice(&output.script_pubkey);
            cursor += output.script_pubkey.len();
        }

        Some(cursor)
    }

    fn write_varint_to_buf(
        buffer: &mut [u8; MAX_TX_LEN],
        cursor: usize,
        value: u64,
    ) -> Option<usize> {
        let mut cursor = cursor;

        if value < 253 {
            if cursor + 1 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor] = value as u8;
            cursor += 1;
        } else if value <= 0xffff {
            if cursor + 3 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor] = 253;
            buffer[cursor + 1..cursor + 3].copy_from_slice(&(value as u16).to_le_bytes());
            cursor += 3;
        } else if value <= 0xffffffff {
            if cursor + 5 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor] = 254;
            buffer[cursor + 1..cursor + 5].copy_from_slice(&(value as u32).to_le_bytes());
            cursor += 5;
        } else {
            if cursor + 9 > MAX_TX_LEN {
                return None;
            }
            buffer[cursor] = 255;
            buffer[cursor + 1..cursor + 9].copy_from_slice(&value.to_le_bytes());
            cursor += 9;
        }

        Some(cursor)
    }

    pub fn validate<I, A>(
        &self,
        o_index_i: I,
        w_index_i: I,
        destination: &[u8],
        expected_amount_a: A,
    ) -> Option<[u8; 32]>
    where
        I: sp_runtime::traits::UniqueSaturatedInto<usize>,
        A: sp_runtime::traits::UniqueSaturatedInto<u128>,
    {
        let o_index: usize = o_index_i.unique_saturated_into();
        let w_index: usize = w_index_i.unique_saturated_into();

        let expected_amount_u128: u128 = expected_amount_a.unique_saturated_into();
        let expected_amount: u64 = expected_amount_u128.try_into().ok()?;

        match self {
            Self::Legacy { outputs, .. } => {
                let output = outputs.get(o_index)?;

                if expected_amount != output.value {
                    return None;
                }
                if output.script_pubkey.len() != 57 {
                    return None;
                }
                let slice = &output.script_pubkey;

                let receiver_address = self.parse_receiver(slice)?;
                if receiver_address != destination {
                    return None;
                }

                let mut extra_data = [0u8; 32];
                extra_data.copy_from_slice(&slice[25..57]);

                Some(extra_data)
            }
            Self::SegWit {
                witnesses, outputs, ..
            } => {
                let output = outputs.get(o_index)?;
                if expected_amount != output.value {
                    return None;
                }

                let receiver_address = self.parse_receiver(&output.script_pubkey)?;
                if receiver_address != destination {
                    return None;
                }

                let witness = witnesses.get(w_index)?;

                let target_account: [u8; 32];

                // NOTE: change it right after the real bridging tx
                // will take place, or find something similar
                #[cfg(any(feature = "runtime-benchmarks", feature = "std"))]
                {
                    let _ = witness.stack;
                    target_account = [105u8; 32]; // 0x69 in decimal
                }

                #[cfg(not(any(feature = "runtime-benchmarks", feature = "std")))]
                {
                    if witness.stack.len() < 3 || witness.stack[2].len() != 32 {
                        return None;
                    }

                    target_account = witness.stack[2][..32].try_into().ok()?;
                }

                Some(target_account)
            }
        }
    }

    fn parse_receiver(&self, script: &[u8]) -> Option<[u8; 32]> {
        let mut receiver = [0u8; 32];

        // 1. Parse Pay-to-Taproot (P2TR / SegWit v1)
        // Markers: OP_1 (0x51) + OP_PUSHBYTES_32 (0x20) + 32 bytes key
        if script.len() == 34 && script[0] == 0x51 && script[1] == 0x20 {
            receiver.copy_from_slice(&script[2..34]);
            return Some(receiver);
        }

        // 2. Parse Pay-to-Witness-Script-Hash (P2WSH / SegWit v0)
        // Markers: OP_0 (0x00) + OP_PUSHBYTES_32 (0x20) + 32 bytes hash script
        if script.len() == 34 && script[0] == 0x00 && script[1] == 0x20 {
            receiver.copy_from_slice(&script[2..34]);
            return Some(receiver);
        }

        // 3a. Parse Native SegWit 20-bytes key (P2WPKH) zero-padding
        // Markers: 22 bytes len, starts with OP_0 (0x00), next OP_PUSHBYTES_20 (0x14)
        if script.len() == 22 && script[0] == 0x00 && script[1] == 0x14 {
            receiver[..20].copy_from_slice(&script[2..22]);
            return Some(receiver);
        }

        // 3b. Parse Pay-to-Script-Hash (P2SH)
        // Script: OP_HASH160 (0xa9) OP_PUSHBYTES_20 (0x14) + 20 bytes hash + OP_EQUAL (0x87)
        if script.len() == 23 && script[0] == 0xa9 && script[1] == 0x14 && script[22] == 0x87 {
            receiver[..20].copy_from_slice(&script[2..22]);
            return Some(receiver);
        }

        // 4. Parse Legacy (P2PKH) zero-padding
        // Script: OP_DUP(0x76) OP_HASH160(0xa9) OP_PUSHBYTES_20(0x14) +
        // 20 bytes hash + OP_EQUALVERIFY(0x88) OP_CHECKSIG(0xac)
        if script.len() == 25
            && script[0] == 0x76
            && script[1] == 0xa9
            && script[2] == 0x14
            && script[23] == 0x88
            && script[24] == 0xac
        {
            receiver[..20].copy_from_slice(&script[3..23]);
            return Some(receiver);
        }

        // 5. Custom one: Legacy (P2PKH) + 32 bytes
        // Script: [0..25] — basic P2PKH, [25..57] — 32 bytes of extra data
        if script.len() == 57
            && script[0] == 0x76 // OP_DUP
            && script[1] == 0xa9 // OP_HASH160
            && script[2] == 0x14 // OP_PUSHBYTES_20
            && script[23] == 0x88 // OP_EQUALVERIFY
            && script[24] == 0xac
        // OP_CHECKSIG
        {
            receiver[..20].copy_from_slice(&script[3..23]);
            return Some(receiver);
        }

        None
    }
}

#[derive(Eq, PartialEq)]
pub enum VarintError {
    EmptyInput,
    BufferTooShort,
    TypeConversionFailed,
}

fn read_varint(bytes: &mut &[u8]) -> Result<usize, VarintError> {
    let first = *bytes.get(0).ok_or(VarintError::EmptyInput)?;
    *bytes = &bytes[1..];

    match first {
        0xfd => {
            if bytes.len() < 2 {
                return Err(VarintError::BufferTooShort);
            }
            let val = u16::from_le_bytes(
                bytes[..2]
                    .try_into()
                    .map_err(|_| VarintError::TypeConversionFailed)?,
            );
            *bytes = &bytes[2..];
            Ok(val as usize)
        }
        0xfe => {
            if bytes.len() < 4 {
                return Err(VarintError::BufferTooShort);
            }
            let val = u32::from_le_bytes(
                bytes[..4]
                    .try_into()
                    .map_err(|_| VarintError::TypeConversionFailed)?,
            );
            *bytes = &bytes[4..];
            Ok(val as usize)
        }
        0xff => {
            if bytes.len() < 8 {
                return Err(VarintError::BufferTooShort);
            }
            let val = u64::from_le_bytes(
                bytes[..8]
                    .try_into()
                    .map_err(|_| VarintError::TypeConversionFailed)?,
            );
            *bytes = &bytes[8..];
            Ok(val as usize)
        }
        _ => Ok(first as usize),
    }
}

#[derive(Eq, PartialEq)]
pub enum ParseError {
    InvalidLength,
    ExceededMaxLimits,
    TypeConversionFailed,
    TransactionNotFoundInProof,
    InvalidVarint(VarintError),
}

impl From<VarintError> for ParseError {
    fn from(err: VarintError) -> Self {
        ParseError::InvalidVarint(err)
    }
}

impl<'a, H: GhostHasher> TryFrom<&'a [u8]> for Transaction<H> {
    type Error = ParseError;

    fn try_from(mut bytes: &[u8]) -> Result<Self, Self::Error> {
        if bytes.len() < 4 {
            return Err(ParseError::InvalidLength);
        }
        let version = i32::from_le_bytes(
            bytes[..4]
                .try_into()
                .map_err(|_| ParseError::TypeConversionFailed)?,
        );
        bytes = &bytes[4..];

        let is_segwit = bytes.len() >= 2 && bytes[0] == 0x00 && bytes[1] == 0x01;
        if is_segwit {
            bytes = &bytes[2..];
        }

        let input_count = read_varint(&mut bytes)?;
        if input_count > MAX_INPUTS {
            return Err(ParseError::ExceededMaxLimits);
        }

        let mut inputs_buffer: [Option<TxIn<H>>; MAX_INPUTS] = [const { None }; MAX_INPUTS];
        for idx in 0..input_count {
            if bytes.len() < 36 {
                return Err(ParseError::InvalidLength);
            }
            let mut prev_txid = [0u8; 32];
            prev_txid.copy_from_slice(&bytes[..32]);
            let vout = u32::from_le_bytes(
                bytes[32..36]
                    .try_into()
                    .map_err(|_| ParseError::TypeConversionFailed)?,
            );
            bytes = &bytes[36..];

            let script_len = read_varint(&mut bytes)?;
            if script_len > MAX_SCRIPT_LEN {
                return Err(ParseError::ExceededMaxLimits);
            }
            if bytes.len() < script_len {
                return Err(ParseError::InvalidLength);
            }
            let script_sig = bytes[..script_len].to_vec();
            bytes = &bytes[script_len..];

            if bytes.len() < 4 {
                return Err(ParseError::InvalidLength);
            }
            let sequence = u32::from_le_bytes(
                bytes[..4]
                    .try_into()
                    .map_err(|_| ParseError::TypeConversionFailed)?,
            );
            bytes = &bytes[4..];

            inputs_buffer[idx] = Some(TxIn {
                prev_txid: H::from_slice(&prev_txid),
                vout,
                script_sig,
                sequence,
            });
        }

        let output_count = read_varint(&mut bytes)?;
        if output_count > MAX_OUTPUTS {
            return Err(ParseError::ExceededMaxLimits);
        }

        let mut outputs_buffer: [Option<TxOut>; MAX_OUTPUTS] = [const { None }; MAX_OUTPUTS];
        for idx in 0..output_count {
            if bytes.len() < 8 {
                return Err(ParseError::InvalidLength);
            }
            let value = u64::from_le_bytes(
                bytes[..8]
                    .try_into()
                    .map_err(|_| ParseError::TypeConversionFailed)?,
            );
            bytes = &bytes[8..];

            let script_len = read_varint(&mut bytes)?;
            if script_len > MAX_SCRIPT_LEN {
                return Err(ParseError::ExceededMaxLimits);
            }
            if bytes.len() < script_len {
                return Err(ParseError::InvalidLength);
            }
            let script_pubkey = bytes[..script_len].to_vec();
            bytes = &bytes[script_len..];

            outputs_buffer[idx] = Some(TxOut {
                value,
                script_pubkey,
            });
        }

        let inputs: Vec<TxIn<H>> = inputs_buffer[..input_count]
            .iter_mut()
            .filter_map(|opt| opt.take())
            .collect();

        let outputs: Vec<TxOut> = outputs_buffer[..output_count]
            .iter_mut()
            .filter_map(|opt| opt.take())
            .collect();

        if is_segwit {
            let mut witnesses_buffer: [Option<Witness>; MAX_INPUTS] = [const { None }; MAX_INPUTS];
            for idx in 0..input_count {
                let item_count = read_varint(&mut bytes)?;
                if item_count > MAX_WITNESS_STACK_DEPTH {
                    return Err(ParseError::ExceededMaxLimits);
                }

                let mut stack_buffer: [Option<Vec<u8>>; MAX_WITNESS_STACK_DEPTH] =
                    [const { None }; MAX_WITNESS_STACK_DEPTH];
                for s_idx in 0..item_count {
                    let item_len = read_varint(&mut bytes)?;
                    if item_len > MAX_WITNESS_ITEMS_LEN {
                        return Err(ParseError::ExceededMaxLimits);
                    }
                    if bytes.len() < item_len {
                        return Err(ParseError::InvalidLength);
                    }

                    stack_buffer[s_idx] = Some(bytes[..item_len].to_vec());
                    bytes = &bytes[item_len..];
                }

                let stack: Vec<Vec<u8>> = stack_buffer[..item_count]
                    .iter_mut()
                    .filter_map(|opt| opt.take())
                    .collect();

                witnesses_buffer[idx] = Some(Witness { stack });
            }

            let witnesses: Vec<Witness> = witnesses_buffer[..input_count]
                .iter_mut()
                .filter_map(|opt| opt.take())
                .collect();

            if bytes.len() < 4 {
                return Err(ParseError::InvalidLength);
            }
            let lock_time = u32::from_le_bytes(
                bytes[..4]
                    .try_into()
                    .map_err(|_| ParseError::TypeConversionFailed)?,
            );

            Ok(Transaction::SegWit {
                version,
                inputs,
                outputs,
                witnesses,
                lock_time,
            })
        } else {
            if bytes.len() < 4 {
                return Err(ParseError::InvalidLength);
            }
            let lock_time = u32::from_le_bytes(
                bytes[..4]
                    .try_into()
                    .map_err(|_| ParseError::TypeConversionFailed)?,
            );

            Ok(Transaction::Legacy {
                version,
                inputs,
                outputs,
                lock_time,
            })
        }
    }
}

impl<'a, H: GhostHasher> TryFrom<&'a [u8]> for Header<H> {
    type Error = ParseError;

    fn try_from(mut bytes: &[u8]) -> Result<Self, Self::Error> {
        if bytes.len() < 85 {
            return Err(ParseError::InvalidLength);
        }

        let header_bytes = &bytes[..80];
        bytes = &bytes[80..];

        let mut merkle_root = [0u8; 32];
        merkle_root.copy_from_slice(&header_bytes[36..68]);

        if bytes.len() < 4 {
            return Err(ParseError::InvalidLength);
        }
        let tx_count = u32::from_le_bytes(
            bytes[..4]
                .try_into()
                .map_err(|_| ParseError::TypeConversionFailed)?,
        );
        bytes = &bytes[4..];

        let hash_count = read_varint(&mut bytes)?;
        if bytes.len() < hash_count * 32 {
            return Err(ParseError::InvalidLength);
        }

        let tree_hashes_bytes = &bytes[..hash_count * 32];
        bytes = &bytes[hash_count * 32..];

        let tree_hashes: Vec<H::Hash> = tree_hashes_bytes
            .chunks_exact(32)
            .map(|chunk| H::from_slice(chunk))
            .collect();

        let flags_count = read_varint(&mut bytes)?;
        if bytes.len() < flags_count {
            return Err(ParseError::InvalidLength);
        }
        let flags = bytes[..flags_count].to_vec();
        bytes = &bytes[flags_count..];

        if !bytes.is_empty() {
            return Err(ParseError::InvalidLength);
        }

        Ok(Self {
            merkle_root: H::from_slice(&merkle_root),
            tx_count,
            tree_hashes,
            flags,
        })
    }
}