import struct
import json
import base64
from datetime import datetime
from .Tool import read_var_int, hash_hash256

# Read the header and transaction count.
# Reference: https://en.bitcoin.it/wiki/Block#Block_structure
class Header:
  def __init__(self, blk_data: bytes):
    # Magic no, block size, and version
    (self.magic_no, self.block_size) = struct.unpack('<LL', blk_data[0:8])
    self.version = blk_data[8:12]
    # Hash values.
    self.hash_prev_block = blk_data[12:44][::-1]  # 32 bytes, little-endian
    self.hash_merkle_root = blk_data[44:76][::-1] # 32 bytes, little-endian
    # Block time, bits, and nonce
    (self.block_time, self.bits, self.nonce) = struct.unpack('<LLL', blk_data[76:88])
    self.block_dt = datetime.fromtimestamp(self.block_time)
    # Transaction count.
    (self.txn_count, txn_count_length) = read_var_int(blk_data, 88)
    self.num_bytes = 88 + txn_count_length
    # Calculate the hash value.
    # Reference: https://en.bitcoin.it/wiki/Block_hashing_algorithm
    self.hash_value = hash_hash256(blk_data[0:self.num_bytes][8:88])[::-1]
  
  def __len__(self):
    return self.num_bytes
    
  def __str__(self):
    return self.to_json(indent=2)
  
  def to_dict(self):
    return {
      'magic_no':         self.magic_no,
      'block_size':       self.block_size,
      'version':          self.version.hex(),
      'hash_prev_block':  self.hash_prev_block.hex(),
      'hash_merkle_root': self.hash_merkle_root.hex(),
      'block_time':       self.block_time,
      'bits':             self.bits,
      'nonce':            self.nonce,
      'txn_count':        self.txn_count,
      'hash_value':       self.hash_value.hex(),
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)
  
  def get_metadata(self):
    return {
      'version':          self.version.hex(),
      'hash_value':       self.hash_value.hex(),
    }

  
# Build classes to represent the Txin, Txout, and Transaction.
# Reference: https://github.com/bitcoin/bips/blob/master/bip-0144.mediawiki
class Txin:
  def __init__(self, blk_data: bytes, start_index: int):
    cur_index = start_index
    # Previous transaction hash.
    self.prev_txn_hash = blk_data[cur_index:cur_index+32][::-1]
    cur_index += 32
    # Previous txout index.
    (self.prev_txout_index,) = struct.unpack('<L', blk_data[cur_index:cur_index+4])
    cur_index += 4
    # Txin script length.
    (self.txin_script_len, var_int_len) = read_var_int(blk_data, cur_index)
    cur_index += var_int_len
    # Txin script (scriptSig, unlocking script).
    self.txin_script = Script(blk_data[cur_index:cur_index+self.txin_script_len], is_locking_script=False)
    cur_index += self.txin_script_len
    # Sequence no.
    self.seq_no = blk_data[cur_index:cur_index+4].hex()
    cur_index += 4
    # Record length.
    self.num_bytes = cur_index - start_index # Indicate how many bytes are read for this record.
  
  def __len__(self):
    return self.num_bytes

  def __str__(self):
    return self.to_json(indent=2)
  
  def to_dict(self):
    return {
      'prev_txn_hash':    self.prev_txn_hash.hex(),
      'prev_txout_index': self.prev_txout_index,
      'txin_script':      self.txin_script.to_dict(),
      'seq_no':           self.seq_no,
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)


class Txout:
  def __init__(self, blk_data: bytes, start_index: int):
    cur_index = start_index
    # Value, non negative integer giving the number of Satoshis(BTC/10^8) to be transferred.
    (self.num_of_satoshis,) = struct.unpack('<Q', blk_data[cur_index:cur_index+8])
    cur_index += 8
    # Txout script length.
    (self.txout_script_len, var_int_len) = read_var_int(blk_data, cur_index)
    cur_index += var_int_len
    # Txout script (scriptPubKey, locking script).
    self.txout_script = Script(blk_data[cur_index:cur_index+self.txout_script_len], is_locking_script=True)
    cur_index += self.txout_script_len
    # Record length.
    self.num_bytes = cur_index - start_index
    
  def __len__(self):
    return self.num_bytes
    
  def __str__(self):
    return self.to_json(indent=2)

  def to_dict(self):
    return {
      'num_of_satoshis':  self.num_of_satoshis,
      'txout_script':     self.txout_script.to_dict(),
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)


# Represent segwit items corresponding to one txin.
# Reference: https://github.com/bitcoin/bips/blob/master/bip-0144.mediawiki
class Witness:
  def __init__(self, blk_data: bytes, start_index: int):
    cur_index = start_index
    (self.num_items, var_int_len) = read_var_int(blk_data, cur_index)
    cur_index += var_int_len
    self.item_list = []
    for _ in range(self.num_items): # Read out all segwit items of current txin.
      (item_len, var_int_len) = read_var_int(blk_data, cur_index)
      cur_index += var_int_len
      self.item_list.append(blk_data[cur_index:cur_index+item_len])
      cur_index += item_len
    self.num_bytes = cur_index - start_index
  
  def __len__(self):
    return self.num_bytes
    
  def __str__(self):
    return self.to_json(indent=2)
  
  def to_dict(self):
    return {
      'num_items':        self.num_items,
      'item_base64_list': [base64.b64encode(item).decode('utf-8') for item in self.item_list],
      'item_str_list':    [item.hex() for item in self.item_list],
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)


# Represent a script in txin or txout.
class Script:
  def __init__(self, script_bytes: bytes, is_locking_script: bool):
    self.script_bytes: bytes = script_bytes            # The script in the format of byte array.
    self.is_locking_script: bool= is_locking_script  # True if the script is a locking one, otherwise it's an unlocking one.
    self.lock_time = None
    
  def __len__(self):
    return len(self.script_bytes)

  def __str__(self):
    return self.to_json(indent=2)

  def set_lock_time(self, lock_time):
    self.lock_time = lock_time
    
  def to_dict(self):
    return {
      'script_base64':      base64.b64encode(self.script_bytes).decode('utf-8'),
      'script_str':         self.script_bytes.hex(),
      'is_locking_script':  self.is_locking_script,
      'lock_time':          self.lock_time,
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)


class Transaction:
  def read_all_txin(self, blk_data: bytes, start_index: int):
    # Read in `self.in_counter` records.
    # Return the number of all bytes that we read.
    self.txin_list: list[Txin] = []
    cur_index = start_index
    for _ in range(self.in_counter):
      txin = Txin(blk_data, cur_index)
      self.txin_list.append(txin)
      txin_len = len(txin)
      cur_index += txin_len
    return cur_index - start_index
  
  def read_all_txout(self, blk_data, start_index):
    # Read in `self.out_counter` records.
    # Return the number of all bytes that we read.
    self.txout_list: list[Txout] = []
    cur_index = start_index
    for _ in range(self.out_counter):
      txout = Txout(blk_data, cur_index)
      self.txout_list.append(txout)
      txout_len = len(txout)
      cur_index += txout_len
    return cur_index - start_index
  
  def __init__(self, blk_data, start_index):
    cur_index = start_index
    # According to https://developer.bitcoin.org/reference/transactions.html,
    # Not all fields are used to compute the hash value of a transaction.
    # `hash_bytes_data` is used to represent the bytes data used for hashing.
    hash_bytes_data = b''
    # Version no.
    (self.version_no,) = struct.unpack('<L', blk_data[cur_index:cur_index+4])
    hash_bytes_data += blk_data[cur_index:cur_index+4]
    cur_index += 4
    # Flag.
    self.flag = (blk_data[cur_index] == 0x00) and (blk_data[cur_index+1] == 0x01)
    if self.flag:
      cur_index += 2
    # In-counter.
    (self.in_counter, var_int_len) = read_var_int(blk_data, cur_index)
    hash_bytes_data += blk_data[cur_index:cur_index+var_int_len]
    cur_index += var_int_len
    # List of inputs.
    inputs_len = self.read_all_txin(blk_data, cur_index)
    hash_bytes_data += blk_data[cur_index:cur_index+inputs_len]
    cur_index += inputs_len
    # Out-counter.
    (self.out_counter, var_int_len) = read_var_int(blk_data, cur_index)
    hash_bytes_data += blk_data[cur_index:cur_index+var_int_len]
    cur_index += var_int_len
    # List of outputs.
    outputs_len = self.read_all_txout(blk_data, cur_index)
    hash_bytes_data += blk_data[cur_index:cur_index+outputs_len]
    cur_index += outputs_len
    # Witnesses
    self.segwit_list: list[Witness] = []
    if self.flag: # https://github.com/bitcoin/bips/blob/master/bip-0144.mediawiki
      for _ in range(self.in_counter): # Each txin has its own segwit.
        witness = Witness(blk_data, cur_index)
        cur_index += len(witness)
        self.segwit_list.append(witness)
    # Lock time.
    self.lock_time = blk_data[cur_index:cur_index+4].hex()
    hash_bytes_data += blk_data[cur_index:cur_index+4]
    cur_index += 4
    # Transaction num_bytes
    self.num_bytes = cur_index - start_index
    # Store the raw data.
    self.raw_txn_bytes = blk_data[start_index:cur_index]
    # Compute the hash value.
    self.hash_value = hash_hash256(hash_bytes_data)[::-1]
    # Assume this transaction is not a coinbase one at this moment.
    # The constructor of Block will take care of set the first transaction in each block to be a coinbase one.
    self.is_coinbase = False
    # Remember to tell the Script objects the value of lock time.
    for txin in self.txin_list:
      txin.txin_script.set_lock_time(self.lock_time)
    for txout in self.txout_list:
      txout.txout_script.set_lock_time(self.lock_time)
  
  def __len__(self):
    return self.num_bytes

  def __str__(self):
    return self.to_json(indent=2)
    
  def to_dict(self):
    return {
      'version_no':  self.version_no,
      'segwit_flag': self.flag,
      'in_count':    self.in_counter,
      'txin_list':   [txin.to_dict() for txin in self.txin_list],
      'out_count':   self.out_counter,
      'txout_list':  [txout.to_dict() for txout in self.txout_list],
      'segwit_list': [wit.to_dict() for wit in self.segwit_list],
      'lock_time':   self.lock_time,
      'hash_value':  self.hash_value.hex(),
      'is_coinbase': self.is_coinbase,
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)
  
  def get_metadata(self):
    return {
      'hash_value':     self.hash_value.hex(),
      'raw_txn_bytes':  self.raw_txn_bytes.hex(),
    }


# Build the class representing a block.
class Block:
  def __init__(self, blk_data: bytes):
    self.header = Header(blk_data)
    self.txn_list: list[Transaction] = []
    cur_index = len(self.header)
    for i in range(self.header.txn_count):
      txn = Transaction(blk_data, cur_index)
      if i == 0:
        txn.is_coinbase = True # The first transaction in each block is a coinbase one.
      self.txn_list.append(txn)
      cur_index += len(txn)
  
  def __len__(self):
    result = len(self.header)
    for txn in self.txn_list:
      result += len(txn)
    return result

  def __str__(self):
    return self.to_json(indent=2)

  def to_dict(self):
    return {
      'header':   self.header.to_dict(),
      'txn_list': [txn.to_dict() for txn in self.txn_list],
    }
  
  def to_json(self, **kwargs):
    return json.dumps(self.to_dict(), **kwargs)
  
  def get_metadata(self):
    return {
      'header':   self.header.get_metadata(),
      'txn_list': [txn.get_metadata() for txn in self.txn_list],
    }


# Split the file content into block data.
def split_dat_file(f_content: bytes) -> list[bytes]:
  blk_data_list = []
  cur_index = 0
  while cur_index < len(f_content):
    (blk_size,) = struct.unpack('<L', f_content[cur_index+4:cur_index+8])
    blk_size += 8
    blk_data_list.append(f_content[cur_index:cur_index+blk_size])
    cur_index += blk_size
  return blk_data_list


# Read the content of the block data file (blkxxxxx.dat).
def read_blk_data_in_dat_file(path_to_file: str) -> list[bytes]: 
  f = open(path_to_file, 'rb')
  f_content = f.read()
  f.close()
  return split_dat_file(f_content)

# Count up the number of blocks in the raw data file (blkxxxxx.dat).
def count_up_blk_in_dat_file(path_to_file: str) -> int:
  return len(read_blk_data_in_dat_file(path_to_file))
