# -*- coding: utf-8 -*-
import os, sys
from math import ceil

def cpm_bdos_dir(dpb, im_file):
  """ Lists the contents of a CP/M disk using the disk parameters from the
      provided disk parameter block and the disk image loaded from im_file """
  fr = file(im_file, "rb")
  fr.seek(dpb['spt'] * 128 * dpb['off'])

  files = []
  for i in range(dpb['drm']):
    UU = fr.read(1)                     # user id 0-15 on most systems
    FN = fr.read(8)                     # file name
    FT = fr.read(3)                     # file type e.g. COM, TXT
    EX = fr.read(1)                     # file extent, for long files
    S2 = fr.read(1)
    S1 = fr.read(1)                     # reserved as 0
    RC = fr.read(1)                     # number of records
    AL = fr.read(16)                    # allocated records
    ALL = []
    if dpb['dsm'] > 255:
      # must be 16 bit block numbers
      for i in range(8):
        ALL.append(ord(AL[i*2]) + 256 * ord(AL[i*2+1]))
      if UU != "\xe5":
        print repr(AL)
        print ALL
    else:
      # 8 bit block numbers
      for i in range(16):
        ALL.append(ord(AL[i]))
    files.append({"UU": UU, "FN": FN, "FT": FT, "EX": ord(EX), "S2": ord(S2), "RC": ord(RC), "AL": ALL})
  fr.close()
  return files

def cpm_dir(dpb, im_file):
  files = cpm_bdos_dir(dpb, im_file)

  col = 0
  for f in files:
    if(f["UU"] != "\xE5"):                   # test for deleted file
      if ((32*f["S2"]) + f["EX"]) / (dpb['exm'] + 1) == 0:
        # if this isn't zero then it's a later entry
        if col % 4 == 0:
          print "A:",
        print f["FN"].upper() + " " + f["FT"].upper(),
        col += 1
        if col % 4 == 0:
          print " "
        else:
          print ":",
  print " "
      
def cpm_pip(cmd, dpb, im_file):
  if cmd == "PIP":
    # no options so go into interactive mode
    while True:
      a = raw_input("*")
      if a == "":
        return
      cpm_pip("PIP " + a.upper(), dpb, im_file)
  else:
    if cmd[4:8] == "PUN:":
      # send something to the "punch" device.  In this simulation that means
      # save the contents to the host filesystem
      b = cmd.split("=")[1]
      b = b.split(",")
      # open a file with the same name as the first argument
      if os.path.isfile(b[0]):
        a = raw_input("Overwrite old file? (Y/N) ")
        if a != "Y":
          return
      fw = file(b[0], "wb")
      for e in b:
        cat(dpb, im_file, e, fw)
      fw.close()

def cat(dpb, im_file, fn, fw):
  # find the file
  blocks = cpm_bdos_search(dpb, im_file, fn)
  fr = file(im_file, "rb")
  for block in blocks:
    fr.seek((block + dpb['off']) * dpb['spt'] * 128)
    fw.write(fr.read(dpb['spt'] * 128))
  fr.close()
  return

def cpm_bdos_search(dpb, im_file, fn):
  files = cpm_bdos_dir(dpb, im_file)
  parts = []
  for f in files:
    if f['FN'].strip() + "." + f['FT'].strip() == fn:
      parts.append(f)

  print parts

  # need to make an ordered list of the blocks for that file now
  blocks_part = [0] * len(parts)
  for part in parts:
    blocks_part[(32 * part['S2'] + part['EX']) / (dpb['exm'] + 1)] = part['AL'][0:int(ceil(((part['EX'] & dpb['exm']) * 128 + part['RC'])/32))]
  blocks = []
  blocks_part.sort()
  for p in blocks_part:
    blocks.extend(p)

  print blocks
  return blocks
  

if __name__ == "__main__":
  dpb = {}
  dpb['spt'] = 0x0020                   # 128 byte sectors ber track
  dpb['bsh'] = 0x05                     # block shift; log2(tracklen/128)
  dpb['blm'] = 0x1F                     # block mask; tracklen/128 - 1
  dpb['exm'] = 0x01                     # extent mask
  dpb['dsm'] = 0x07F9                   # blocks on disk - 1
  dpb['drm'] = 0x03ff                   # directory entries - 1
  dpb['al0'] = 0xFF                     # directory alloc bitmap byte 1
  dpb['al1'] = 0x00                     # directory alloc bitmap byte 2
  dpb['cks'] = 0x0000                   # checksum (0 for a fixed disk)
  dpb['off'] = 0x0006                   # offset to directory

  ## Note on al0 al1:
  ## They form a bitmask of tracks following the offset that are used for the
  ## root directory, so total can be up to 16 tracks
  ## in this example 8 tracks are used (8 * 4096 = 32768; 32768/32 = 1024 which
  ## is 0x400, i.e. number of directory entries.)

  ## reserved tracks are used for booting the system.

  if len(sys.argv) < 2:
    print "Usage: %s <diskimg>"
    sys.exit(1)

  while True:
    # do a command interpreter
    a = raw_input("A>")
    a = a.upper()
    if a == "DIR":
      cpm_dir(dpb, sys.argv[1])
    elif a[0:3] == "PIP":
      cpm_pip(a, dpb, sys.argv[1])
    elif a == "EXIT":
      sys.exit(0)
    else:
      print a + "?\n"
 