import re
import os
import os.path
import sys
import time
import struct
import socket
import pathlib
import logging
import argparse

import tqdm
import cbor2

from Crypto.Hash import SHA256
from Crypto.Cipher import Salsa20
from Crypto.Random import get_random_bytes

BYAFN_ADMINSOCKET = os.getenv(
  'BYAFN_ADMINSOCKET',
  default='./adminsocket'
)

logging.basicConfig(
  format='%(asctime)s %(levelname)s: %(message)s',
  level=logging.INFO
)

class Metafile:
  HEADER = b'\x80BYAFN-METAFILE\x00\x00'
  HEADER_LEN = len(HEADER)

  def __init__(self, filename, size, checksum, pieces, key=None):
    self.filename = filename
    self.size = size
    self.checksum = checksum
    self.pieces = pieces
    self.key = key

  def save(self, f):
    f.write(Metafile.HEADER)

    filename = self.filename.encode('UTF-8')
    filename_len = len(filename)
  
    f.write(struct.pack('!I', filename_len))
    f.write(filename)

    f.write(struct.pack('!L', self.size))

    f.write(self.checksum)

    if self.key:
      f.write(b'Y')
      f.write(self.key)
    else:
      f.write(b'N')
      
    pieces_count = len(self.pieces)

    f.write(struct.pack('!L', pieces_count))
    
    for hash in self.pieces:
      f.write(hash)

  @staticmethod
  def load(f):
    header = f.read(Metafile.HEADER_LEN)
    if header != Metafile.HEADER:
      raise ValueError

    filename_len = f.read(4)
    filename_len = struct.unpack('!I', filename_len)[0]

    filename = f.read(filename_len)
    if len(filename) != filename_len:
      raise ValueError

    filename = filename.decode('UTF-8')
    filename = pathlib.Path(filename).name

    size = f.read(4)
    size = struct.unpack('!I', size)[0]

    checksum = f.read(32)
    if len(checksum) != 32:
      raise ValueError

    is_encrypted = f.read(1)
    if is_encrypted == b'Y':
      key = f.read(40)
      if len(key) != 40:
        raise ValueError
    elif is_encrypted == b'N':
      key = None
    else:
      raise ValueError

    pieces_count = f.read(4)
    pieces_count = struct.unpack('!L', pieces_count)[0]

    if pieces_count < 1:
      raise ValueError

    pieces = []

    while pieces_count > 0:
      piece_hash = f.read(32)
      if len(piece_hash) != 32:
        raise ValueError

      pieces.append(piece_hash)

      pieces_count -= 1

    return Metafile(filename, size, checksum, pieces, key)

def sha256(data):
  hash = SHA256.new()
  hash.update(data)

  return hash.digest()

def salsa20_create_encryptor():
  key = get_random_bytes(32)
  cipher = Salsa20.new(key=key)

  return (cipher.nonce + key, cipher)

def salsa20_create_decryptor(key):
  if len(key) != 40:
    raise ValueError

  cipher = Salsa20.new(key=key[8:], nonce=key[:8])
  
  return cipher

def parse_file_size(size):
  units = {
    'b': 1, 
    'kb': 10e3,
    'mb': 10e6
  }

  match = re.match(r'(\d+)([km]b)', size.strip().lower())
  if not match:
    raise ValueError

  size, unit = int(match.group(1)), units[match.group(2)]

  return int(size * unit)

def send_command(**command):
  try:
    conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    conn.connect(args.admin_socket_path)

    command = cbor2.dumps(command)
    
    conn.send(struct.pack('<I', len(command)))
    conn.send(command)

    length = struct.unpack('<I', conn.recv(4))[0]
    response = conn.recv(length)
    response = cbor2.loads(response)

    conn.close()
  except Exception as e:
    logging.error(f'Contacting the admin socket `{args.admin_socket_path}\': {e}')

    sys.exit(1)

  return response

parser = argparse.ArgumentParser()
parser.add_argument(
  '-G', '--genconf',
  help='Generate configuration file',
  action='store_true'
)

parser.add_argument(
  '--listen-address',
  help='Set ListenAddress (use together with --genconf)',
  type=str, default='0.0.0.0'
)

parser.add_argument(
  '-p', '--listen-port',
  help='Set ListenPort (use together with --genconf)',
  type=int, default=42424
)

parser.add_argument(
  '-s', '--storage-path',
  help='Set StoragePath (use together with --genconf)',
  type=str, default='./storage'
)

parser.add_argument(
  '-a', '--admin-socket-path',
  help='Set AdminSocketPath (use together with --genconf, --share or --query)',
  type=str, default=BYAFN_ADMINSOCKET
)

parser.add_argument(
  '-S', '--share',
  help='Share a file',
  type=argparse.FileType('rb')
)

parser.add_argument(
  '-e', '--encrypt',
  help='Encrypt file before sharing (use together with --share)',
  action='store_true'
)

parser.add_argument(
  '-Q', '--query',
  help='Query a file',
  type=argparse.FileType('rb')
)

parser.add_argument(
  '-o', '--output',
  help='Set output file',
  type=argparse.FileType('wb')
)

parser.add_argument(
  '-P', '--piece-size',
  help='Set piece size',
  type=parse_file_size, default='16kb'
)

parser.add_argument(
  '--shutdown',
  help='Request graceful shutdown of the peer',
  type=int
)

parser.add_argument(
  '--cleanup',
  help='Remove every piece from the storage if they weren\'t accessed more than for specified amount of seconds',
  type=int
)

args = parser.parse_args()

if args.genconf:
  import ecies

  key = ecies.utils.generate_eth_key()
  pk = key.public_key.to_hex()
  sk = key.to_hex()

  if not args.output:
    output = sys.stdout.buffer
  else:
    output = args.output

  output.write(bytes(f'''{{
  Peers: []

  ListenAddress: {args.listen_address}
  ListenPort: {args.listen_port}

  Key: {pk}
  Secret: {sk}

  StoragePath: {os.path.abspath(args.storage_path)}
  AdminSocketPath: {os.path.abspath(args.admin_socket_path)}
}}''', 'UTF-8'))

if args.share:
  with args.share as f:
    filename = pathlib.Path(f.name).name

    f.seek(0, os.SEEK_END)
    file_size = f.tell()
    f.seek(0)

    checksum = SHA256.new()

    if args.encrypt:
      key, cipher = salsa20_create_encryptor()

    pieces_count = max(1, file_size // args.piece_size)
    progress = tqdm.trange(pieces_count)
    pieces = []

    while True:
      piece = f.read(args.piece_size)

      if not piece:
        break

      checksum.update(piece)

      if args.encrypt:
        piece = cipher.encrypt(piece)

      hash = send_command(store={
        'piece': piece
      })['hash']

      pieces.append(hash)
      progress.update(1)

    metafile = Metafile(
      filename,
      file_size,
      checksum.digest(),
      pieces,
      key=key if args.encrypt else None
    )

    if not args.output:
      output = open(f'{filename}.byafn', 'wb')
    else:
      output = args.output

    metafile.save(output)

    output.close()

if args.query:
  with args.query as f:
    metafile = Metafile.load(f)

    if not args.output:
      output = open(metafile.filename, 'wb')
    else:
      output = args.output

    if metafile.key:
      cipher = salsa20_create_decryptor(metafile.key)

    checksum = SHA256.new()

    progress = tqdm.trange(len(metafile.pieces))

    for piece in metafile.pieces:
      interval = 10

      while True:
        data = send_command(query={
          'hash': piece
        })

        if 'piece' not in data:
          logging.error(f'Piece `{piece.hex()}\' is not available')
          logging.info(f'Retrying in {interval}sec.')

          time.sleep(interval)

          if interval < 120:
            interval += 10

          continue

        break

      piece = data['piece']

      if metafile.key:
        piece = cipher.decrypt(piece)

      checksum.update(piece)

      if checksum.digest() != metafile.checksum:
        logging.error('Integrity check failed')

        sys.exit(1)

      output.write(piece)
      progress.update(1)

    output.close()

if args.shutdown is not None:
  send_command(shutdown={
    'delay': args.shutdown
  })

  logging.info('Shutdown command sent')

if args.cleanup is not None:
  response = send_command(cleanup={
    'age': args.cleanup
  })

  logging.info(f'Removed {response["removed_count"]} piece(s).')