#!/bin/env python

##############################################################################
#
# $RCSfile: hd24.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: hd24.py,v 1.33 2007/12/18 12:46:32 jfp Exp $
#
##############################################################################

_FILE_REVISION      = "$Revision: 1.33 $"

import os, sys
import time

from adatfst import AdatFst, getProjectName, getSongName, getNumTracks, HD24Error, CLUSTER_MULT

try:
    set, frozenset
except NameError:
    from sets import Set as set, ImmutableSet as frozenset

from optparse import OptionParser, OptionValueError

VERSION_NUMBER      = "0.0.1"
LICENCE_STRING      = "Copyright (C) 2006 John Popplewell"
VERSION_STRING      = "Version %s\n%s"%(VERSION_NUMBER, LICENCE_STRING)
DESCRIPTION_STRING  = """This program is used to display and extract projects and
songs directly from a hard-drive used by the Alesis ADAT HD24 digital recorder.
"""

if sys.platform == 'win32':
    USAGE_STRING = "usage: %prog [options] \\\\.\\PHYSICALDRIVE[0-3]"
else:
    USAGE_STRING = "usage: %prog [options] /dev/[hs]d[a-z] or a binary image"

def _safe_mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)

def _human_readable(n):
    if n < 1024: return "%dB"%n
    if n < 1024*1024: return "%1.1fKB"%(float(n)/1024)
    if n < 1024*1024*1024: return "%1.1fMB"%(float(n)/(1024*1024))
    return "%1.1fGB"%(float(n)/(1024*1024*1024))

def _s(n):
    if n != 1: return "s"
    return ""


class NullSpinner:
    def __init__(self):
        pass
    def message(self, msg):
        pass
    def tick(self, total, current):
        pass

class Spinner:
    def __init__(self, width):
        self.width = width
        self.length_prev = -1
        self.percent_prev = -1

    def message(self, msg):
        print msg

    def tick(self, total, current):
        if total == current:
            sys.stdout.write("\r%s\r"%(" "*self.width))
            sys.stdout.flush()
            return
        width = self.width-9
        length = width*current/total
        percent = 100*current/total
        if percent == self.percent_prev and length == self.length_prev:
            return
        self.length_prev = length
        self.percent_prev = percent
        sys.stdout.write("\r  %-4s[%s]"%("%d%%"%percent, (("="*length)+">").ljust(width)))
        sys.stdout.flush()


class AdatFstUtils(AdatFst):
    def __init__(self, filename, spinner=NullSpinner()):
        AdatFst.__init__(self, filename)
        self.create_all = False
        self.overwrite_all = False
        self.spinner = spinner

    def _showSong(self, verbose, idx, song, tracks=None):
        if verbose:
            print "    name=%-34s %2dtrk tag='%s'"%(
                getSongName(idx, song), song['num-tracks'], song['tag'])
            print "      %(sample-rate)5dHz total-samples=%(num-samples)d"%song,
            if song['num-samples']:
                print "file-size=%9d (%7s)"%(song['file-size'], _human_readable(song['file-size']))
                print "      clusters:"
                for start, length in song['clusters']:
                    print "        start=%#010x length=%#04x"%(start*CLUSTER_MULT, length)
                print "      Song size = %6s"%(_human_readable(song['song-size']),)
            else:
                print
        else:
            print "    name=%-34s %2dtrk, %5dHz"%(
                getSongName(idx, song), song['num-tracks'], song['sample-rate']),
            if song['num-samples']:
                print "file-size=%7s song-size=%7s"%(_human_readable(song['file-size']),
                                                     _human_readable(song['song-size']))
            else:
                print
        return getNumTracks(song, tracks)*song['file-size']

    def _showProject(self, verbose, idx, proj):
        if verbose:
            print "  name='%s' tag='%s'\n  %d songs:"%(
                        getProjectName(idx, proj), proj['tag'], proj['number-of-songs'])
        else:
            print "  name='%s'  %d songs:"%(
                        getProjectName(idx, proj), proj['number-of-songs'])
        if not proj['number-of-songs']:
            return
        for i, song in enumerate(proj['songs']):
            self._showSong(verbose, i, song)
        print "    Project size = %s"%_human_readable(proj['project-size'])

    def _showInfo(self, verbose):
        if verbose:
            print "DRIVE:\n  fs-type='%s' fs-version='%s'\n  drive-name='%s' num-projects=%d total-songs=%d\n" % (
                        self.info['fs-type'], self.info['fs-version'], self.info['drive-name'], 
                        self.info['num-projects'], self.info['total-songs'])
        else:
            print "DRIVE: drive-name='%s' num-projects=%d total-songs=%d"%(
                        self.info['drive-name'], self.info['num-projects'],
                        self.info['total-songs'])
        print "PROJECTS:"
        for idx, proj in enumerate(self.info['projects']):
            self._showProject(verbose, idx, proj)
        print "Total disk space required = %s"%_human_readable(self.info['total-size'])

    def _start_timing(self):
        self._start = time.time()

    def _finish_timing(self, size):
        duration = time.time()-self._start
        print "\nExtracted %s in %1.2f secs. Average transfer rate was %s/sec."%(
            _human_readable(size), duration, _human_readable(size/duration))

    def _getSong(self, proj, idx):
        try:
            song = proj['songs'][idx]
        except IndexError:
            raise HD24Error("There is no Song-%02d. Index out of range."%(idx+1))
        except KeyError:
            raise HD24Error("This project has no songs.")
        return song

    def _getProject(self, idx):
        try:
            proj = self.info['projects'][idx]
        except IndexError:
            raise HD24Error("There is no Project-%02d. Index out of range."%(idx+1))
        return proj

    def _showInfoSong(self, verbose, proj_index, song_index, track_list):
        proj = self._getProject(proj_index)
        song = self._getSong(proj, song_index)
        print "PROJECT:\n  name='%s':"%getProjectName(proj_index, proj)
        length = self._showSong(verbose, song_index, song, track_list)
        if length == song['song-size']:
            print "  Total song size = %s"%_human_readable(length)
        else:
            ntracks = getNumTracks(song, track_list)
            print "  Total song size (%d track%s) = %s"%(ntracks, _s(ntracks), _human_readable(length))

    def _showInfoProject(self, verbose, proj_index, *_):
        proj = self._getProject(proj_index)
        print "PROJECT:"
        self._showProject(verbose, proj_index, proj)

    def _showInfoAll(self, verbose, *_):
        self._showInfo(verbose)

    def showInfo(self, options):
        i = bool(options.project != None) + bool(options.song != None)
        showInfoFN = (self._showInfoAll, self._showInfoProject, self._showInfoSong)
        showInfoFN[i](options.verbose, options.project, options.song, options.tracks)

    def _choice(self, prompt):
        res = raw_input(prompt).lower()
        if res == '': return ''
        return res[0]

    def _create_missing_dir(self, path):
        if not os.path.exists(path):
            if self.create_all:
                _safe_mkdir(path)
                return True
            while 1:
                res = self._choice("Output directory '%s' doesn't exist.\nCreate (Yes/Cancel/All)?"%path)
                if res == '': continue
                if res == 'y':
                    _safe_mkdir(path)
                    return True
                if res == 'a':
                    _safe_mkdir(path)
                    self.create_all = True
                    return True
                if res == 'c':
                    return False
        return True

    def _query_existing_dir(self, path):
        if os.path.exists(path):
            if self.overwrite_all:
                return True
            while 1:
                res = self._choice("Directory '%s' already exists.\nOverwrite (Yes/No/Cancel/All)?"%path)
                if res == '': continue
                if res == 'y':
                    return 1
                if res == 'n':
                    return 0
                if res == 'a':
                    self.overwrite_all = True
                    return 1
                if res == 'c':
                    return -1
        _safe_mkdir(path)
        return True

    def _writeSong(self, projdir, song_index, song, track_list=None):
        songname = getSongName(song_index, song)
        if not song['num-samples']:
            self.spinner.message("  Skipping   '%s': nothing recorded."%songname)
            return 1
        outdir = os.path.join(projdir, songname)
        choice = self._query_existing_dir(outdir)
        if choice != 1:
            return choice
        self.spinner.message("  Extracting '%s'..."%songname)
        self.spinner.tick(100, 0)
        self.writeTracks(outdir, song_index, song, track_list, self.spinner)
        return 1

    def _writeProject(self, outdir, proj_index, proj):
        projname = getProjectName(proj_index, proj)
        projdir  = os.path.join(outdir, projname)
        choice = self._query_existing_dir(projdir)
        if choice != 1:
            return choice
        self.spinner.message("Processing '%s'"%projname)
        for song_index, song in enumerate(proj.get('songs', [])):
            choice = self._writeSong(projdir, song_index, song)
            if choice == -1:
                return choice
        return 1

    def _extractSong(self, outdir, proj_index, song_index, track_list):
        if not self._create_missing_dir(outdir):
            return 0
        proj = self._getProject(proj_index)
        song = self._getSong(proj, song_index)
        projname = getProjectName(proj_index, proj)
        projdir  = os.path.join(outdir, projname)
        choice = self._query_existing_dir(projdir)
        if choice != 1:
            return choice
        self._start_timing()
        choice = self._writeSong(projdir, song_index, song, track_list)
        if choice != 1:
            return choice
        self._finish_timing(getNumTracks(song, track_list)*song['file-size'])
        return 1

    def _extractProject(self, outdir, proj_index, *_):
        if not self._create_missing_dir(outdir):
            return 0
        proj = self._getProject(proj_index)
        self._start_timing()
        choice = self._writeProject(outdir, proj_index, proj)
        if choice != 1:
            return choice
        self._finish_timing(proj['project-size'])
        return 1

    def _extractAll(self, outdir, *_):
        if not self._create_missing_dir(outdir):
            return 0
        self._start_timing()
        for idx, proj in enumerate(self.info['projects']):
            choice = self._writeProject(outdir, idx, proj)
            if choice == -1:
                return choice
        self._finish_timing(self.info['total-size'])
        return 1

    def extract(self, options):
        i = bool(options.project != None) + bool(options.song != None)
        extractFN = (self._extractAll, self._extractProject, self._extractSong)
        return extractFN[i](options.outdir, options.project, options.song, options.tracks)


def _main():
    def index_rangeCB(option, opt_str, value, parser, lower, upper):
        if value < lower or value > upper:
            raise OptionValueError("option %s: integer value out of range [%d-%d]."%(opt_str, lower, upper))
        setattr(parser.values, option.dest, value-1)

    def index_range_appendCB(option, opt_str, value, parser, lower, upper):
        if value < lower or value > upper:
            raise OptionValueError("option %s: integer value out of range [%d-%d]."%(opt_str, lower, upper))
        parser.values.ensure_value(option.dest, set()).add(value-1)

    parser = OptionParser(usage=USAGE_STRING, version=VERSION_STRING, description=DESCRIPTION_STRING)
    parser.add_option("-l", "--list",    action="store_true", dest="show", default=False, help="list info, projects and songs")
    parser.add_option("-x", "--extract", action="store_true", dest="extract", default=False, help="extract all songs from all projects")
    parser.add_option("-p", "--project", action="callback", callback=index_rangeCB, callback_args=(1, 99), type="int", dest="project", metavar="P", help="select project P")
    parser.add_option("-s", "--song",    action="callback", callback=index_rangeCB, callback_args=(1, 99), type="int", dest="song", metavar="S", help="select song S in project P")
    parser.add_option("-t", "--track",   action="callback", callback=index_range_appendCB, callback_args=(1, 24), type="int", dest="tracks", metavar="T", help="select track T in song S in project P")
    def_out_dir = "."
    parser.add_option("-o", "--outdir",  dest="outdir",  default=def_out_dir, metavar="OUTDIR", help="use output directory OUTDIR [default: '%s']"%(def_out_dir,))
    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False, help="be verbose")
    (options, args) = parser.parse_args()

    if not options.show and not options.extract:
        options.show = True
    if options.song != None and options.project == None:
        parser.print_help()
        print "\n%s: error: option -s|--song: to select a song, you must also select a project."%parser.get_prog_name()
        return
    if options.tracks != None and (options.song == None or options.project == None):
        parser.print_help()
        print "\n%s: error: option -t|--track: to select a track, you must also select a song and a project."%parser.get_prog_name()
        return
    if not len(args):
        parser.print_help()
        return

    try:
        hd24 = AdatFstUtils(args[0], Spinner(66))
    except HD24Error, msg:
        print "Error on Device '%s': %s"%(args[0], msg)
        return
    try:
        if options.show:
            hd24.showInfo(options)
        if options.extract:
            hd24.extract(options)
    except KeyboardInterrupt:
        print "\nOperation cancelled."
    except HD24Error, msg:
        print "Error: %s"%msg
        return


if __name__ == '__main__':
    _main()