#!/usr/bin/env python
# -*- coding: utf-8 -*-

# 2011-01-27 11:32:42

###############################################################################
# Copyright (c) 2011, Vadim Shlyakhov
#
#  Permission is hereby granted, free of charge, to any person obtaining a
#  copy of this software and associated documentation files (the "Software"),
#  to deal in the Software without restriction, including without limitation
#  the rights to use, copy, modify, merge, publish, distribute, sublicense,
#  and/or sell copies of the Software, and to permit persons to whom the
#  Software is furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included
#  in all copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
#  OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
#  THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#  DEALINGS IN THE SOFTWARE.
###############################################################################

from __future__ import with_statement
from __future__ import print_function

from PIL import Image
import zlib
import mmap
import struct
import itertools

def flatten(two_level_list):
    return list(itertools.chain(*two_level_list))

class OzfImg(object):

#ozf_header_1:
#  short magic;         // set it to 0x7780 for ozfx3 and 0x7778 for ozf2
#  long locked;         // if set to 1, than ozi refuses to export the image (doesn't seem to work for ozfx3 files though); just set to 0
#  short tile_width;    // set always to 64
#  short version;       // set always to 1
#  long old_header_size;// set always to 0x436; this has something to do with files having magic 0x7779
#                       // (haven't seen any of those, they are probably rare, but ozi has code to open them)
    hdr1_fmt='<HIHHI'
    hdr1_size=struct.calcsize(hdr1_fmt)
    hdr1_fields=('magic','locked','tile_width','version','old_header_size')

#ozf_header_2:
#?  int header_size;?  // always 40
#?  int image_width;?  // width of image in pixels
#?  int image_height;?  // height of image in pixels
#?  short depth;?  ?  // set to 1
#?  short bpp;?  ?  ?  // set to 8
#?  int reserved1;?  ?  // set to 0
#?  int memory_size;?  // height * width; probably not used and contains junk if it exceeds 0xFFFFFFFF (which is perfectly ok)
#?  int reserved2;?  ?  // set to 0
#?  int reserved3;?  ?  // set to 0
#?  int unk2;?  ?  ?  // set to 0x100
#?  int unk2;?  ?  ?  // set to 0x100

    hdr2_fmt='<iIIhhiiiiii'
    hdr2_size=struct.calcsize(hdr2_fmt)
    hdr2_fields=('header_size','image_width','image_height','depth','bpp',
        'reserved1','memory_size','reserved2','reserved3','unk2','unk3')

#zoom_level_hdr:
#   uint width;
#   uint height;
#   ushort tiles_x;
#   ushort tiles_y;

    def __init__(self,img_fname,ignore_decompression_errors=False):
        self.ignore_decompression_errors=ignore_decompression_errors
        self.errors=[]
        self.fname=img_fname
        self.f=open(img_fname,'r+b')
        self.mmap=mmap.mmap(self.f.fileno(),0,access=mmap.ACCESS_READ)
        self.mmap_pos=0
        long_fmt='<I'
        long_size=4
        oziread=self.read

        # header 1
        hdr1=oziread(self.hdr1_fmt,self.hdr1_fields)
        magic=hdr1['magic']
        try:
            self.ozfx3 = { 0x7780: True, 0x7778: False}[magic]
        except KeyError:
            raise Exception('Invalid file: %s' % img_fname)

        hdr2_ofs=self.tell() # remember current pos

        if self.ozfx3: # initialize descrambling
            magic_size=oziread('<B')[0]
            magic_lst=oziread('<%dB' % magic_size)
            self.new_seed(magic_lst[0x93])
            magic2=oziread(long_fmt)[0]
#             if magic2 != 0x07A431F1: pass
#                 pf('m',end='')
            try:
                magic2_prog={
                    0x07A431F1: 'Ozf2Img 3.00+', # ; unrestricted saving; encryption depth: 16',
                    0xE592C118: 'MapMerge 1.05+',# ; no saving; encryption depth: 16',
                    0xC208F454: 'MapMerge 1.04', # ; no saving; encryption depth: full',
                    }[magic2]
            except:
                magic2_prog='?'
            # re-read hdr1
            hdr2_ofs=self.tell() # remember current pos
            self.seek(0)
            hdr1=oziread(self.hdr1_fmt,self.hdr1_fields)

            # find a seed which decodes hdr2
            pattern=struct.pack('<I',self.hdr2_size) # 1st dword is hdr2 size
            src=self.mmap[hdr2_ofs:hdr2_ofs+4]
            for i in range(256):
                seed=(i+magic_lst[0x93]+0x8A) & 0xFF # start from the well known seed
                if self.descramble(src,seed=seed) == pattern :
                    break
            else:
                raise Exception('seed not found')
            foo=seed - magic_lst[0x93]
            if foo < 0: foo+=256
#             if foo not in magic_lst:
#                 pf('s',end='')
            self.new_seed(seed)

        # continue with header 1
        tw=hdr1['tile_width']
        self.tile_sz=(tw,tw)
        assert tw == 64
        assert hdr1['version'] == 1
        assert hdr1['old_header_size'] == 0x436

        # header 2
        self.seek(hdr2_ofs)
        hdr2=oziread(self.hdr2_fmt,self.hdr2_fields)
        assert hdr2['header_size'] == self.hdr2_size
        assert hdr2['bpp'] == 8
        assert hdr2['depth'] == 1

        # pointers to zoom level tilesets
        self.seek(self.mmap.size()-long_size)
        zoom_lst_ofs=oziread(long_fmt)[0]
        zoom_cnt=(self.mmap.size()-zoom_lst_ofs)//long_size-1
        self.seek(zoom_lst_ofs)
        zoom0_ofs=oziread(long_fmt)[0]

        # zoom 0 level hdr, unscramble individually
        self.seek(zoom0_ofs)
        w,h,tiles_x,tiles_y=[oziread('<'+f)[0] for f in 'IIHH']
        self.size=(w,h)
        self.t_range=(tiles_x,tiles_y)

        # palette
        p_raw=oziread('<%dB' % (256*long_size))
        self.palette=flatten(zip(p_raw[2::4],p_raw[1::4],p_raw[0::4]))
        try:
            assert not any(p_raw[3::4]), 'pallete pad is not zero'
        except:
	         pass
#             pf('p',end='')
        #self.max_color=0
        #ld('palette',self.palette)

        # tile offsets, unscramble individually
        self.tile_ofs=[oziread(long_fmt)[0] for i in range(self.t_range[0]*self.t_range[1]+1)]
        #ld(self.tile_ofs)
#         pf('.',end='')

    def close(self):
        self.mmap.close()
        self.f.close()
        return self.fname,self.errors

    def tile_data(self,x,y,flip=True):
        idx=x+y*self.t_range[0]
        ofs=self.tile_ofs[idx]
        nofs=self.tile_ofs[idx+1]
        try:
            tile=zlib.decompress(self.descramble(self.mmap[ofs:nofs],descr_len=16))
        except zlib.error as exc:
            err_msg='tile %d,%d %s' % (x,y,exc.args[0])
            self.errors.append(err_msg)
            if self.ignore_decompression_errors:
                tx,ty=self.tile_sz
                tile='\x00'*(tx*ty)
            else:
                raise exc
        #self.max_color=max(self.max_color,max(tile))
        if flip:
            tx,ty=self.tile_sz
            tile=''.join([tile[i*tx:(i+1)*ty] for i in range(ty-1,-1,-1)])
        return tile

    def read(self,fmt,fields=None):
        sz=struct.calcsize(fmt)
        src=self.mmap[self.mmap_pos:self.mmap_pos+sz]
        res=struct.unpack(fmt,self.descramble(src))
        self.mmap_pos+=sz
        return dict(zip(fields,res)) if fields else res

    def tell(self):
        return self.mmap_pos

    def seek(self,pos):
        self.mmap_pos=pos

#     ozfx3_key=bytearray('\x2D\x4A\x43\xF1\x27\x9B\x69\x4F\x36\x52\x87\xEC\x5F'
#                         '\x42\x53\x22\x9E\x8B\x2D\x83\x3D\xD2\x84\xBA\xD8\x5B')
    ozfx3_key='\x2D\x4A\x43\xF1\x27\x9B\x69\x4F\x36\x52\x87\xEC\x5F\x42\x53\x22\x9E\x8B\x2D\x83\x3D\xD2\x84\xBA\xD8\x5B'
    ozfx3_key_ln=len(ozfx3_key)

    def descramble(self,src,descr_len=None,seed=None):
        return src

    def ozfx3_descramble(self,src,descr_len=None,seed=None):
        if seed is None:
            seed=self.seed
        if descr_len is None:
            descr_len=len(src)
#         res=bytearray(src[:descr_len])
        res=src[:descr_len]
        key=self.ozfx3_key
        kln=self.ozfx3_key_ln
        for i in range(descr_len):
            res[i] ^= (key[i % kln] + seed) & 0xFF
        return str(res)+src[descr_len:]

    def new_seed(self,seed):
        self.seed=seed
        self.descramble=self.ozfx3_descramble

    def tile(self,x,y):
        data=self.tile_data(x,y)#,flip=False)
        img=Image.fromstring('L',self.tile_sz,data,'raw','L',0,1)
        #img.putpalette(self.palette)
        return img

    def image(self):
        img=Image.new('L',self.size)
        for x in range(self.t_range[0]):
            for y in range(self.t_range[1]):
                tile=self.tile(x,y)
                img.paste(tile,(x*self.tile_sz[0],y*self.tile_sz[1]))
        img.putpalette(self.palette)
        return img
