#
# utility for use teledisk images
# put sectors in 'logical' order
#

class stream:
    def __init__(self, name, mode='rb'):
        self._f = open(name, mode)
        self._crc = 0
    def crc0(self):
        self._crc = 0
    def byte(self):
        return ord(self._f.read(1))
    def word(self):
        return ord(self._f.read(1))+ord(self._f.read(1))*256
    def rle_decode(self, debug=False):
        _data = ''
        _len = self.word()
        while (_len > 0):
            _code = self.byte()
            _len -= 1
            if (debug):
                print 'RLE decode : code %02x len %d' % (_code, _len,)
            if (_code == 0):
                while (_len > 0):
                    _data += chr(self.byte())
                    _len -= 1
            elif (_code == 1):
                _len1 = self.word()
                _len -= 2
                _d1 = self.byte()
                _len -= 1
                _d2 = self.byte()
                _len -= 1
                while (_len1 > 0):
                    _data += chr(_d1)
                    _data += chr(_d2)
                    _len1 -= 1
            elif (_code == 2):
                while (_len > 0):
                    _code2 = self.byte()
                    _len -= 1;
                    _len2 = self.byte()
                    _len -= 1;
                    if (debug):
                        print '               subcode %02x sublen %d (len %d)' % (_code2, _len2, _len)
                    if (_code2 == 0):
                        while (_len2 > 0):
                            _data += chr(self.byte())
                            _len2 -= 1
                            _len -= 1
                    elif (_code2 == 1):
                        _d1 = self.byte()
                        _len -= 1
                        _d2 = self.byte()
                        _len -= 1
                        while (_len2 > 0):
                            _data += chr(_d1)
                            _data += chr(_d2)
                            _len2 -= 1
                    else:
                        print 'RLE decode : strange subcode : %02x' % _code2
            else:
                print 'RLE decode : strange code : %02x' % _code
            if (debug):
                print '             len : %d' % _len
        return _data[:]
    def close(self):
        self._f.close()

class disk:
    """A class to deal with disk images"""
    def __init__(self, name=None, debug=False, dump=False):
        if name == None:
            self.header = []
            self.sectors = 16
            self.heads = 2
            self.cylinders = 33
            self.size = 256
            self._disk = {}
            for i in range(self.cylinders):
                for j in range(self.heads):
                    for k in range(self.sectors):
                        ref = (i,j,k)
                        self._disk[ref] = chr(255)*self.size
        elif name[-4:] == '.TD0':
            self._readtd(name, debug, dump)
        elif name[-4:] == '.IMG':
            pass
    def show(self):
        _h1 = '0123456789'*26
        _h2 = ''
        for i in range(26):
            if (i < 10):
                _h2 += '%s         ' % i
            else:
                _h2 += '%s        ' % i
        maps = '  cyl %s\nsec h %s\n' % (_h2[:self.cylinders], _h1[:self.cylinders],)  
        for _sect in range(self.sectors):
            for _head in range(self.heads):
                _maps = ''
                for _cyl in range(self.cylinders):
                    _ref = (_cyl, _head, _sect)
                    if (_ref in self._disk.keys()):
                        if (len(self._disk[_ref]) == 256):
                            _d = '='
                        elif (len(self._disk[_ref]) < 256):
                            _d = '-'
                        else:
                            _d = '+'
                    else:
                        _d = ' '
                    _maps += _d
                if (_head == 0):
                    maps += '%03d %d %s\n' % (_sect, _head, _maps)
                else:
                    maps += '    %d %s\n' % (_head, _maps,)
        return maps[:]
    def _readtd(self, name, debug=False, dump=False):
        self.header = []
        self.sectors = 16
        self.heads = 2
        self.cylinders = 33
        self.size = 256
        self._disk = {}
        f = stream(name, 'rb')
        # header
        _hdr_TXT = "%c%c" % (f.byte(), f.byte())
        _hdr_SeqVal = f.byte()
        _hdr_ChkSig = f.byte()
        _hdr_TDVer = f.byte()
        _hdr_Dens = f.byte()
        _hdr_DrvType = f.byte()
        _hdr_TrkDens = f.byte()
        _hdr_DosMode = f.byte()
        _hdr_Surface = f.byte()
        _hdr_CRC = f.word()
        if (_hdr_TrkDens & 0x80):
            _hdr_Com_Crc = f.word()
            _hdr_Com_Len = f.word()
            _hdr_Com_YMD = (f.byte(), f.byte(), f.byte())
            _hdr_Com_HMS = (f.byte(), f.byte(), f.byte())
            _hdr_Com_Com = ''
            for i in range(_hdr_Com_Len):
                _hdr_Com_Com += chr(f.byte())
            _hdr_Comment = '%s (Crc: %04x Date: %d/%d/%d, %d:%02d:%02d)' % (_hdr_Com_Com, _hdr_Com_Crc, _hdr_Com_YMD[2], _hdr_Com_YMD[1], _hdr_Com_YMD[0]+1900, _hdr_Com_HMS[0], _hdr_Com_HMS[1], _hdr_Com_HMS[2])
        else:
            _hdr_Comment = 'None'
        print 'HEADER'
        print '  TXT: %2s, SeqVal: %02x, ChkSig: %02x, TDVer: %02x, Dens: %02x' % (_hdr_TXT, _hdr_SeqVal, _hdr_ChkSig, _hdr_TDVer, _hdr_Dens)
        print '  DrvType: %02x, TrkDens: %02x, DosMode: %02x, Surface: %02x, CRC: %04x' % (_hdr_DrvType, _hdr_TrkDens, _hdr_DosMode, _hdr_Surface, _hdr_CRC)
        print '  COM: %s' % _hdr_Comment
        # tracks (cylinders)
        do_tracks = True
        while (do_tracks):
            _trk_SecPerTrk = f.byte()
            _trk_PhisCyl = f.byte()
            _trk_PhisSide = f.byte()
            _trk_Crc = f.byte()
            if (debug):
                print 'TRACK: %d, Side: %d, %d sectors (Crc:%02x)' % (_trk_PhisCyl, _trk_PhisSide, _trk_SecPerTrk, _trk_Crc)
            if (_trk_SecPerTrk == 0xFF):
                do_tracks = False
            if (do_tracks):
                for i in range(_trk_SecPerTrk):
                    _sec_Cyl = f.byte()
                    _sec_Side = f.byte()
                    _sec_SNum = f.byte()
                    _sec_SLen = (2**(f.byte()))*128
                    _sec_Syndrome = f.byte()
                    _sec_Crc = f.byte()
                    _sec_DataP = True
                    if ((_sec_Cyl != _trk_PhisCyl) or (_sec_Side != _trk_PhisSide)):
                        print 'TRACK  %d SIDE %d : sector for track %d side %d' % (_trk_PhisCyl, _trk_PhisSide, _sec_Cyl, _sec_Side)
                    if (debug):
                        print'  SECTOR: %02d (Cyl: %d, Side: %d, Len: %d) Syndrome: %02x (Crc: %02x)' % (_sec_SNum, _sec_Cyl, _sec_Side, _sec_SLen, _sec_Syndrome, _sec_Crc)
                    if (_sec_Syndrome & 1):
                        if (debug):
                            print '    more than once sector number appear'
                    if (_sec_Syndrome & 2):
                        if (debug):
                            print '    crc error'
                    if (_sec_Syndrome & 4):
                        if (debug):
                            print '    deleted control mark'
                    if (_sec_Syndrome & 16):
                        if (debug):
                            print '    DOS sector copy -> no data'
                        _sec_DataP = False
                    if (_sec_Syndrome & 32):
                        if (debug):
                            print '    data field missing -> no data'
                        _sec_DataP = False
                    if (_sec_Syndrome & 64):
                        if (debug):
                            print '    no ID address present'
                    if (_sec_DataP):
                        _sec_Data = ''
                        _sec_Data_Hex = ''
                        _sec_Data_Chr = ''
                        # _rle_len = f.word()
                        data = f.rle_decode()
                        if (dump):
                            print 'RLE decode : len %d' % len(data)
                            for d in data:
                                _sec_Data += d
                                _sec_Data_Hex += '%02x ' % ord(d)
                                if ((ord(d) >= 32) and (ord(d) < 128)):
                                    _sec_Data_Chr += d
                                else:
                                    _sec_Data_Chr += '.'
                                for k in range(len(_sec_Data) / 16):
                                    print '    %04x %s %s' % (k*16, _sec_Data_Hex[k*48:k*48+48], _sec_Data_Chr[k*16:k*16+16])
                        ref = (_sec_Cyl, _sec_Side, _sec_SNum)
                        self._disk[ref] = data[:]

        f.close()
        _max_cyl = 0
        _max_head = 0
        _max_sect = 0
        for _ref in self._disk.keys():
            _max_cyl = max(_max_cyl, _ref[0])
            _max_head = max(_max_head, _ref[1])
            _max_sect = max(_max_sect, _ref[2])
        self.cylinders = _max_cyl+1
        self.heads = _max_head+1
        self.sectors = _max_sect+1
    def _readimg(self, name):
        pass
    def dump_img0(self, name):
        #
        # dump from an 70 cyl, 1 side, 16 sectors
        #
        _g = open(name, 'wb')
        for _cyl in range(66):
            for _sect in range(16):
                if ((_cyl, 0, _sect) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 0, _sect)])
        _g.close()
        print 'done'
    def dump_img1(self, name):
        #
        # dump from an 35 cyl, 2 side, 16 sectors
        #
        _g = open(name, 'wb')
        for _cyl in range(33):
            for _sect in range(16):
                if ((_cyl, 0, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 0, _sect+1)])
            for _sect in range(16):
                if ((_cyl, 1, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 1, _sect+1)])
        _g.close()
        print 'done'
    def dump_img2(self, name):
        #
        # dump from an 70 cyl, 1 side, 16 sectors from sector 1
        #
        _g = open(name, 'wb')
        for _cyl in range(66):
            for _sect in range(16):
                if ((_cyl, 0, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 0, _sect+1)])
        _g.close()
        print 'done'
    def dump_img3(self, name):
        #
        # dump from an 77 cyl, 2 sides, 16 256b sectors from sector 1
        #
        _g = open(name, 'wb')
        for _cyl in range(77):
            for _sect in range(16):
                if ((_cyl, 0, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 0, _sect+1)])
            for _sect in range(16):
                if ((_cyl, 1, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 1, _sect+1)])
        _g.close()
        print 'done'
    def dump_img4(self, name):
        #
        # dump from an 77 cyl, 2 sides, 16 256b sectors from sector 1
        # only the first 264 kb
        _g = open(name, 'wb')
        for _cyl in range(33):
            for _sect in range(16):
                if ((_cyl, 0, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 0, _sect+1)])
            for _sect in range(16):
                if ((_cyl, 1, _sect+1) not in self._disk.keys()):
                    _g.write(chr(0)*256)
                else:
                    _g.write(self._disk[(_cyl, 1, _sect+1)])
        _g.close()
        print 'done'
#
# end
#
name='bas2utl2'
vc = disk(name+'.TD0', debug=False)
# vc = disk('WORDSTAR.TD0', debug=False, dump=False)
m = vc.show()
print m
vc.dump_img3(name+'.BIN')
