import logging
import functools
import os
import contextlib
import multiprocessing as mp
import multiprocessing.pool

import numpy as np

from buzzard._actors.message import Msg
from buzzard._actors.pool_job import MaxPrioJobWaiting, PoolJobWorking
from buzzard._gdal_file_raster import BackGDALFileRaster
from buzzard._tools import conv
from buzzard._footprint import Footprint

LOGGER = logging.getLogger(__name__)

class ActorFileChecker(object):
    """Actor that takes care of performing various checks on a cache file from a pool"""

    def __init__(self, raster):
        self._raster = raster
        self._back_ds = raster.back_ds
        self._alive = True
        io_pool = raster.io_pool
        if io_pool is not None:
            if isinstance(io_pool, mp.pool.ThreadPool):
                self._same_address_space = True
            elif isinstance(io_pool, mp.pool.Pool):
                self._same_address_space = False
            else: # pragma: no cover
                assert False, 'Type should be checked in facade'
            self._waiting_room_address = '/Pool{}/WaitingRoom'.format(id(io_pool))
            self._working_room_address = '/Pool{}/WorkingRoom'.format(id(io_pool))
        self._waiting_jobs = set()
        self._working_jobs = set()
        self.address = '/Raster{}/FileChecker'.format(self._raster.uid)

    @property
    def alive(self):
        return self._alive

    # ******************************************************************************************* **
    def receive_infer_cache_file_status(self, cache_fp, path):
        msgs = []

        if self._raster.io_pool is None:
            work = Work(self, cache_fp, path)
            status = work.func()
            msgs += [Msg(
                'CacheSupervisor', 'inferred_cache_file_status', cache_fp, path, status
            )]
        else:
            wait = Wait(self, cache_fp, path)
            self._waiting_jobs.add(wait)
            msgs += [Msg(self._waiting_room_address, 'schedule_job', wait)]

        return msgs

    def receive_token_to_working_room(self, job, token):
        self._waiting_jobs.remove(job)
        work = Work(self, job.cache_fp, job.path)
        self._working_jobs.add(work)
        return [
            Msg(self._working_room_address, 'launch_job_with_token', work, token)
        ]

    def receive_job_done(self, job, status):
        self._working_jobs.remove(job)
        return [
            Msg('CacheSupervisor', 'inferred_cache_file_status', job.cache_fp, job.path, status)
        ]

    def receive_die(self):
        assert self._alive
        self._alive = False
        msgs = []
        for job in self._waiting_jobs:
            msgs += [Msg(self._waiting_room_address, 'unschedule_job', job)]
        for job in self._working_jobs:
            msgs += [Msg(self._working_room_address, 'cancel_job', job)]
        self._waiting_jobs.clear()
        self._working_jobs.clear()
        self._raster = None
        self._back_ds = None
        return msgs

    # ******************************************************************************************* **

class Wait(MaxPrioJobWaiting):
    def __init__(self, actor, cache_fp, path):
        self.cache_fp = cache_fp
        self.path = path
        super().__init__(actor.address)

class Work(PoolJobWorking):
    def __init__(self, actor, cache_fp, path):
        self.cache_fp = cache_fp
        self.path = path
        if actor._raster.io_pool is None or actor._same_address_space:
            func = functools.partial(
                _cache_file_check,
                cache_fp, path, len(actor._raster), actor._raster.dtype,
                actor._back_ds
            )
        else:
            func = functools.partial(
                _cache_file_check,
                cache_fp, path, len(actor._raster), actor._raster.dtype,
                None,
            )
        actor._raster.debug_mngr.event('object_allocated', func)
        super().__init__(actor.address, func)

def _checksum(fname, buffer_size=512 * 1024, dtype='uint64'):
    # https://github.com/earthcube-lab/buzzard/pull/39/#discussion_r239071556
    dtype = np.dtype(dtype)
    dtypesize = dtype.itemsize
    assert buffer_size % dtypesize == 0
    assert np.issubdtype(dtype, np.unsignedinteger)

    acc = dtype.type(0)
    with open(fname, "rb") as f:
        with np.warnings.catch_warnings():
            np.warnings.filterwarnings('ignore', r'overflow encountered')

            for chunk in iter(lambda: f.read(buffer_size), b""):
                head = np.frombuffer(chunk, dtype, count=len(chunk) // dtypesize)
                head = np.add.reduce(head, dtype=dtype, initial=acc)
                acc += head

                tailsize = len(chunk) % dtypesize
                if tailsize > 0:
                    # This should only be needed for file's tail
                    tail = chunk[-tailsize:] + b'\0' * (dtypesize - tailsize)
                    tail = np.frombuffer(tail, dtype)
                    acc += tail
        return '{:016x}'.format(acc.item())

def _cache_file_check(cache_fp, path, channel_count, dtype, back_ds_opt):
    checksum = path
    checksum = checksum.split('.')[-2]
    checksum = checksum.split('_')[-1]
    new_checksum = _checksum(path)
    if new_checksum != checksum:
        if back_ds_opt is not None:
            back_ds_opt.deactivate(path)
        LOGGER.warning('Removing {} because invalid checksum ({} instead of {})'.format(
            path, new_checksum, checksum,
        ))
        os.remove(path)
        return False

    allocator = lambda: BackGDALFileRaster.open_file(path, 'GTiff', [], 'r') # This may raise
    with contextlib.ExitStack() as stack:
        try:
            if back_ds_opt is None:
                gdal_ds = allocator()
            else:
                gdal_ds = stack.enter_context(back_ds_opt.acquire_driver_object(path, allocator))

            file_fp = Footprint(
                gt=gdal_ds.GetGeoTransform(),
                rsize=(gdal_ds.RasterXSize, gdal_ds.RasterYSize),
            )
            file_dtype = conv.dtype_of_gdt_downcast(gdal_ds.GetRasterBand(1).DataType)
            file_len = gdal_ds.RasterCount
            if file_fp != cache_fp: # pragma: no cover
                raise RuntimeError('invalid Footprint of {}({} instead of {})'.format(
                    path, file_fp, cache_fp
                ))
            if file_dtype != dtype: # pragma: no cover
                raise RuntimeError('invalid dtype of {}({} instead of {})'.format(
                    path, file_dtype, dtype
                ))
            if file_len != channel_count: # pragma: no cover
                raise RuntimeError('invalid channel_count of {}({} instead of {})'.format(
                    path, file_len, channel_count
                ))
            del gdal_ds
        except Exception:
            # Those exceptions should not trigger a cache file removal, because it might originate
            # from a mistake in the code that does not mean that those files are corrupted. For exemple:
            # - Maximum number of file descriptors reach
            # - Mismatch in cache directories path
            back_ds_opt.deactivate(path)
            raise

    return True
