#!/usr/bin/python
#
# bakeDisk64.py v1.0 by Christopher Jam 2013-2016
#
# generates c64 disk images from .prg files and raw blocks
#
# Requires Numpy, tested with Python 2.7.10 and 3.5.2
#
# for usage:
# python bakeDisk64.py -h
#
from __future__ import print_function
from __future__ import division
import numpy as np
import sys
import os

DEFAULT_INTERLEAVE=5

g_interleave=DEFAULT_INTERLEAVE
g_ID="ID"

def tracks():
    return range(1,36)

def sectorsInTrack(track):
    if(track<18): return 21
    if(track<25): return 19
    if(track<31): return 18
    return 17

def sectors():
    for tk in tracks():
        for se in range(sectorsInTrack(tk)):
            yield tk,se

filetypecodes= {
    'SEQ':0x81,
    'PRG':0x82,
}
def petscii(c):
    v=ord(c)
    if(97<=v<=97+26):
        v-=32
    elif(65<=v<=65+26):
        v+=193-65
    return v

class CDirEnt:
    def __init__(self, tk,se,fnam,typ):
        d=np.zeros(30,np.uint8)
        d[0]=filetypecodes[typ]
        d[1:3]=tk,se
        d[3:19]=(list(map(petscii,fnam))+[160,]*16)[:16]
        self.d=d
    def setSize(self, nblocks):
        self.d[28]=nblocks%256
        self.d[29]=nblocks/256
    def data(self):
        return self.d
    def tk(self):
        return self.d[1]
    def se(self):
        return self.d[2]


class CDisk:
    def __init__(self, disknam):

        disknam=(disknam+chr(160)*16)[:16]

        self.blocks={}
        self.bam={}
        for (tk,se) in sectors():
            self.blocks[(tk,se)]=np.zeros(256, np.uint8)

        self.blocksFree={}
        for tk in tracks():
            self.blocksFree[tk]=sectorsInTrack(tk)
            for se in range(24):
                self.bam[(tk,se)]=(se<sectorsInTrack(tk))

        block0=self.blocks[(18,0)]
        block0[0:4]=18,1,65,0
        block0[144:160]=(list(map(petscii,disknam))+[160,]*16)[:16]
        block0[160:162]=160,160
        block0[162:164]=list(map(ord,g_ID))
        block0[164:171]=160,50,65, 160,160,160,160
        self.dirents=[]

    def getFreeBlockNear(self,tk,se, offset=None):
        if offset==None:
            offset=g_interleave
        while(tk==18 or self.blocksFree[tk]==0):
            tk=tk+(tk>18)*2-1
            if(tk==0):
                tk=19
            se=0
            assert(tk in tracks())
        ntracks=sectorsInTrack(tk)

        if(self.bam[(tk,se)]!=1):   # '1' means free
            se=(se+ntracks+offset)%ntracks
            while(self.bam[(tk,se)]!=1):   # '1' means free
                se=(se+1)%ntracks
        self.markUsed(tk,se)
        return tk,se
    def markUsed(self,tk,se):
        if self.bam[(tk,se)]==0:
            print("block %d,%d already in use!"%(tk,se))
            sys.exit(1)
        self.bam[(tk,se)]=0
        self.blocksFree[tk]-=1
        

    def reserveStartingBlock(self, fs):
        self.markUsed(fs.tk,fs.se)

    def addFile(self, fs):
        path=fs.path
        fnam=fs.fnam
        fdata=np.fromstring(open(path,'rb').read(),np.uint8)
        tk=fs.tk
        se=fs.se
        interleave=fs.interleave

        if interleave==None:
            interleave=g_interleave

        if tk==None:
            (tk,se)=self.getFreeBlockNear(18,1, offset=1)
        nblocks=1

        print("starting file %s at %d,%d (%d blocks, interleave %d)"%(fnam,tk,se, (len(fdata)+253)//254, interleave))
        self.dirents.append(CDirEnt(tk,se,fnam,'PRG'))
        while(len(fdata)>254):
            self.blocks[(tk,se)][2:]=fdata[:254]
            fdata=fdata[254:]
            ntk,nse=self.getFreeBlockNear(tk,se,interleave)
            nblocks+=1
            self.blocks[(tk,se)][:2]=ntk,nse
            tk,se=ntk,nse
        self.blocks[(tk,se)][2:2+len(fdata)]=fdata
        self.blocks[(tk,se)][:2]=0,1+len(fdata)
        self.dirents[-1].setSize(nblocks)

    def writeBlock(self, fn,data,tk,se):
        data=np.fromstring(data,np.uint8)
        nb=len(data)
        print("writing %3d bytes from %15s to %d,%d"%(nb,fn,tk,se))
        assert(nb<257)
        assert(self.bam[(tk,se)]==1)
        self.blocks[(tk,se)][:nb]=data

        self.bam[(tk,se)]=0
        self.blocksFree[tk]-=1
            
            

    def updateBAM(self):
        self.markUsed(18,0)
        block0=self.blocks[(18,0)]
        fb=0
        for tk in tracks():
            block0[tk*4+0]=self.blocksFree[tk]
            block0[tk*4+1]=np.add.reduce([ (2**x)*self.bam[tk, 0+x] for x in range(8)])
            block0[tk*4+2]=np.add.reduce([ (2**x)*self.bam[tk, 8+x] for x in range(8)])
            block0[tk*4+3]=np.add.reduce([ (2**x)*self.bam[tk,16+x] for x in range(8)])
            if(tk!=18):
                fb+=self.blocksFree[tk]
        print(fb,"blocks free")

    def updateDir(self):
        pos=2
        tk,se=18,1
        print("adding %d,%d to directory"%(tk,se))
        cblock=self.blocks[(tk,se)]
        self.markUsed(tk,se)
        cblock[0:2]=0,255    #mark as last block - fix this with ref to next if such exists
        for de in self.dirents:
            #print "updating dir ent at %d pointing to block %d,%d" %(pos,de.tk(),de.se())
            cblock[pos:pos+30]=de.data()
            pos+=32
            if(pos>256):
                pos=2
                se=se+3
                cblock[0:2]=tk,se
                print("adding %d,%d to directory"%(tk,se))
                self.markUsed(tk,se)
                cblock=self.blocks[(tk,se)]
                cblock[0:2]=0,255
            
        if 0:
            for i in range(0,256,32):
                sd=cblock[i:i+32]
                for j in sd:
                    print("%02x"%(j,),end=' ')
                print("".join(map( lambda c: ['.',chr(c)][31<c<127], sd)))

    def writeD64(self,fnam):
        def dumpBlock(tk,se):
            print("block %2d.%02d:"%(tk,se))
            for i in range(8):
                bd=self.blocks[(tk,se)][i*32:][:32]
                print("    %s"%(" ".join(['%02x'%x for x in bd])))
        self.updateDir()
        self.updateBAM()
        fo=open(fnam,'wb')
        for (tk,se) in sectors():
            fo.write(self.blocks[(tk,se)].astype(np.int8).tostring())
        fo.close()
        (tk,se)=self.getFreeBlockNear(18,1)
        print("lowest track = ",tk)



def parseD2OrDie(digits, context=""):
    if len(context)>0:context=context+": "
    if not (len(digits) in (1,2)): usage(context+"Expected 1-2 digit number, found '%s'"%digits)
    if not (digits[0] in "0123456789"): usage(context+"Expected decimal number, found '%s'"%digits)
    if not (digits[-1] in "0123456789"): usage(context+"Expected decimal number, found '%s'"%digits)
    return int(digits)

class FileSpec:
    def __init__(self,path,fnam=None,tk=None,se=None,interleave=None):
        if fnam==None:
            fnam=os.path.basename(path)
            if(fnam.lower().endswith('.prg')):
                fnam=fnam[:-4]
        self.path=path
        self.fnam=fnam
        self.tk=tk
        self.se=se
        self.interleave=interleave

    @classmethod
    def fromArg(self,arg, seperator=","):
        argl=arg.split(seperator)
        if len(argl)>5: usage("expected no more than five '%s' seperated items in filespec, found %s"%(seperator,arg))
        path=argl[0]
        fnam=None
        tk=se=None
        interleave=None
        if len(argl)>1 and len(argl[1])>0:
            fnam=argl[1]
        if len(argl)==3:
            interleave=parseD2OrDie(argl[2], arg)
        if len(argl)>3:
            if len(argl[2])>0 and len(argl[3])>=0:
                tk=parseD2OrDie(argl[2], arg)
                se=parseD2OrDie(argl[3], arg)
        if len(argl)==5:
            interleave=parseD2OrDie(argl[4], arg)
        return FileSpec(path,fnam,tk,se,interleave)


def usage(err=None):
    if err!=None:
        sys.stderr.write("\nError: %s\n\n"%err)
    sys.stderr.write("""Usage:
    %s -o filename.d64 -n diskname [-l interleave] [-i ID] [ sourcefile.prg[,[c64filename][,tk,se|,,][,interleave] ] ]* [-wb block.bin,tk,se]*

    For each source file, you can optionally specify as many as you like of
    - c64 filename
    - starting track,sector (specify either both or neither)
    - desired block interleave

    You can also specify a default interleave for files that don't have one specified, otherwise the default will be %d
    
    d64 name and diskname are required.

    block.bin files should be at most 256 bytes long (they'll be zero padded if under)
    Blocks written will be marked as allocated in the BAM

    eg:
    %s -o testdisk.d64 -n my\ disk -l 5  boot.prg,,17,0 part1.prg,1part,4 -wb loader.bin,17,1

    boot.prg will be written as "boot", starting at 17,0, with interleave of 5
    loader.bin will be placed in 17,1
    part1.prg will be written as "1part", starting at first available block, with interleave of 4\n\n"""% (sys.argv[0], DEFAULT_INTERLEAVE, sys.argv[0]))
    sys.exit(1)

files=[]
blockWrites=[]
dnam=None
fnam=None
args=iter(sys.argv[1:])

for a in args:
    if   a=='-o':  fnam=next(args)
    elif a=='-n':  dnam=next(args)
    elif a=='-h':  usage()
    elif a=='-i':  g_ID=next(args)
    elif a=='-l':  g_interleave=parseD2OrDie(next(args))
    elif a=='-wb': blockWrites.append(next(args).split(','))
    else:         files.append(FileSpec.fromArg(a))

if fnam==None: usage("No output filename specified")
if dnam==None: usage("No diskname specified")


d=CDisk(dnam)
for fn,tk,se in blockWrites:
    d.writeBlock(fn,open(fn,'rb').read(),int(tk),int(se))

for fs in files:
    if fs.tk!=None:
        d.reserveStartingBlock(fs)

for fs in files:
    d.addFile(fs)

d.writeD64(fnam)