#!/usr/bin/env python3
# Databankr: a CLI tool to encode arbitrary data to Casio Databank/Telememo 
# watches and restore it from there
# Created by Luxferre in 2024, released into public domain

import sys, math, json, re

# universal base conversion methods

def to_base(number, base, charset):
  if not number:
    return charset[0]
  res = ''
  while number > 0:
    res = charset[number % base] + res
    number //= base
  return res

def from_base(number, base, charset):
  if number == charset[0]:
    return 0
  res = 0
  for c in number:
    ind = charset.find(c)
    if ind > -1:
      res = res * base + ind
    else:
      return 0
  return res

# record preparation methods

# create a padded field from a numeric value
def val_to_field(n, charset, padlen):
  return to_base(n, len(charset), charset).rjust(padlen, charset[0])

# main encoding method
def encode(data: bytes, config):
  effective_namelen = config['namelen'] - 1
  alphabase = len(config['alpha'])
  numbase = len(config['digit'])
  indexsize = len(config['index'])
  # calculate record estimation
  name_part_bits = int(effective_namelen * math.log2(alphabase))
  number_part_bits = int(config['numberlen'] * math.log2(numbase))
  record_bits = name_part_bits + number_part_bits # single record capacity
  max_bits = record_bits * indexsize # overall databank capacity
  # start processing
  bitstr = ''.join(f'{c:08b}' for c in data) # create an aligned bitstring
  bitlen = len(bitstr) # overall bitstring length
  if bitlen > max_bits: # truncate the excess
    bitlen = max_bits
    bitstr = bitstr[:max_bits]
  rec_len = int(math.ceil(bitlen / record_bits)) # message length in records
  records = [] # list of lists
  pos = 0 # current position tracker
  for i in range(0, rec_len): # slice over the bitstring
    namebin = bitstr[pos:pos+name_part_bits].ljust(name_part_bits, '0')
    pos += name_part_bits
    numbin = bitstr[pos:pos+number_part_bits].ljust(number_part_bits, '0')
    pos += number_part_bits
    # now we only got binary representation of both parts
    # let's convert them to bigints and then to the actual records
    namefield = config['index'][i] + val_to_field(int(namebin, 2), 
                  config['alpha'], effective_namelen)
    numfield = val_to_field(int(numbin, 2), 
                  config['digit'], config['numberlen'])
    records.append([namefield, numfield])
  return records

# main decoding method
def decode(records, config, expected=0):
  alphabase = len(config['alpha'])
  numbase = len(config['digit'])
  effective_namelen = config['namelen'] - 1
  name_part_bits = int(effective_namelen * math.log2(alphabase))
  number_part_bits = int(config['numberlen'] * math.log2(numbase))
  bitstr = '' # bit string storage
  for rec in records: # iterate over records
    nameval = from_base(rec[0][1:], alphabase, config['alpha'])
    numval = from_base(rec[1], numbase, config['digit'])
    bitstr += format(nameval, '0b').zfill(name_part_bits)
    bitstr += format(numval, '0b').zfill(number_part_bits)
  # reconstruct the raw data from the bitstring and return it
  datalen = int(math.ceil(len(bitstr)/8)) # estimate the data length in bytes
  if expected > 0: # truncate if expected length is specified
    datalen = expected
  data = b'' # raw data placeholder
  for i in range(0, datalen): # iterate over byte slices
    ind = i << 3
    data += int(bitstr[ind:ind+8].ljust(8, '0'), 2).to_bytes(1, 'big')
  return data

def auto_int(x): # helps to convert from any base natively supported in Python
    return int(x,0)

if __name__ == '__main__': # main app start
  from argparse import ArgumentParser
  parser = ArgumentParser(description='Databankr: Casio Databank/Telememo record format encoder/decoder', epilog='(c) Luxferre 2024 --- No rights reserved <https://unlicense.org>')
  parser.add_argument('mode', help='Operation mode (enc/dec)')
  parser.add_argument('-t', '--type', type=str, default='bin', help='Data type (bin/hex, default bin)')
  parser.add_argument('-i', '--input-file', type=str, default='-', help='Source input file (default "-", stdin)')
  parser.add_argument('-o', '--output-file', type=str, default='-', help='Result output file (default "-", stdout)')
  parser.add_argument('-c', '--config', type=str, default='config.json', help='Configuration JSON file path (default config.json in current working directory)')
  parser.add_argument('-m', '--module', type=str, default='2515-lat', help='Module configuration code according to your watch (default 2515-lat)')
  parser.add_argument('-l', '--expected-length', type=auto_int, default=0, help='Expected decoded data length in bytes (default 0 - no limits)')
  args = parser.parse_args()

  # detect the mode
  if args.mode == 'enc':
    flow = 'enc'
  elif args.mode == 'dec':
    flow = 'dec'
  else:
    print('Invalid mode! Please specify enc or dec!')
    exit(1)

  # load the configuration file
  try:
    f = open(args.config)
    confdata = json.load(f)
    f.close()
  except:
    print('Config file missing or invalid!')
    exit(1)

  # load the module config
  if args.module in confdata:
    moduleconfig = confdata[args.module]
    print('Loaded the configuration for %s' % moduleconfig['description'])
  else:
    print('Module configuration %s not found in the config file!' % args.module)
    exit(1)

  # load the input data
  try:
    if args.input_file == '-':
      infd = sys.stdin
    else:
      infd = open(args.input_file, mode='rb')
    indata = infd.read()
    if infd == sys.stdin: # perform additional conversion
      indata = indata.encode('utf-8')
    else:
      infd.close()
  except:
    print('Error reading the input data!')
    print(sys.exc_info())
    exit(1)

  # run the selected flow
  if flow == 'enc': # encoding flow
    if args.type == 'hex': # convert the input data if the type is hex
      indata = bytes.fromhex(re.sub(r"[^0-9a-fA-F]", "" ,indata.decode('utf-8')))
    records = encode(indata, moduleconfig)
    outdata = ''
    for rec in records: # separate each record with double newline
      outdata += rec[0] + '\n' + rec[1] + '\n\n'
  else: # decoding flow
    # parse the records
    rawrecs = indata.decode('utf-8').split('\n\n')
    records = [] # records will be stored here
    for pairstr in rawrecs:
      if len(pairstr) > 0: # exclude empty records
        pair = pairstr.split('\n') # get raw pair and then left-adjust
        records.append([pair[0].ljust(moduleconfig['namelen'], ' '),
                        pair[1].ljust(moduleconfig['numberlen'], ' ')])
    # decode the records
    outdata = decode(records, moduleconfig, args.expected_length)
    if args.type == 'hex': # convert the output data if the type is hex
      outdata = outdata.hex()
  # now, write the output file
  try:
    if args.output_file == '-':
      outfd = sys.stdout
      if type(outdata) == type(b''):
        outdata = outdata.decode('utf-8')
      outfd.write(outdata)
    else:
      outfd = open(args.output_file, mode='wb')
      if type(outdata) == type(''):
        outdata = outdata.encode('utf-8')
      outfd.write(outdata)
    if outfd != sys.stdout:
      outfd.close()
  except:
    print('Error writing the output file!')
    print(sys.exc_info())
    exit(1)

  print('\nOperation complete')