import functools as ft
import hashlib as hl
import mimetypes
import multiprocessing as mp
import os
import os.path as op
import re

from wcpan.logger import INFO, WARNING, DEBUG
import wcpan.worker as ww

from .api import Client
from .database import Database, Node
from .network import NetworkError
from .util import Settings, GoogleDriveError, stream_md5sum, FOLDER_MIME_TYPE, CHUNK_SIZE


FILE_FIELDS = 'id,name,mimeType,trashed,parents,createdTime,modifiedTime,md5Checksum,size'
CHANGE_FIELDS = 'nextPageToken,newStartPageToken,changes(fileId,removed,file({0}))'.format(FILE_FIELDS)
EMPTY_MD5SUM = 'd41d8cd98f00b204e9800998ecf8427e'


off_main_thread = ww.off_main_thread_method('_pool')


class Drive(object):

    def __init__(self, settings_path=None):
        self._settings = Settings(settings_path)
        self._client = Client(self._settings)
        self._db = Database(self._settings)
        self._pool = ww.create_thread_pool()

    def close(self):
        self._pool.shutdown()
        self._db.close()

    async def initialize(self):
        await self._client.initialize()
        self._db.initialize()

    async def sync(self):
        assert self._client.authorized

        INFO('wcpan.drive.google') << 'sync begin'

        try:
            check_point = self._db.get_metadata('check_point')
        except KeyError:
            check_point = '1'

        # first time, get root node
        if check_point == '1':
            rv = await self._client.files.get('root', fields=FILE_FIELDS)
            rv = rv.json_
            rv['name'] = None
            rv['parents'] = []
            node = Node.from_api(rv)
            self._db.insert_node(node)

        new_start_page_token = None
        changes_list_args = {
            'page_token': check_point,
            'page_size': 1000,
            'restrict_to_my_drive': True,
            'fields': CHANGE_FIELDS,
        }

        while new_start_page_token is None:
            rv = await self._client.changes.list_(**changes_list_args)
            rv = rv.json_
            next_page_token = rv.get('nextPageToken', None)
            new_start_page_token = rv.get('newStartPageToken', None)
            changes = rv['changes']

            check_point = next_page_token if next_page_token is not None else new_start_page_token

            await self._apply_changes(changes, check_point)
            changes_list_args['page_token'] = check_point

            INFO('wcpan.drive.google') << 'applied' << len(changes) << 'changes'

        INFO('wcpan.drive.google') << 'sync end'
        return True

    @property
    def root_node(self):
        return self._db.root_node

    @off_main_thread
    def get_node_by_id(self, node_id):
        return self._db.get_node_by_id(node_id)

    @off_main_thread
    def get_node_by_path(self, path):
        return self._db.get_node_by_path(path)

    @off_main_thread
    def get_path(self, node):
        return self._db.get_path_by_id(node.id_)

    @off_main_thread
    def get_path_by_id(self, node_id):
        return self._db.get_path_by_id(node_id)

    @off_main_thread
    def get_child_by_id(self, node_id, name):
        return self._db.get_child_by_id(node_id, name)

    @off_main_thread
    def get_child(self, node, name):
        return self._db.get_child_by_id(node.id_, name)

    @off_main_thread
    def get_children(self, node):
        return self._db.get_children_by_id(node.id_)

    @off_main_thread
    def get_children_by_id(self, node_id):
        return self._db.get_children_by_id(node_id)

    @off_main_thread
    def find_nodes_by_regex(self, pattern):
        return self._db.find_nodes_by_regex(pattern)

    @off_main_thread
    def find_duplicate_nodes(self):
        return self._db.find_duplicate_nodes()

    @off_main_thread
    def find_orphan_nodes(self):
        return self._db.find_orphan_nodes()

    async def download_file_by_id(self, node_id, path):
        node = await self.get_node_by_id(node_id)
        return await self.download_file(node, path)

    async def download_file(self, node, path):
        # sanity check
        if not node:
            raise ValueError('node is none')
        if node.is_folder:
            raise ValueError('node should be a file')
        if not op.isdir(path):
            raise ValueError('{0} does not exist'.format(path))

        # check if exists
        complete_path = op.join(path, node.name)
        if op.isfile(complete_path):
            return True

        # exists but not a file
        if op.exists(complete_path):
            msg = '{0} exists but is not a file'.format(complete_path)
            raise DownloadError(msg)

        # resume download
        tmp_path = complete_path + '.__tmp__'
        if op.isfile(tmp_path):
            offset = op.getsize(tmp_path)
        elif op.exists(tmp_path):
            msg = '{0} exists but is not a file'.format(complete_path)
            raise DownloadError(msg)
        else:
            offset = 0
        range_ = (offset, node.size)

        with open(tmp_path, 'ab') as fout:
            writer = ft.partial(file_consumer, fout)

            api = self._client.files
            rv = await api.download(file_id=node.id_, range_=range_,
                                    consumer=writer)

        # rename it back if completed
        os.rename(tmp_path, complete_path)

        return True

    async def create_folder(self, parent_node, folder_name):
        # sanity check
        if not parent_node:
            raise UploadError('invalid parent node')
        if not parent_node.is_folder:
            raise UploadError('invalid parent node')

        # do not create again if there is a same file
        node = await self.fetch_child_by_id(parent_node.id_, folder_name)
        if node:
            INFO('wcpan.drive.google') << 'skipped (existing)' << folder_name
            return node

        api = self._client.files
        rv = await api.create_folder(folder_name=folder_name,
                                     parent_id=parent_node.id_)
        rv = rv.json_
        node = await self.fetch_node_by_id(rv['id'])

        return node

    async def upload_file(self, file_path, parent_node, exist_ok=False):
        # sanity check
        if not parent_node:
            raise UploadError('invalid parent node')
        if not parent_node.is_folder:
            raise UploadError('invalid parent node')
        if not op.isfile(file_path):
            raise UploadError('invalid file path')

        api = self._client.files
        file_name = op.basename(file_path)

        # do not upload if remote exists a same file
        node = await self.fetch_child_by_id(parent_node.id_, file_name)
        if node:
            if exist_ok:
                INFO('wcpan.drive.google') << 'skipped (existing)' << file_path
                return node
            else:
                raise UploadConflictedError(file_path)

        total_file_size = op.getsize(file_path)
        mt, e = mimetypes.guess_type(file_path)
        if total_file_size <= 0:
            rv = await api.create_empty_file(file_name=file_name,
                                             parent_id=parent_node.id_,
                                             mime_type=mt)
            local_md5 = EMPTY_MD5SUM
        else:
            args = {
                'file_path': file_path,
                'file_name': file_name,
                'total_file_size': total_file_size,
                'parent_id': parent_node.id_,
                'mime_type': mt,
            }
            rv, local_md5 = await self._inner_upload_file(**args)

        rv = rv.json_
        node = await self.fetch_node_by_id(rv['id'])

        if node.md5 != local_md5:
            raise UploadError('md5 mismatch')

        return node

    async def fetch_child_by_id(self, node_id, name):
        safe_name = re.sub(r"[\\']", r"\\\g<0>", name)
        query = "'{0}' in parents and name = '{1}'".format(node_id,
                                                           safe_name)
        fields = 'files({0})'.format(FILE_FIELDS)
        while True:
            try:
                rv = await self._client.files.list_(q=query, fields=fields)
                break
            except NetworkError as e:
                if e.status == '400':
                    DEBUG('wcpan.drive.google') << 'failed query:' << query
                if e.fatal:
                    raise

        rv = rv.json_
        files = rv['files']
        if not files:
            return None

        node = Node.from_api(files[0])
        self._db.insert_node(node)
        return node

    async def fetch_node_by_id(self, node_id):
        rv = await self._client.files.get(node_id, fields=FILE_FIELDS)
        rv = rv.json_
        node = Node.from_api(rv)
        self._db.insert_node(node)
        return node

    async def trash_node_by_id(self, node_id):
        if node_id == self.root_node.id_:
            return
        await self._client.files.update(node_id, trashed=True)
        node = await self.get_node_by_id(node_id)
        node.is_trashed = True
        self._db.insert_node(node)
        return node

    async def trash_node(self, node):
        return await self.trash_node_by_id(node.id_)

    async def _inner_upload_file(self, file_path, file_name, total_file_size,
                                 parent_id, mime_type):
        api = self._client.files

        rv = await api.initiate_uploading(file_name=file_name,
                                          total_file_size=total_file_size,
                                          parent_id=parent_id,
                                          mime_type=mime_type)

        url = rv.get_header('Location')

        with open(file_path, 'rb') as fin:
            hasher = hl.md5()
            reader = ft.partial(file_producer, fin, hasher)
            uploader = ft.partial(self._inner_try_upload_file,
                                  url=url, producer=reader,
                                  total_file_size=total_file_size,
                                  mime_type=mime_type)

            retried = False
            offset = 0
            while True:
                ok, rv = await uploader(offset=offset)
                if ok:
                    break
                offset = rv
                fin.seek(offset, os.SEEK_SET)
                retried = True

            if retried:
                fin.seek(0, os.SEEK_SET)
                local_md5 = stream_md5sum(fin)
            else:
                local_md5 = hasher.hexdigest()

        return rv, local_md5

    async def _inner_try_upload_file(self, url, producer, offset,
                                     total_file_size, mime_type):
        api = self._client.files

        try:
            rv = await api.upload(url, producer=producer, offset=offset,
                                  total_file_size=total_file_size,
                                  mime_type=mime_type)
            return True, rv
        except NetworkError as e:
            if e.status == '404':
                raise UploadError('the upload session has been expired')
            if e.fatal:
                raise

        while True:
            try:
                rv = await api.get_upload_status(url, total_file_size)
                break
            except NetworkError as e:
                if e.fatal:
                    raise

        if rv.status != '308':
            raise UploadError('invalid upload status')
        rv = rv.get_header('Range')
        if not rv:
            raise UploadError('invalid upload status')
        rv = re.match(r'bytes=(\d+)-(\d+)', rv)
        if not rv:
            raise UploadError('invalid upload status')
        rv = int(rv.group(2))

        return False, rv

    @off_main_thread
    def _apply_changes(self, changes, check_point):
        self._db.apply_changes(changes, check_point)


class DownloadError(GoogleDriveError):

    def __init__(self, message):
        self._message = message

    def __str__(self):
        return self._message


class UploadError(GoogleDriveError):

    def __init__(self, message):
        self._message = message

    def __str__(self):
        return self._message


class UploadConflictedError(UploadError):

    def __init__(self, file_path):
        self._file_path = file_path

    def __str__(self):
        return 'remote file already exists: ' + self._file_path


async def file_producer(fin, hasher, write):
    while True:
        chunk = fin.read(CHUNK_SIZE)
        if not chunk:
            break
        hasher.update(chunk)
        await write(chunk)


def file_consumer(fout, chunk):
    fout.write(chunk)
