"""
parser.local.utils module (imdb package).

This module provides miscellaneous utilities used by
the imdb.parser.local classes.

Copyright 2004-2006 Davide Alberani <da@erlug.linux.it>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program 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 General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
"""

from stat import ST_SIZE
from os import stat
from sys import maxint
from struct import unpack, calcsize

from imdb._exceptions import IMDbDataAccessError


_int_type = type(0)

# List of typical binary types used in the imdb local database indexes.
POSITION = 1
BYTE = 2
TITLE_ID = 3
ATTR = 4
INT = 5
OFFSET = 6
FULL_OFFSET = 7
NAME_ID = 8
COUNT = 9

# The size of these binary types, in bytes.
_data_size = {POSITION: 1, BYTE: 1, INT: 2, TITLE_ID: 3, ATTR: 3,
            OFFSET: 3, NAME_ID: 3, FULL_OFFSET: 4, COUNT: 3}

# Groups of one or more contiguous binary types, with specific meanings.
_dataset_list = (('plot', (TITLE_ID, FULL_OFFSET)),
                ('rating', (TITLE_ID, ('s', 10), INT, BYTE)),
                ('idx2fidx', (TITLE_ID, FULL_OFFSET)),
                ('idx2idx', (TITLE_ID, OFFSET)),
                ('akatdb', (TITLE_ID, TITLE_ID, ATTR)),
                ('akandb', (NAME_ID, NAME_ID)),
                ('akatidx', (TITLE_ID, TITLE_ID)),
                ('akanidx', (NAME_ID, NAME_ID)),
                ('personID', (NAME_ID,)),
                ('movieID', (TITLE_ID,)),
                ('attrID', (ATTR,)),
                ('position', (POSITION,)),
                ('fulloffset', (FULL_OFFSET,)),
                ('filmcount', (COUNT,)),
                ('length', (BYTE,)),
                ('orderset', (BYTE, BYTE, BYTE)),
                ('mlinks', (TITLE_ID, BYTE, TITLE_ID)),
                ('moviedata', (TITLE_ID, INT, INT, ATTR)))

def _buildStruct(items):
    """Build a tuple of ('struct_string', struct_size, struct_items)."""
    res_l = []
    for item in items:
        # If the entry is expressed with an integer, it will be read
        # from a char.
        if type(item) is _int_type:
            length = _data_size.get(item, 0)
            item = 's'
        else:
            item, length = item
        res_l.append('%s%s' % (length, item))
    res_struct = ' '.join(res_l)
    return res_struct, calcsize(res_struct), items

# Indexes in the dataset_dict.
STRUCT_STR = 0
STRUCT_SIZE = 1
STRUCT_ITEMS = 2
# The dictionary used to recover information about data sets.
dataset_dict = dict([(kind, _buildStruct(dset))
                    for kind, dset in _dataset_list])


def _buildChrIntsList(length):
    """Build a list of dictionaries with 256 keys, one for every
    char in range 0-255 with the value set to its integer (or long)
    value to the power of i, where i is the index in the
    resulting list."""
    chi = []
    for j in xrange(length):
        chi.append({})
        for count in xrange(256):
            val = count << j*8L
            if val <= maxint: val = int(val)
            chi[j][chr(count)] = val
    return chi

_chr_ints = _buildChrIntsList(4)

# XXX: using the values in pre-built _chr_ints list is
#      equivalent - but faster - than using this function:
#def _StoL(s):
#    """Convert the string s to an integer value."""
#    data = 0L
#    for i in xrange(len(s)):
#        data |= ord(s[i]) << i*8L
#    return data


# FIXME: it's slow!  use _only_ the unpack function.
def convBin(s, kind):
    """Convert the string s in a tuple where int/long values
    are converted, while other types are left untouched.
    The number of items in the tuple depends on the structure
    specified with the kind argument.

    If the result contains a single elements, this is returned
    instead of the tuple."""
    kentry = dataset_dict[kind]
    out = unpack(kentry[STRUCT_STR], s)
    res = []
    app = res.append
    for i in xrange(len(out)):
        datum = out[i]
        if type(kentry[STRUCT_ITEMS][i]) is _int_type:
            value = 0
            for k in xrange(len(datum)):
                value |= _chr_ints[k][datum[k]]
            app(value)
        else:
            app(datum)
    if len(res) == 1: return res[0]
    return res


# FIXME: it's slow!  A better (a _real_!) algorithm is required!
def getFullIndex(fname, key, kind='idx2fidx', index=0, rindex=1,
                default=None, multi=0, mode='rb'):
    """Search fname (which contains data in the kind format),
    for key at the given index, returning the rindex or default
    if key is not found.  The file is opened in mode mode; if
    multi is set, it returns a list of entries where key matches;
    if rindex is None, the whole structure is returned.

    It's basically used to convert from an offset (TITLE_ID or
    NAME_ID) to a FULL_OFFSET."""
    try:
        fptr = open(fname, mode)
    except IOError, e:
        raise IMDbDataAccessError, str(e)
    fread = fptr.read
    fseek = fptr.seek
    rsize = dataset_dict[kind][STRUCT_SIZE]
    curpos = 1L * key * rsize
    fsize = stat(fname)[ST_SIZE]
    multires = []
    if curpos >= (fsize - rsize):
        curpos = fsize - (rsize * 100L)
        if curpos < rsize: curpos = 0
    # Try to seek at a point where checkdata is lower
    # than the searched key.
    while 1:
        if curpos < rsize:
            curpos = 0
            break
        fseek(curpos)
        rdata = convBin(fread(rsize), kind)
        checkdata = rdata[index]
        if checkdata > key: curpos -= rsize * 1000L
        elif checkdata < key: break
        elif checkdata == key:
            if not multi:
                fptr.close()
                if rindex is None: return rdata
                else: return rdata[rindex]
            else:
                curpos -= rsize * 50L
    fseek(curpos)
    # Try to seek nearer the key we're searching.
    while 1:
        newpos = curpos + (rsize * 500L)
        if newpos >= fsize - rsize: break
        fseek(newpos)
        bdata = fread(rsize)
        if len(bdata) != rsize: break
        checkdata = convBin(bdata, kind)[index]
        if checkdata >= key: break
        curpos = newpos
    fseek(curpos)
    # Read every entry until the key value is found.
    while 1:
        bdata = fread(rsize)
        if len(bdata) != rsize: break
        rdata = convBin(bdata, kind)
        checkdata = rdata[index]
        if checkdata == key:
            if not multi:
                fptr.close()
                if rindex is None: return rdata
                else: return rdata[rindex]
            else:
                if rindex is None: multires.append(rdata)
                else: multires.append(rdata[rindex])
        elif checkdata > key: break
    fptr.close()
    return multires or default


def getRawData(dataF, offset, doCast=0, doWriters=0):
    """Collect data from the '.data' file dataF, reading the
    entry at offset; if doCast or doWriters are set, the related
    information are read.
    Return the personID and a list of dictionaries (of movieID and
    various other attributes) he worked in."""
    try:
        fptr = open(dataF, 'rb')
    except IOError, e:
        raise IMDbDataAccessError, str(e)
    fread = fptr.read
    fptr.seek(offset)
    piddata = fread(3)
    if len(piddata) != 3:
        fptr.close()
        return (None, ())
    personID = convBin(piddata, 'personID')
    filmcount = convBin(fread(3), 'filmcount')
    noWith = (filmcount >> 12) & 4095
    noWithout = filmcount & 4095
    resList = []
    resapp = resList.append
    for i in xrange(noWith):
        tmpd = {}
        movieID = convBin(fread(3), 'movieID')
        tmpd['movieID'] = movieID
        attrID = convBin(fread(3), 'attrID')
        tmpd['attributeID'] = attrID
        if doCast:
            length = convBin(fread(1), 'length')
            if length > 0:
                tmpd['currentRole'] = latin2utf(fread(length))
            tmpd['position'] = convBin(fread(1), 'position')
        if doWriters:
            orderset = convBin(fread(3), 'orderset')
            try:
                tmpd['position'] = orderset[2]*1000 + orderset[1]*100 + \
                                    orderset[0]
            except TypeError:
                pass
        resapp(tmpd)
    for i in xrange(noWithout):
        movieID = convBin(fread(3), 'movieID')
        tmpd = {'movieID': movieID}
        if doCast:
            length = convBin(fread(1), 'length')
            if length > 0:
                tmpd['currentRole'] = latin2utf(fread(length))
            tmpd['position'] = convBin(fread(1), 'position')
        if doWriters:
            orderset = convBin(fread(3), 'orderset')
            try:
                tmpd['position'] = orderset[2]*1000 + orderset[1]*100 + \
                                    orderset[0]
            except TypeError:
                pass
        resapp(tmpd)
    fptr.close()
    return personID, resList


def getLabel(ind, indexF, keyF):
    """Return the person name or movie title for the personID or
    movieID ind, searching the indexF and keyF files."""
    try:
        ifile = open(indexF, 'rb')
    except IOError, e:
        raise IMDbDataAccessError, str(e)
    ifile.seek(4L*ind)
    fodata = ifile.read(4)
    if len(fodata) != 4:
        ifile.close()
        return None
    idx = convBin(fodata, 'fulloffset')
    try:
        kfile = open(keyF, 'rt')
    except IOError, e:
        raise IMDbDataAccessError, str(e)
    kfile.seek(idx)
    line = kfile.readline().split('|')
    ifile.close()
    kfile.close()
    if line: return latin2utf(line[0])
    return None


class KeyFScan:
    """Create an index for the entries in the given file."""
    def __init__(self, keyF, granularity=1000):
        """Scan lines in keyF and create an index with granularity-1 entries.
        """
        posdb = []
        posdbapp = posdb.append
        self.keyF = keyF
        self.granularity = granularity
        try:
            kfile = open(keyF, 'rt')
        except IOError, e:
            raise IMDbDataAccessError, str(e)
        kfilesize = stat(keyF)[ST_SIZE]
        step = divmod(kfilesize, granularity)[0]
        for i in xrange(1, granularity):
            kfile.seek(step*i)
            kfile.readline()
            pos = kfile.tell()
            ton = latin2utf(kfile.readline().split('|')[0])
            posdbapp((ton.lower(), pos))
        kfile.close()
        self.posdb = posdb

    def getID(self, ton):
        """Return the ID for the given name or title or None if not found."""
        ton = ton.lower()
        cpos = 0
        for dbTon, dbPos in self.posdb:
            if ton <= dbTon: break
            cpos = dbPos
        try:
            kfile = open(self.keyF, 'rt')
        except IOError, e:
            raise IMDbDataAccessError, str(e)
        kfile.seek(cpos)
        line = kfile.readline()
        while line:
            curTon, curID = line.split('|')
            curTon = latin2utf(curTon.lower())
            if ton == curTon:
                kfile.close()
                return long(curID, 16)
            elif ton < curTon:
                kfile.close()
                return None
            line = kfile.readline()
        kfile.close()
        return None


def latin2utf(s):
    """Convert a latin_1 string to unicode."""
    return unicode(s, 'latin_1', 'replace')