##############################################################################
#
# $RCSfile: adatfst.py,v $
#
# Copyright (C) 2006 John Popplewell
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
# 
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
# 
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
#
# Author : John Popplewell
# Email  : john@johnnypops.demon.co.uk
# Website: http://www.johnnypops.demon.co.uk/
#
# $Id: adatfst.py,v 1.2 2006/08/30 04:45:02 jfp Exp $
#
##############################################################################

_FILE_REVISION      = "$Revision: 1.2 $"

import struct
import itertools
import os, sys
import wave

if sys.platform == 'win32':
    import  win32disk
    open_device  = win32disk.File
else:
    open_device  = open

# File system constants
BYTES_PER_SAMPLE    =   3
STRING_LENGTH       =   64
SONG_INDEX_MULT     =   0x200
CLUSTER_MULT        =   0x200
BLOCK_SIZE          =   0x90000
NUM_LOCATIONS       =   25
SONG_NAME_OFFSET    =   0x03B8
START_OFFSET        =   0x0000
TOTAL_SONGS_OFFSET  =   0x0068
NUM_PROJECTS_OFFSET =   0x020C
DRIVE_NAME_OFFSET   =   0x03B8
PROJECTS_OFFSET     =   0x27F8

# Wave file headers
RIFF_HEADER_LEN     =   8
WAVE_HEADER_LEN     =   36

def _grouper(n, iterable, padvalue=None):
    "_grouper(3, 'abcdefg', 'x') --> ('a','b','c'), ('d','e','f'), ('g','x','x')"
    return itertools.izip(*[itertools.chain(iterable, itertools.repeat(padvalue, n-1))]*n)

def _packedstring(ps):
    s = ""
    for i,j,k,l in _grouper(4, ps, '\x00'):
        if not ord(l): break
        s += l
        if not ord(k): break
        s += k
        if not ord(j): break
        s += j
        if not ord(i): break
        s += i
    return s

def _packeddigits(ps):
    s = ""
    for i,j in _grouper(2, ps, '\x00'):
        if not ord(j): break
        s += j
        if not ord(i): break
        s += i
    return s

def _packed24bit(ps):
    return ps[0]+ps[1]*256+ps[2]*65536

class HD24Error(Exception):
    pass

class DummyWaveFile:
    def writeframesraw(self, data):
        pass
    def close(self):
        pass

def getTrackName(idx):
    return "Track%02d.wav"%(idx+1,)

def getSongName(idx, song):
    return "Song%02d - %s"%(idx+1, song['name'])

def getProjectName(idx, proj):
    return "Project%02d - %s"%(idx+1, proj['name'])

def getNumTracks(song, tracks):
    if tracks == None or not len(tracks):
        return song['num-tracks']
    return len(tracks)


class AdatFst:
    def __init__(self, filename):
        self.fp = open_device(filename, "rb")
        self.filename = filename
        self.info = self._readInfo()

    def _readInfo(self):
        info = {}
        self.fp.seek(START_OFFSET)
        fmt = struct.unpack("8s4s", self.fp.read(12))
        info['fs-type'] = _packedstring(fmt[0])
        if info['fs-type'] != 'ADAT FST':
            raise HD24Error("Not an 'ADAT FST' file-system.")
        info['fs-version'] = _packedstring(fmt[1])
        if info['fs-version'][:3] != '110':
            raise HD24Error("Incorrect version '%s'. Only 110 supported."%info['fs-version'])
        self.fp.seek(TOTAL_SONGS_OFFSET)
        fmt = struct.unpack("<1i", self.fp.read(4))
        info['total-songs'] = fmt[0]
        self.fp.seek(NUM_PROJECTS_OFFSET)
        fmt = struct.unpack("<1i", self.fp.read(4))
        info['num-projects'] = fmt[0]
        self.fp.seek(DRIVE_NAME_OFFSET)
        fmt = struct.unpack("64s", self.fp.read(STRING_LENGTH))
        info['drive-name'] = _packedstring(fmt[0])
        self.fp.seek(PROJECTS_OFFSET)
        projects = []
        for i in range(info['num-projects']):
            projects.append(self._readProjectInfo())
        info['projects'] = projects
        total_size = 0
        for proj in info['projects']:
            project_size = 0
            if proj['number-of-songs']:
                songs = []
                for i in proj['song-indexes']:
                    song = self._readSongInfo(i)
                    project_size += song['song-size']
                    songs.append(song)
                proj['songs'] = songs
            proj['project-size'] = project_size
            total_size += project_size
        info['total-size'] = total_size
        return info

    def _readProjectInfo(self):
        project = {}
        fmt = struct.unpack("8s8s2s2s", self.fp.read(20))
        project['tag'] = _packedstring(fmt[1]) + _packeddigits(fmt[2])
        fmt = struct.unpack("<i16s", self.fp.read(20))
        project['number-of-songs'] = nsongs = fmt[0]
        indexes = []
        for i in range(102):
            fmt = struct.unpack("<i", self.fp.read(4))
            if i < nsongs:
                indexes.append(fmt[0])
        project['song-indexes'] = indexes
        fmt = struct.unpack("64s", self.fp.read(STRING_LENGTH))
        project['name'] = _packedstring(fmt[0])
        return project

    def _readSongInfo(self, index):
        song = {}
        self.fp.seek(index*SONG_INDEX_MULT)
        fmt = struct.unpack("<40s10s1B2s3B1I124s", self.fp.read(184))
        song['tag'] = _packedstring(fmt[1])
        song['num-tracks'] = fmt[2]
        song['sample-rate'] = _packed24bit(fmt[4:7])
        song['num-samples'] = fmt[7]
        song['file-size'] = (((song['num-samples']*BYTES_PER_SAMPLE) + 1) & ~1) + RIFF_HEADER_LEN + WAVE_HEADER_LEN
        song['song-size'] = song['file-size']*song['num-tracks']
        locs = []
        for i in range(NUM_LOCATIONS):
            fmt = struct.unpack("<1I8s", self.fp.read(12))
            locs.append((_packedstring(fmt[1]), fmt[0]))
        song['locations'] = locs
        self.fp.seek((index*SONG_INDEX_MULT)+SONG_NAME_OFFSET)
        fmt = struct.unpack("64s", self.fp.read(STRING_LENGTH))
        song['name'] = _packedstring(fmt[0])
        self.fp.read(8)
        clusters = []
        while 1:
            fmt = struct.unpack("<Ii", self.fp.read(8))
            if not fmt[0]:
                break
            clusters.append(fmt)
        song['clusters'] = clusters
        return song

    def writeTracks(self, outdir, idx, song, tracks, spinner=None):
        fps = []
        for i in range(song['num-tracks']):
            if tracks == None or not len(tracks) or i in tracks:
                f = wave.open(os.path.join(outdir, getTrackName(i)), "wb")
                f.setnchannels(1)
                f.setsampwidth(BYTES_PER_SAMPLE)
                f.setframerate(song['sample-rate'])
                f.setnframes(0)
            else:
                f = DummyWaveFile()
            fps.append(f)

        totalsamples = datalen = song['num-samples']*BYTES_PER_SAMPLE
        ntracks = song['num-tracks']
        chunk = BLOCK_SIZE/ntracks
        clusters = song['clusters']
        cluster_idx = cluster_cnt = 0
        while datalen:
            if spinner: spinner.tick(totalsamples, totalsamples-datalen)
            if not cluster_cnt:
                self.fp.seek(clusters[cluster_idx][0]*CLUSTER_MULT)
                cluster_cnt = clusters[cluster_idx][1]
                cluster_idx += 1

            data = self.fp.read(BLOCK_SIZE)
            cluster_cnt -= 1
                
            if datalen > chunk: blk = chunk
            else: blk = datalen
            for i in range(ntracks):
                offset = i*chunk
                fps[i].writeframesraw(data[offset:offset+blk])
            datalen -= blk

        for f in fps: f.close()

        if totalsamples & 1:
            # work-round 'bug' in wave module:
            #   make the file length even and fix the RIFF header
            #   keeping the data length the same
            for i in range(song['num-tracks']):
                if tracks == None or not len(tracks) or i in tracks:
                    f = open(os.path.join(outdir, getTrackName(i)), "r+b")
                    f.seek(0x0004)
                    f.write(struct.pack("<I",totalsamples+1+WAVE_HEADER_LEN))
                    f.seek(0, 2)
                    f.write("\x00")
                    f.close()
        if spinner: spinner.tick(totalsamples, totalsamples)

    def getWavFile(self, song, track):
        class WavFile:
            def __init__(self, fp, song, track):
                self.fp = fp
                self.song = song
                self.header = self._make_header()
                self.pos = 0
                self.datalen = song['num-samples']*BYTES_PER_SAMPLE
                self.chunk = BLOCK_SIZE/song['num-tracks']
                self.clusters = song['clusters']
                self.track_start = self.chunk*track
                self.track_end = self.track_start+self.chunk

            def _make_header(self):
                header = 'RIFF'
                nframes = self.song['num-samples']
                nchannels = 1
                framerate = self.song['sample-rate']
                sampwidth = BYTES_PER_SAMPLE
                datalength = self.song['num-samples']*BYTES_PER_SAMPLE*nchannels
                datalength_pad = (datalength + 1) & ~1
                header += struct.pack('<l4s4slhhllhh4sl',
                    WAVE_HEADER_LEN + datalength_pad, 'WAVE', 'fmt ', 16,
                    1, nchannels, framerate,
                    nchannels * framerate * sampwidth,
                    nchannels * sampwidth,
                    sampwidth * 8, 'data', datalength)
                return header

            def close(self):
                self.fp.close()

            def read(self, size):
                data = ''
                pos = self.pos - len(self.header)
                if pos < 0:
                    data = self.header[self.pos:self.pos+size]
                    size -= len(data)
                    pos = 0
                if size:
                    chunk_start = pos/self.chunk
                    chunk_end   = (pos + size + self.chunk-1)/self.chunk
                    block = self._read_chunks(chunk_start, chunk_end - chunk_start)
                    pos -= chunk_start*self.chunk
                    data += block[pos:pos+size]
                self.pos += len(data)
                return data

            def _read_chunks(self, index, n):
                for cluster_idx, (pos, cluster_cnt) in enumerate(self.clusters):
                    if index < cluster_cnt:
                        self.fp.seek(pos*CLUSTER_MULT + index*BLOCK_SIZE)
                        cluster_cnt -= index
                        cluster_idx += 1
                        break
                    index -= cluster_cnt

                buff = []
                for i in range(n):
                    if not cluster_cnt:
                        self.fp.seek(self.clusters[cluster_idx][0]*CLUSTER_MULT)
                        cluster_cnt = self.clusters[cluster_idx][1]
                        cluster_idx += 1
                    buff.append(self.fp.read(BLOCK_SIZE)[self.track_start:self.track_end])
                    cluster_cnt -= 1
                return ''.join(buff)

            def seek(self, where, whence=0):
                if whence == 0:
                    self.pos = where
                elif whence == 1:
                    self.pos += where
                elif whence == 2:
                    self.pos = self.song['file-size'] - where
                else:
                    raise IOError('[Errno 22] Invalid argument')

            def tell(self):
                return self.pos

        return WavFile(open_device(self.filename, "rb"), song, track)