use std::{
    marker::PhantomData,
    io::{
        Write,
        Cursor
    },
    fmt::{Debug, Display, Formatter}
};

use f4jumble;

use borsh::{BorshDeserialize, BorshSerialize};

use haskell_ffi::{
    error::Result,
    from_haskell::{marshall_from_haskell_var, marshall_from_haskell_fixed},
    to_haskell::{marshall_to_haskell_var, marshall_to_haskell_fixed},
    FromHaskell, HaskellSize, ToHaskell
};

use zcash_primitives::{
    zip32::Scope as SaplingScope,
    transaction::components::sapling::{
        read_zkproof,
        GrothProofBytes,
        OutputDescription,
        CompactOutputDescription
    },
    sapling::{
        value::ValueCommitment as SaplingValueCommitment,
        keys::{
            FullViewingKey as SaplingViewingKey,
            PreparedIncomingViewingKey as SaplingPreparedIncomingViewingKey
        },
        note_encryption::SaplingDomain,
        PaymentAddress,
        note::ExtractedNoteCommitment as SaplingExtractedNoteCommitment
    },
    transaction::Transaction,
    consensus::{
        BranchId::Nu5,
        MainNetwork,
        BlockHeight
    }
};

use zcash_address::{
    Network,
    unified::{Address, Encoding, Ufvk, Container, Fvk},
    ZcashAddress
};

use zcash_client_backend::keys::sapling::ExtendedFullViewingKey;

use orchard::{
    Action,
    keys::{FullViewingKey, PreparedIncomingViewingKey, Scope},
    note::{Nullifier, TransmittedNoteCiphertext, ExtractedNoteCommitment},
    note_encryption::OrchardDomain,
    primitives::redpallas::{VerificationKey, SpendAuth, Signature},
    value::ValueCommitment
};

use zcash_note_encryption::EphemeralKeyBytes;

use bech32::{
    decode,
    u5,
    FromBase32,
    ToBase32,
    Variant
};

pub enum RW {}
pub const RW: PhantomData<RW> = PhantomData;

#[derive(BorshSerialize, BorshDeserialize)]
pub struct RawData {
    hrp: Vec<u8>,
    bytes: Vec<u8>
}

impl<RW> ToHaskell<RW> for RawData {
    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
        self.serialize(writer)?;
        Ok(())
    }
}

//impl<RW> FromHaskell<RW> for RawData {
    //fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
        //let x = RawData::deserialize(buf)?;
        //Ok(x)
    //}
//}

#[derive(BorshSerialize, BorshDeserialize)]
pub struct HshieldedOutput {
    cv: Vec<u8>,
    cmu: Vec<u8>,
    eph_key: Vec<u8>,
    enc_txt: Vec<u8>,
    out_txt: Vec<u8>,
    proof: Vec<u8>
}

impl<RW> FromHaskell<RW> for HshieldedOutput {
    fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
        let x = HshieldedOutput::deserialize(buf)?;
        Ok(x)
    }
}

#[derive(BorshSerialize, BorshDeserialize)]
pub struct Haction {
    nf: Vec<u8>,
    rk: Vec<u8>,
    cmx: Vec<u8>,
    eph_key: Vec<u8>,
    enc_txt: Vec<u8>,
    out_txt: Vec<u8>,
    cv: Vec<u8>,
    auth: Vec<u8>
}

impl<RW> FromHaskell<RW> for Haction {
    fn from_haskell(buf: &mut &[u8], _tag: PhantomData<RW>) -> Result<Self> {
        let x = Haction::deserialize(buf)?;
        Ok(x)
    }
}

#[derive(BorshSerialize, BorshDeserialize)]
pub struct Hnote {
    note: u64,
    recipient: Vec<u8>,
    memo: Vec<u8>
}

impl<RW> ToHaskell<RW> for Hnote {
    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
        self.serialize(writer)?;
        Ok(())
    }
}


#[derive(BorshSerialize, BorshDeserialize)]
pub struct Hufvk {
    net: u8,
    orchard: Vec<u8>,
    sapling: Vec<u8>,
    transparent: Vec<u8>
}

impl<RW> ToHaskell<RW> for Hufvk {
    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
        self.serialize(writer)?;
        Ok(())
    }
}

impl Hufvk {
    fn add_key_section(&mut self, fvk: &Fvk) {
        if let Fvk::Orchard(v) = fvk {
            self.orchard = v.to_vec();
        }
        if let Fvk::Sapling(w) = fvk {
            self.sapling = w.to_vec();
        }
        if let Fvk::P2pkh(x) = fvk {
            self.transparent = x.to_vec();
        }
    }
}

#[derive(BorshSerialize, BorshDeserialize)]
pub struct Hsvk {
    vk: Vec<u8>,
    ovk: Vec<u8>
}

impl<RW> ToHaskell<RW> for Hsvk {
    fn to_haskell<W: Write>(&self, writer: &mut W, _tag: PhantomData<RW>) -> Result<()> {
        self.serialize(writer)?;
        Ok(())
    }
}

fn to_array<T, const N: usize>(v: Vec<T>) -> [T; N] {
    v.try_into().unwrap_or_else(|v: Vec<T>| panic!("Expected a Vec of length {} but it was {}", N, v.len()))
}

#[no_mangle]
pub extern "C" fn rust_wrapper_f4jumble(
    input: *const u8,
    input_len: usize,
    out: *mut u8,
    out_len: &mut usize) {
    let input: Vec<u8> = marshall_from_haskell_var(input, input_len, RW);
    let result = f4jumble::f4jumble(&input).unwrap();
    marshall_to_haskell_var(&result, out, out_len, RW);
}

#[no_mangle]
pub extern "C" fn rust_wrapper_f4unjumble(
    input: *const u8,
    input_len: usize,
    out: *mut u8,
    out_len: &mut usize) {
    let input: Vec<u8> = marshall_from_haskell_var(input, input_len, RW);
    let result = f4jumble::f4jumble_inv(&input).unwrap();
    marshall_to_haskell_var(&result, out, out_len, RW);
}

#[no_mangle]
pub extern "C" fn rust_wrapper_ua_decode(
    input: *const u8,
    input_len: usize,) -> bool {
    let input: String = marshall_from_haskell_var(input, input_len, RW);
    Address::decode(&input).is_ok()
    //marshall_to_haskell_var(&result, out, out_len, RW);
}

#[no_mangle]
pub extern "C" fn rust_wrapper_shielded_decode(
    input: *const u8,
    input_len: usize) -> bool {
    let input: String = marshall_from_haskell_var(input, input_len, RW);
    ZcashAddress::try_from_encoded(&input).is_ok()
}

#[no_mangle]
pub extern "C" fn rust_wrapper_bech32decode(
    input: *const u8,
    input_len: usize,
    out: *mut u8,
    out_len: &mut usize
    ) {
    let input: String = marshall_from_haskell_var(input, input_len, RW);
    let decodedBytes = bech32::decode(&input);
    match decodedBytes {
        Ok((hrp, bytes, variant)) => {
            let rd = RawData {hrp: hrp.into(), bytes: Vec::<u8>::from_base32(&bytes).unwrap()};
            marshall_to_haskell_var(&rd, out, out_len, RW);
        }
        Err(_e) => {
            let rd1 = RawData {hrp: "fail".into(), bytes: vec![0]};
            marshall_to_haskell_var(&rd1, out, out_len, RW);
        }
    }
}

#[no_mangle]
pub extern "C" fn rust_wrapper_svk_decode(
    input: *const u8,
    input_len: usize
    ) -> bool {
    let input: Vec<u8> = marshall_from_haskell_var(input, input_len, RW);
    let svk = ExtendedFullViewingKey::read(&*input);
    match svk {
        Ok(k) => {
            true
        }
        Err(e) => {
            print!("{}", e);
            false
        }
    }
}

#[no_mangle]
pub extern "C" fn rust_wrapper_svk_check_address(
    key_input: *const u8,
    key_input_len: usize,
    address_input: *const u8,
    address_input_len: usize
    ) -> bool {
    let key_input: Vec<u8> = marshall_from_haskell_var(key_input, key_input_len, RW);
    let address_input: Vec<u8> = marshall_from_haskell_var(address_input, address_input_len, RW);
    let svk = ExtendedFullViewingKey::read(&*key_input);
    let sa = PaymentAddress::from_bytes(&to_array(address_input)).unwrap();
    match svk {
        Ok(k) => {
            let (div_index, def_address) = k.default_address();
            sa == def_address
        }
        Err(e) => {
            false
        }
    }
}

#[no_mangle]
pub extern "C" fn rust_wrapper_ufvk_decode(
    input: *const u8,
    input_len: usize,
    out: *mut u8,
    out_len: &mut usize
    ) {
    let input: String = marshall_from_haskell_var(input, input_len, RW);
    let dec_key = Ufvk::decode(&input);
    match dec_key {
        Ok((n, ufvk)) => {
            let x = match n {
                Network::Main => 1,
                Network::Test => 2,
                Network::Regtest => 3
                };
            let mut hk = Hufvk { net: x, orchard: vec![0], sapling: vec![0], transparent: vec![0] };
            let fvks = ufvk.items();
            fvks.iter().for_each(|k| hk.add_key_section(k));
            marshall_to_haskell_var(&hk, out, out_len, RW);
        }
        Err(_e) => {
            let hk0 = Hufvk { net: 0, orchard: vec![0], sapling: vec![0], transparent: vec![0] };
            marshall_to_haskell_var(&hk0, out, out_len, RW);
        }
    }
}

#[no_mangle]
pub extern "C" fn rust_wrapper_sapling_note_decrypt(
    key: *const u8,
    key_len: usize,
    note: *const u8,
    note_len: usize,
    out: *mut u8,
    out_len: &mut usize
    ){
    let evk: Vec<u8> = marshall_from_haskell_var(key, key_len, RW);
    let note_input: Vec<u8> = marshall_from_haskell_var(note,note_len,RW);
    let mut note_reader = Cursor::new(note_input);
    let svk = ExtendedFullViewingKey::read(&*evk);
    match svk {
        Ok(k) => {
            let domain = SaplingDomain::for_height(MainNetwork, BlockHeight::from_u32(2000000));
            let action2: Transaction = Transaction::read(&mut note_reader, Nu5).unwrap();
            let bundle = action2.sapling_bundle().unwrap();
            let sh_out = bundle.shielded_outputs();
            let fvk = k.to_diversifiable_full_viewing_key().to_ivk(SaplingScope::External);
            let pivk = SaplingPreparedIncomingViewingKey::new(&fvk);
            let result = zcash_note_encryption::try_note_decryption(&domain, &pivk, &sh_out[0]);
            match result {
                Some((n, r, m)) => {
                    let hn = Hnote {note: n.value().inner(), recipient: r.to_bytes().to_vec(), memo: m.as_slice().to_vec() };
                    marshall_to_haskell_var(&hn, out, out_len, RW);
                }
                None => {
                    let hn0 = Hnote { note: 0, recipient: vec![0], memo: vec![0] };
                    marshall_to_haskell_var(&hn0, out, out_len, RW);
                }
            }
        }
        Err(_e) => {
            let hn0 = Hnote { note: 0, recipient: vec![0], memo: vec![0] };
            marshall_to_haskell_var(&hn0, out, out_len, RW);
        }
    }
}

#[no_mangle]
pub extern "C" fn rust_wrapper_orchard_note_decrypt(
    key: *const u8,
    key_len: usize,
    note: *const u8,
    note_len: usize,
    out: *mut u8,
    out_len: &mut usize
    ){
    let fvk_input: Vec<u8> = marshall_from_haskell_var(key, key_len, RW);
    let note_input: Haction = marshall_from_haskell_var(note, note_len, RW);
    let action: Action<Signature<SpendAuth>> = Action::from_parts(
        Nullifier::from_bytes(&to_array(note_input.nf)).unwrap(), 
        VerificationKey::try_from(to_array(note_input.rk)).unwrap(), 
        ExtractedNoteCommitment::from_bytes(&to_array(note_input.cmx)).unwrap(), 
        TransmittedNoteCiphertext {epk_bytes: to_array(note_input.eph_key), enc_ciphertext: to_array(note_input.enc_txt), out_ciphertext: to_array(note_input.out_txt)}, 
        ValueCommitment::from_bytes(&to_array(note_input.cv)).unwrap(), 
        Signature::from(to_array(note_input.auth)));
    let fvk_array = to_array(fvk_input);
    let domain = OrchardDomain::for_nullifier(*action.nullifier());
    let dec_fvk = FullViewingKey::from_bytes(&fvk_array);
    match dec_fvk {
        Some(fvk) => {
            let ivk = fvk.to_ivk(Scope::External);
            let pivk = PreparedIncomingViewingKey::new(&ivk);
            let result = zcash_note_encryption::try_note_decryption(&domain, &pivk, &action);
            match result {
                Some((n, r, m)) => {
                    let hn = Hnote {note: n.value().inner(), recipient: r.to_raw_address_bytes().to_vec(), memo: m.to_vec() };
                    marshall_to_haskell_var(&hn, out, out_len, RW);
                }
                None => {
                    let hn0 = Hnote { note: 0, recipient: vec![0], memo: vec![0] };
                    marshall_to_haskell_var(&hn0, out, out_len, RW);
                }
            }
        },
        None => {
            let hn0 = Hnote { note: 0, recipient: vec![0], memo: vec![0] };
            marshall_to_haskell_var(&hn0, out, out_len, RW);
        }
    }
}