#!/usr/bin/env python2
# -*- coding: UTF-8 -*-
# File: res.py
# Date: Fri Jan 09 11:52:53 2015 +0800
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>

import glob
import os
import re
# TODO: perhaps we don't need to introduce PIL and numpy. libjpeg might be enough
import Image
import cStringIO
import base64
import logging
import imghdr
from multiprocessing import Pool
logger = logging.getLogger(__name__)

import pysox

from .avatar import AvatarReader
from .utils import timing

LIB_PATH = os.path.dirname(os.path.abspath(__file__))
INTERNAL_EMOJI_DIR = os.path.join(LIB_PATH, 'static', 'internal_emoji')
VOICE_DIRNAME = 'voice2'
IMG_DIRNAME = 'image2'
EMOJI_DIRNAME = 'emoji'
JPEG_QUALITY = 50

def do_get_voice_mp3(amr_fpath):
    """ return base64 string, and voice duration"""
    assert amr_fpath.endswith('.amr')
    assert os.path.isfile(amr_fpath), amr_fpath
    mp3_file = os.path.join('/tmp',
                            os.path.basename(amr_fpath)[:-4] + '.mp3')

    infile = pysox.CSoxStream(amr_fpath)
    outfile = pysox.CSoxStream(mp3_file, 'w', infile.get_signal())
    chain = pysox.CEffectsChain(infile, outfile)
    chain.flow_effects()
    outfile.close()

    signal = infile.get_signal().get_signalinfo()
    duration = signal['length'] * 1.0 / signal['rate']
    mp3_string = Resource.get_file_b64(mp3_file)
    os.unlink(mp3_file)
    return mp3_string, duration

class Resource(object):
    """ multimedia resources in chat"""
    def __init__(self, res_dir):
        assert os.path.isdir(res_dir), "No such directory: {}".format(res_dir)
        self.res_dir = res_dir
        self.img_dir = os.path.join(res_dir, IMG_DIRNAME)
        self.emoji_dir = os.path.join(res_dir, EMOJI_DIRNAME)
        self.avt_reader = AvatarReader(self.res_dir)
        self.init()

    def init(self):
        """ load some index in memory"""
        self.speak_data = {}
        for root, subdirs, files in os.walk(
            os.path.join(self.res_dir, VOICE_DIRNAME)):
            if subdirs:
                continue
            for f in files:
                if not f.endswith('amr'):
                    continue
                full_path = os.path.join(root, f)
                key = f[4:-4]  # msg_xxxxx.amr
                assert len(key) == 26, \
                    "Error interpreting the protocol, this is potentially a bug!"
                self.speak_data[key] = full_path

    @timing(total=True)
    def get_voice_mp3(self, imgpath):
        idx = self.voice_cache_idx.get(imgpath)
        if idx is None:
            return do_get_voice_mp3(self.speak_data[imgpath])
        return self.voice_cache[idx].get()

    def cache_voice_mp3(self, voice_paths):
        """ for speed """
        self.voice_cache_idx = {k: idx for idx, k in enumerate(voice_paths)}
        pool = Pool(3)
        self.voice_cache = [pool.apply_async(do_get_voice_mp3,
                                             (self.speak_data[k],)) for k in voice_paths]

    @staticmethod
    def get_file_b64(fname):
        data = open(fname, 'rb').read()
        return base64.b64encode(data)

    def get_avatar(self, username):
        """ return base64 string"""
        ret = self.avt_reader.get_avatar(username)
        if ret is None:
            return ""
        im = Image.fromarray(ret)
        buf = cStringIO.StringIO()
        im.save(buf, 'JPEG', quality=JPEG_QUALITY)
        jpeg_str = buf.getvalue()
        return base64.b64encode(jpeg_str)

    def _get_img_file(self, fnames):
        """ fnames: a list of filename to search for
            return (filename, filename) of (big, small) image.
            could be empty string.
        """
        cands = []
        for fname in fnames:
            dir1, dir2 = fname[:2], fname[2:4]
            dirname = os.path.join(self.img_dir, dir1, dir2)
            if not os.path.isdir(dirname):
                logger.warn("Directory not found: {}".format(dirname))
                continue
            for f in os.listdir(dirname):
                if fname in f:
                    full_name = os.path.join(dirname, f)
                    size = os.path.getsize(full_name)
                    if size > 0:
                        cands.append((full_name, size))
        if not cands:
            return ("", "")
        cands = sorted(cands, key=lambda x: x[1])

        def name_is_thumbnail(name):
            return os.path.basename(name).startswith('th_') \
                    and not name.endswith('hd')
        if len(cands) == 1:
            name = cands[0][0]
            if name_is_thumbnail(name):
                # thumbnail
                return ("", name)
            else:
                logger.warn("Found big image but not thumbnail: {}".format(fname))
                return (name, "")
        big = cands[-1]
        ths = filter(name_is_thumbnail, [k[0] for k in cands])
        if not ths:
            return (big[0], "")
        return (big[0], ths[0])


    def get_img(self, fnames):
        """ return two base64 jpg string"""
        fnames = [k for k in fnames if k]   # filter out empty string
        big_file, small_file = self._get_img_file(fnames)

        def get_jpg_b64(img_file):
            if not img_file:
                return None
            if not img_file.endswith('jpg') and \
               imghdr.what(img_file) != 'jpeg':
                im = Image.open(open(img_file, 'rb'))
                buf = cStringIO.StringIO()
                im.save(buf, 'JPEG', quality=JPEG_QUALITY)
                return base64.b64encode(buf.getvalue())
            return Resource.get_file_b64(img_file)
        big_file = get_jpg_b64(big_file)
        small_file = get_jpg_b64(small_file)

        if big_file and not small_file:
            # TODO resize big to a thumbnail
            pass
        return (big_file, small_file)

    def get_emoji(self, md5, pack_id):
        path = self.emoji_dir
        if pack_id:
            path = os.path.join(path, pack_id)
        candidates = glob.glob(os.path.join(path, '{}*'.format(md5)))
        candidates = [k for k in candidates if \
                      not k.endswith('_thumb') and not k.endswith('_cover')]
        if len(candidates) > 1:
            # annimation
            candidates = [k for k in candidates if not re.match('.*_[0-9]+$', k)]
            # only one file is the gif in need, others are frames or cover
            if len(candidates) == 0:
                # TODO stitch frames to gif
                logger.warning("Cannot find emoji: {}".format(md5))
                return None, None
        if not candidates:
            return None, None
        fname = candidates[0]
        return Resource.get_file_b64(fname), imghdr.what(fname)

    def get_internal_emoji(self, fname):
        f = os.path.join(INTERNAL_EMOJI_DIR, fname)
        return Resource.get_file_b64(f), imghdr.what(f)


