#!/usr/bin/env python

import sys
import re
import os, os.path
import cksum
import zipfile
try:
    import rarfile
except:
    print 'warning: no rar support (get "rarfile" from cheeseshop)'
    rarfile = False
try:
    import py7zlib
except:
    print 'warning: no 7zip support (get "pylzma" from cheeseshop)'
    py7zlib = False
from optparse import OptionParser

region_map_1l = {
    'E'  : 'Europe',
    'J'  : 'Japan',
    'U'  : 'USA',
    'I'  : 'Italy',
    'G'  : 'Germany',
    'F'  : 'France',
    'S'  : 'Spain',
    'C'  : 'China',
    'N'  : 'Netherlands',
    'K'  : 'Korea',
    'A'  : 'Australia',
}

region_map_2l = {
    'EU'  : 'Europe',
    'JP'  : 'Japan',
    'US'  : 'USA',
    'IT'  : 'Italy',
    'DE'  : 'Germany',
    'FR'  : 'France',
    'ES'  : 'Spain',
    'CN'  : 'China',
    'NE'  : 'Netherlands',
    'KO'  : 'Korea',
    'AU'  : 'Australia',
}

region_map  = region_map_1l.copy()
region_map.update(region_map_2l)

def map_region(region):
    for key,val in region_map_1l.items():
        if val == region:
            return key

region_keys = '|'.join(region_map.keys())

phr_info_format = """clrmamepro (
        name "%s"
        description "%s (fixed)"
        category %s
        version (%s-fixed)
        author %s+gbaname
)

"""

phr_game_format = """game (
        name "%s - %s"
        description "%s - %s"
        rom ( name "%s - %s" crc %s )
)

"""

# 0463 - MAX Media Loader
# 0799 - NinjaPass X9TF
phr_hwo_list = [ 463, 799 ]

# source for ldist from BBand: 
# http://groups.google.com/group/comp.lang.python/browse_frm/thread/43f3ef0bf88f40d2
def ldist(a,b):
    "Calculates the Levenshtein distance between two strings"
    c = {}
    n = len(a); m = len(b)
    for i in range(0,n+1):
        c[i,0] = i
    for j in range(0,m+1):
        c[0,j] = j

    for i in range(1,n+1):
        for j in range(1,m+1):
            x = c[i-1,j]+1
            y = c[i,j-1]+1
            if a[i-1] == b[j-1]:
                z = c[i-1,j-1]
            else:
                z = c[i-1,j-1]+1
            c[i,j] = min(x,y,z)
    return c[n,m]

# this code is very naive; it will quit if zip files contain
# more than one file in them...
def zip_extract(file, outfile):
    try:
        zf = zipfile.ZipFile(file)
    except:
        print 'error in zip file "%s"' % (file)
        return False
    return extract(zf, 'zip', outfile)

def rar_extract(file, outfile):
    try:
        rf = rarfile.RarFile(file)
    except:
        print 'error in rar file "%s"' % (file)
        return False
    return extract(rf, 'rar', outfile)

def lzma_extract(file, outfile):
    try:
        f = open(file, 'rb')
        zf = py7zlib.Archive7z(f)
    except:
        print 'error in 7zip file "%s"' % (file)
        return False
    for af in zf.files:
        ext = af.filename.split('.')[-1]
        if ext.lower() in ('nds', 'gba'):
            found = af
        if not found:
            print '%s file (%s) w/o a file in these types: %s' % ('7zip', zf, ('nds', 'gba'))
            return False
        outfile.write(found.read())
        outfile.seek(0)
        return outfile

def extract(cf, ftype, outfile):
    found = None
    for name in cf.namelist():
        ext = name.split('.')[-1]
        if ext.lower() in ('nds', 'gba'):
            found = name
    if not found:
        print '%s file (%s) w/o a file in these types: %s' % (ftype, cf, ('nds', 'gba'))
        return False
    outfile.write(cf.read(found))
    outfile.seek(0)
    return outfile

class Rom:
    def __init__(self, filename):
        self.filename = filename
        self.name = filename.replace('_', ' ')
        self.number = int(re.compile('.*?(?P<num>[0-9]{4}).*').match(self.name).group('num'))
        s = self.name.split(str(self.number))
        if len(s[0]) > len(s[1]): tmp = s[0]
        else: tmp = s[1]
        self.cleanname = ''.join(tmp.split('.')[:-1])
        self.extension = self.filename[-3:]
        if tmp[-7:] in ('.tar.gz', 'tar.bz2'): self.extension = tmp[-7:].strip('.')
        # bother with this?
        self.dupe = False
        self.rename = False
    
    def get_crc(self):
        tmp = False
        if self.extension in ('gba', 'nds'):
            #print ' > performing checksum calculation on %s' % (self.filename),
            self.crc = cksum.CheckSum(get_path() + self.filename).getfilecrc()
            #print self.crc
        elif self.extension.lower() == 'zip':
            tmp = zip_extract(get_path() + self.filename, os.tmpfile())
        elif self.extension.lower() == 'rar':
            tmp = rar_extract(get_path() + self.filename, os.tmpfile())
        elif self.extension.lower() == '.7z':
            tmp = lzma_extract(get_path() + self.filename, os.tmpfile())
        else:
            print 'unrecognized file: %s (%s)' % (self.filename, self.extension)
        if tmp:
            self.crc = cksum.CheckSum(tmp).getfilecrc()
            tmp.close()
            del tmp
        else:
            self.crc = None

    def __str__(self):
        return "rom file \"%s\"\n\tname: %s\n\trelease no.: %s\n\tclean: %s\n\tformat: %s" % (self.filename, self.name, self.number, self.cleanname, self.extension)

class Game:
    def __init__(self, gamestr):
        global region_map, region_keys
        for line in gamestr:
            if line.startswith('name'):
                self.name = line.split('"')[1]
                self.filename = self.name
                self.cleanname = '-'.join(self.name.split('-')[1:]).split('(')[0].strip()
                try:
                    self.number = int(self.name.split('-')[0].strip())
                except ValueError: # most likely, number = 'xxxx'
                    self.number = self.name.split('-')[0].strip()
            elif line.startswith('description'):
                self.desc = line.split('"')[1]
            elif line.startswith('rom'):
                self.rom = self.romparse(line)
                self.crc = self.rom['crc']
                
        # match out the region and dumper info from the name attribute
        rd_re = re.compile('\((?P<region>('+region_keys+'))\).*?(\((?P<dumper>[^\)]+)\))?')
        m = rd_re.search(self.name)
        try:
            self.region = region_map[m.group('region')]
            self.dumper = m.group('dumper')
        except:
            self.region = '?'
            self.dumper = '?'
            printv('problem parsing region/dumper in %s; is region not one of [%s]?' % (self.name, region_keys))
        # set when we have checked a rom file against this record
        self.checked = False 
        
    def romparse(self, line):
        # a group for the name (between "'s) and crc (crc (8 hex digits))
        rom_re = re.compile('.*?"(?P<name>.*?)".*?crc (?P<crc>[0-9a-fA-F]{8})')
        size_re = re.compile('size (?P<size>[0-9]*)')
        m = rom_re.search(line)
        m2 = size_re.search(line)
        if not m2:
            return { 'name' : m.group('name'), 'crc' : m.group('crc') }
        else:
            return { 'name' : m.group('name'), 'crc' : m.group('crc'), 'size' : m2.group('size') }

    def format(self):
        return ("%04d" % (self.number),
                "-".join(self.name.split('-')[1:]).strip(), # rom (C)(dumper)
                "%04d" % (self.number),
                "-".join(self.desc.split('-')[1:]).strip(), # same for desc
                "%04d" % (self.number),
                "-".join(self.rom['name'].split('-')[1:]).strip(), # same for rom name
                self.rom['crc']
        )

    def __str__(self):
        return """%s\n\tclean: %s\n\trelease no.: %04d\n\tdesc: %s\n\tcrc: %s\n\tregion: %s\n\tdumper: %s""" % (self.name, self.cleanname, self.number, self.desc, self.rom['crc'], self.region, self.dumper)

    def __repr__(self):
        return "%s - %s (%s)" % (self.number, self.cleanname, map_region(self.region))

class NameInfo:
    def __init__(self, dat):
        "Takes an open dat file and creates an info list."
        self.lines = [l.strip() for l in dat.readlines()]
        self.header = []
        for line in self.lines:
            if len(line):
                self.header.append(line)
            else: break
        self.read_header()
        self.gameinfo = self.lines[len(self.header)+1:]
        self.games = []
        for chunk in self.chunklines():
            self.games.append(Game(chunk))
        self.games.sort(rom_cf)
        # for now, strip out the nukes from ADVANScENe dats
        self.games = [game for game in self.games if game.number != 'xxxx']
    
    def read_header(self):
        "Sets the header information for the PH info list file."
        for line in self.header:
            if line.startswith('name'):
                nre = re.compile(r'^name "?(?P<name>.*?)"?$')
                self.name = nre.match(line).group('name')
            elif line.startswith('description'):
                dre = re.compile(r'^description "?(?P<desc>.*?)"?$')
                self.desc = dre.match(line).group('desc')
            elif line.startswith('category'):
                self.category = ' '.join(line.split()[1:])
            elif line.startswith('version'): 
                # usually an int but just in case...
                vre = re.compile(r'\(?(?P<version>[\d\w_-]+)\)?')
                self.version = vre.search(line).group('version') 
            elif line.startswith('author'):
                self.author = ' '.join(line.split()[1:])

    def sanity_check(self):
        "Checks the game list indicies against game.number's"
        for idx,game in enumerate(self.games):
            if game.number != idx + 1:
                print "Error >>> Index (%s) contains rom (%s)." % (idx + 1, game.number)
                sys.exit(1)
    
    def chunklines(self):
        chunk = []
        for line in self.gameinfo:
            if len(line): chunk.append(line)
            else:
                yield chunk
                chunk = []
        # yield the last chunk
        if len(chunk): yield chunk
   
    def remove_game(self, game_number):
        "remove and renumber the game list"
        i = game_number - 1
        del self.games[i]
        while i < len(self.games):
            self.games[i].number-=1
            i+=1

    def format(self): return (self.name, self.desc, self.category, self.version, self.author)

    def export(self, fileobj=sys.stdout):
        global phr_info_format, phr_game_format
        fileobj.write(phr_info_format % self.format())
        for game in self.games:
            fileobj.write(phr_game_format % game.format())

# a rom comparisson function
def rom_cf(x,y):
    if x.number > y.number:
        return 1
    elif x.number == y.number:
        #printv('Note:  duplicate rom found for %s ("%s" and "%s")' % (x.number, x.filename, y.filename))
        x.dupe = True
        y.dupe = True
        return 0
    else:
        return -1

def get_path():
    d = args[1]
    if d[-1] != '/': d += '/'
    return d

def get_name(rom, games, crc=True):
    game = games[rom.number -1]
    game.checked = True
    if crc:
        return '%04d - %s (%s).%s' % (game.number, game.cleanname, map_region(game.region), rom.extension)
    else:
        return '%04d - %s (%s)(!).%s' % (game.number, game.cleanname, map_region(game.region), rom.extension)

def check_rom(rom, games):
    game = games[rom.number -1]
    rom.get_crc()
    if not rom.crc:
        printv('Could not obtain CRC for %s' % (game.name))
        return True
    if rom.crc.upper() != game.crc.upper():
        print '! %s : crc differs (got "%s", expected "%s")' % (rom.filename, rom.crc.lower(), game.crc.lower())
        possible = [game for game in games if game.crc.upper() == rom.crc.upper()]
        if possible:
            print "   > possible matches: %r" % (possible)
        return False
    printv('+ %s : crc matches (got "%s", expected "%s")' % (game.name, rom.crc.lower(), game.crc.lower()))
    return True

def rename_rom(rom, games, crc=True):
    newname = get_name(rom, games, crc)
    if crc: print "> %s  >>> %s" % (d+rom.filename, d + newname)
    else: print "# %s  ### %s" % (d+rom.filename, d + newname)
    os.rename(d + rom.filename, d + newname)

def printv(s):
    if options.verbose: print s

if __name__ == '__main__':
    # sorry global namespace!
    global options, args
    
    usage = "usage: %prog [options] datfile [romdir]"
    version = "%prog 0.1"
    
    parser = OptionParser(usage=usage, version=version)
    parser.add_option('-c', '--crc', action='store_true', dest='crc', help='perform crc checking')
    parser.add_option('-r', '--rename', action='store_true', dest='rename', help='perform rom renaming')
    parser.add_option('-p', '--pretend', action='store_true', dest='pretend', help='do not rename but display actions')
    parser.add_option('-m', '--missing', action='store_true', dest='missing', help='show missing roms')
    parser.add_option('-g', '--range', metavar="RANGE", dest='range', help='perform only over range (#-#)')
    parser.add_option('-n', '--nds-clean', action='store_true', dest='nds_clean', help='remove MAX Media Launcher, reorder, etc')
    parser.add_option('-v', '--verbose', action='store_true', dest='verbose', help='ever present verbose mode')
    
    (options, args) = parser.parse_args()
    if len(args) < 1 or (len(args) < 2 and not options.nds_clean):
        parser.print_help()
        sys.exit(0)
    dat = open(args[0])
    info = NameInfo(dat)
    info.sanity_check()

    if options.nds_clean:
        i = 0
        for game_number in phr_hwo_list:
            info.remove_game(game_number-i)
            i+= 1
        info.export()
        sys.exit(1)

    if options.range:
        global rnge
        nr = options.range.split('-')
        if len(nr) != 2:
            print 'Error: range must be formatted: ##-##'
            sys.exit(0)
        low = int(nr[0])
        high = int(nr[1]) + 1
        rnge = range(low, high)
        printv('checking roms in the range of %s' % (options.range))

    if args[0].rfind("NDS") > -1: system = "nintendo DS "
    elif args[0].rfind("GBA") > -1: system = "gameboy advanced "
    else: system = ""

    games = info.games
    #for item in info.games[:10]: print item
    printv("Indexed %d %srom releases." % (len(games), system))
    d = get_path()
    ls = os.listdir(d)
    # read ALL files.. disregard directories and save files
    roms = [Rom(l) for l in ls if os.path.isfile(d+l) and l[-4:] != '.sav']
    roms.sort(rom_cf)

    printv("Indexed %d %sroms in '%s'" % (len(roms), system, d))
    for rom in roms:
        newname = get_name(rom, games)
        if options.range:
            if rom.number not in rnge:
                continue
        if options.crc:
            if check_rom(rom, games):
                if rom.filename != newname and options.rename and not options.pretend:
                    rename_rom(rom, games)
                elif rom.filename != newname and options.rename and options.pretend:
                    print '> %s >> %s' % (rom.name, get_name(rom, games))
            else:
                if rom.filename != get_name(rom, games, crc=False) and options.rename and not options.pretend:
                    rename_rom(rom, games, crc=False)
                elif rom.filename != get_name(rom, games, crc=False) and options.rename and options.pretend:
                    print '> %s >> %s' % (rom.name, get_name(rom, games))
        elif options.rename and not options.pretend:
            if rom.filename != get_name(rom, games):
                rename_rom(rom, games)
        elif options.rename and options.pretend:
            if rom.filename != get_name(rom, games):
                print '> %s >> %s' % (rom.name, get_name(rom, games))
    if options.missing:
        missing = [game for game in games if not game.checked]
        print "missing:"
        for item in missing:
            print '   ', repr(item)