#include <graphene/peerplays_sidechain/bitcoin/serialize.hpp>

#include <fc/io/raw.hpp>

#include <graphene/peerplays_sidechain/bitcoin/sign_bitcoin_transaction.hpp>

namespace graphene { namespace peerplays_sidechain { namespace bitcoin {

const secp256k1_context_t *btc_context() {
   static secp256k1_context_t *ctx = secp256k1_context_create(SECP256K1_CONTEXT_VERIFY | SECP256K1_CONTEXT_SIGN);
   return ctx;
}

fc::sha256 get_signature_hash(const bitcoin_transaction &tx, const bytes &scriptCode, int64_t amount,
                              size_t in_index, int hash_type, bool is_witness) {
   fc::datastream<size_t> ps;
   if (is_witness)
      pack_tx_witness_signature(ps, scriptCode, tx, in_index, amount, hash_type);
   else
      pack_tx_signature(ps, scriptCode, tx, in_index, hash_type);

   std::vector<char> vec(ps.tellp());
   if (!vec.empty()) {
      fc::datastream<char *> ds(vec.data(), vec.size());
      if (is_witness)
         pack_tx_witness_signature(ds, scriptCode, tx, in_index, amount, hash_type);
      else
         pack_tx_signature(ds, scriptCode, tx, in_index, hash_type);
   }

   return fc::sha256::hash(fc::sha256::hash(vec.data(), vec.size()));
}

std::vector<char> privkey_sign(const bytes &privkey, const fc::sha256 &hash, const secp256k1_context_t *context_sign) {
   bytes sig;
   sig.resize(72);
   int sig_len = sig.size();

   FC_ASSERT(secp256k1_ecdsa_sign(
         context_sign,
         reinterpret_cast<unsigned char *>(hash.data()),
         reinterpret_cast<unsigned char *>(sig.data()),
         &sig_len,
         reinterpret_cast<const unsigned char *>(privkey.data()),
         secp256k1_nonce_function_rfc6979,
         nullptr)); // TODO: replace assert with exception

   sig.resize(sig_len);

   return sig;
}

std::vector<bytes> sign_witness_transaction_part(const bitcoin_transaction &tx, const std::vector<bytes> &redeem_scripts,
                                                 const std::vector<uint64_t> &amounts, const bytes &privkey,
                                                 const secp256k1_context_t *context_sign, int hash_type) {
   FC_ASSERT(tx.vin.size() == redeem_scripts.size() && tx.vin.size() == amounts.size());
   FC_ASSERT(!privkey.empty());

   std::vector<bytes> signatures;
   for (size_t i = 0; i < tx.vin.size(); i++) {
      const auto sighash = get_signature_hash(tx, redeem_scripts[i], static_cast<int64_t>(amounts[i]), i, hash_type, true);
      auto sig = privkey_sign(privkey, sighash, context_sign);
      sig.push_back(static_cast<uint8_t>(hash_type));

      signatures.push_back(sig);
   }
   return signatures;
}

void sign_witness_transaction_finalize(bitcoin_transaction &tx, const std::vector<bytes> &redeem_scripts, bool use_mulisig_workaround) {
   FC_ASSERT(tx.vin.size() == redeem_scripts.size());

   for (size_t i = 0; i < tx.vin.size(); i++) {
      if (use_mulisig_workaround)
         tx.vin[i].scriptWitness.insert(tx.vin[i].scriptWitness.begin(), bytes()); // Bitcoin workaround CHECKMULTISIG bug
      tx.vin[i].scriptWitness.push_back(redeem_scripts[i]);
   }
}

bool verify_sig(const bytes &sig, const bytes &pubkey, const bytes &msg, const secp256k1_context_t *context) {
   std::vector<unsigned char> sig_temp(sig.begin(), sig.end());
   std::vector<unsigned char> pubkey_temp(pubkey.begin(), pubkey.end());
   std::vector<unsigned char> msg_temp(msg.begin(), msg.end());

   int result = secp256k1_ecdsa_verify(context, msg_temp.data(), sig_temp.data(), sig_temp.size(), pubkey_temp.data(), pubkey_temp.size());
   return result == 1;
}

std::vector<std::vector<bytes>> sort_sigs(const bitcoin_transaction &tx, const std::vector<bytes> &redeem_scripts,
                                          const std::vector<uint64_t> &amounts, const secp256k1_context_t *context) {
   FC_ASSERT(redeem_scripts.size() == amounts.size());

   using data = std::pair<size_t, bytes>;
   struct comp {
      bool operator()(const data &lhs, const data &rhs) const {
         return lhs.first < rhs.first;
      }
   };

   std::vector<std::vector<bytes>> new_stacks;

   for (size_t i = 0; i < redeem_scripts.size(); i++) {
      const std::vector<bytes> &keys = get_pubkey_from_redeemScript(redeem_scripts[i]);
      const auto &sighash = get_signature_hash(tx, redeem_scripts[i], static_cast<int64_t>(amounts[i]), i, 1, true).str();
      bytes sighash_temp(parse_hex(sighash));

      std::vector<bytes> stack(tx.vin[i].scriptWitness);
      std::vector<bool> marker(tx.vin[i].scriptWitness.size(), false);
      std::set<data, comp> sigs;

      for (size_t j = 0; j < keys.size(); j++) {
         for (size_t l = 0; l < stack.size(); l++) {
            if (!verify_sig(stack[l], keys[j], sighash_temp, context) || marker[l])
               continue;
            sigs.insert(std::make_pair(j, stack[l]));
            marker[l] = true;
            break;
         }
      }

      std::vector<bytes> temp_sig;
      for (auto s : sigs) {
         temp_sig.push_back(s.second);
      }
      new_stacks.push_back(temp_sig);
   }
   return new_stacks;
}

void add_signatures_to_transaction_multisig(bitcoin_transaction &tx, std::vector<std::vector<bytes>> &signature_set) {
   for (unsigned int i = 0; i < signature_set.size(); i++) {
      std::vector<bytes> signatures = signature_set[i];
      FC_ASSERT(signatures.size() == tx.vin.size(), "Invalid signatures number");
      for (unsigned int i = 0; i < tx.vin.size(); i++)
         tx.vin[i].scriptWitness.push_back(signatures[i]);
   }
}

void add_signatures_to_transaction_weighted_multisig(bitcoin_transaction &tx, std::vector<std::vector<bytes>> &signature_set) {
   for (unsigned int i = 0; i < signature_set.size(); i++) {
      std::vector<bytes> signatures = signature_set[i];
      FC_ASSERT(signatures.size() == tx.vin.size(), "Invalid signatures number");
      // push signatures in reverse order because script starts to check the top signature on the stack first
      for (unsigned int i = 0; i < tx.vin.size(); i++)
         tx.vin[i].scriptWitness.insert(tx.vin[i].scriptWitness.begin(), signatures[i]);
   }
}

void add_signatures_to_transaction_user_weighted_multisig(bitcoin_transaction &tx, std::vector<std::vector<bytes>> &signature_set) {
   add_signatures_to_transaction_weighted_multisig(tx, signature_set);
   for (size_t itr = 0; itr < tx.vin.size(); itr++) {
      tx.vin[0].scriptWitness.push_back(bytes());
   }
}

}}} // namespace graphene::peerplays_sidechain::bitcoin
