# coding: utf-8

import os
import time
import logging

from logging.handlers import QueueHandler
from multiprocessing import Process
from multiprocessing.shared_memory import SharedMemory
from queue import Empty

import requests
from requests.adapters import HTTPAdapter, DEFAULT_POOLBLOCK

from legendary.models.chunk import Chunk
from legendary.models.downloading import (
    DownloaderTask, DownloaderTaskResult,
    WriterTask, WriterTaskResult,
    TerminateWorkerTask, TaskFlags
)


class BindingHTTPAdapter(HTTPAdapter):
    def __init__(self, addr):
        self.__attrs__.append('addr')
        self.addr = addr
        super().__init__()

    def init_poolmanager(
            self, connections, maxsize, block=DEFAULT_POOLBLOCK, **pool_kwargs
    ):
        pool_kwargs['source_address'] = (self.addr, 0)
        super().init_poolmanager(connections, maxsize, block, **pool_kwargs)


class DLWorker(Process):
    def __init__(self, name, queue, out_queue, shm, max_retries=7,
                 logging_queue=None, dl_timeout=10, bind_addr=None):
        super().__init__(name=name)
        self.q = queue
        self.o_q = out_queue
        self.session = requests.session()
        self.session.headers.update({
            'User-Agent': 'EpicGamesLauncher/11.0.1-14907503+++Portal+Release-Live Windows/10.0.19041.1.256.64bit'
        })
        self.max_retries = max_retries
        self.shm = SharedMemory(name=shm)
        self.log_level = logging.getLogger().level
        self.logging_queue = logging_queue
        self.dl_timeout = float(dl_timeout) if dl_timeout else 10.0

        # optionally bind an address
        if bind_addr:
            adapter = BindingHTTPAdapter(bind_addr)
            self.session.mount('https://', adapter)
            self.session.mount('http://', adapter)

    def run(self):
        # we have to fix up the logger before we can start
        _root = logging.getLogger()
        _root.handlers = []
        _root.addHandler(QueueHandler(self.logging_queue))

        logger = logging.getLogger(self.name)
        logger.setLevel(self.log_level)
        logger.debug(f'Download worker reporting for duty!')

        empty = False
        while True:
            try:
                job: DownloaderTask = self.q.get(timeout=10.0)
                empty = False
            except Empty:
                if not empty:
                    logger.debug('Queue Empty, waiting for more...')
                empty = True
                continue

            if isinstance(job, TerminateWorkerTask):  # let worker die
                logger.debug('Worker received termination signal, shutting down...')
                break

            tries = 0
            compressed = 0
            chunk = None

            try:
                while tries < self.max_retries:
                    # retry once immediately, otherwise do exponential backoff
                    if tries > 1:
                        sleep_time = 2**(tries-1)
                        logger.info(f'Sleeping {sleep_time} seconds before retrying.')
                        time.sleep(sleep_time)

                    # print('Downloading', job.url)
                    logger.debug(f'Downloading {job.url}')

                    try:
                        r = self.session.get(job.url, timeout=self.dl_timeout)
                        r.raise_for_status()
                    except Exception as e:
                        logger.warning(f'Chunk download for {job.chunk_guid} failed: ({e!r}), retrying...')
                        continue

                    if r.status_code != 200:
                        logger.warning(f'Chunk download for {job.chunk_guid} failed: status {r.status_code}, retrying...')
                        continue
                    else:
                        compressed = len(r.content)
                        chunk = Chunk.read_buffer(r.content)
                        break
                else:
                    raise TimeoutError('Max retries reached')
            except Exception as e:
                logger.error(f'Job for {job.chunk_guid} failed with: {e!r}, fetching next one...')
                # add failed job to result queue to be requeued
                self.o_q.put(DownloaderTaskResult(success=False, **job.__dict__))
            except KeyboardInterrupt:
                logger.warning('Immediate exit requested, quitting...')
                break

            if not chunk:
                logger.warning('Chunk somehow None?')
                self.o_q.put(DownloaderTaskResult(success=False, **job.__dict__))
                continue

            # decompress stuff
            try:
                data = chunk.data
                size = len(data)
                if size > job.shm.size:
                    logger.fatal('Downloaded chunk is longer than SharedMemorySegment!')

                self.shm.buf[job.shm.offset:job.shm.offset + size] = data
                del chunk
                self.o_q.put(DownloaderTaskResult(success=True, size_decompressed=size,
                                                  size_downloaded=compressed, **job.__dict__))
            except Exception as e:
                logger.warning(f'Job for {job.chunk_guid} failed with: {e!r}, fetching next one...')
                self.o_q.put(DownloaderTaskResult(success=False, **job.__dict__))
                continue
            except KeyboardInterrupt:
                logger.warning('Immediate exit requested, quitting...')
                break

        self.shm.close()


class FileWorker(Process):
    def __init__(self, queue, out_queue, base_path, shm, cache_path=None, logging_queue=None):
        super().__init__(name='FileWorker')
        self.q = queue
        self.o_q = out_queue
        self.base_path = base_path
        self.cache_path = cache_path or os.path.join(base_path, '.cache')
        self.shm = SharedMemory(name=shm)
        self.log_level = logging.getLogger().level
        self.logging_queue = logging_queue

    def run(self):
        # we have to fix up the logger before we can start
        _root = logging.getLogger()
        _root.handlers = []
        _root.addHandler(QueueHandler(self.logging_queue))

        logger = logging.getLogger(self.name)
        logger.setLevel(self.log_level)
        logger.debug('Download worker reporting for duty!')

        last_filename = ''
        current_file = None

        while True:
            try:
                try:
                    j: WriterTask = self.q.get(timeout=10.0)
                except Empty:
                    logger.warning('Writer queue empty!')
                    continue

                if isinstance(j, TerminateWorkerTask):
                    if current_file:
                        current_file.close()
                    logger.debug('Worker received termination signal, shutting down...')
                    # send termination task to results halnder as well
                    self.o_q.put(TerminateWorkerTask())
                    break

                # make directories if required
                path = os.path.split(j.filename)[0]
                if not os.path.exists(os.path.join(self.base_path, path)):
                    os.makedirs(os.path.join(self.base_path, path))

                full_path = os.path.join(self.base_path, j.filename)

                if j.flags & TaskFlags.CREATE_EMPTY_FILE:  # just create an empty file
                    open(full_path, 'a').close()
                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue
                elif j.flags & TaskFlags.OPEN_FILE:
                    if current_file:
                        logger.warning(f'Opening new file {j.filename} without closing previous! {last_filename}')
                        current_file.close()

                    current_file = open(full_path, 'wb')
                    last_filename = j.filename

                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue
                elif j.flags & TaskFlags.CLOSE_FILE:
                    if current_file:
                        current_file.close()
                        current_file = None
                    else:
                        logger.warning(f'Asking to close file that is not open: {j.filename}')

                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue
                elif j.flags & TaskFlags.RENAME_FILE:
                    if current_file:
                        logger.warning('Trying to rename file without closing first!')
                        current_file.close()
                        current_file = None
                    if j.flags & TaskFlags.DELETE_FILE:
                        try:
                            os.remove(full_path)
                        except OSError as e:
                            logger.error(f'Removing file failed: {e!r}')
                            self.o_q.put(WriterTaskResult(success=False, **j.__dict__))
                            continue

                    try:
                        os.rename(os.path.join(self.base_path, j.old_file), full_path)
                    except OSError as e:
                        logger.error(f'Renaming file failed: {e!r}')
                        self.o_q.put(WriterTaskResult(success=False, **j.__dict__))
                        continue

                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue
                elif j.flags & TaskFlags.DELETE_FILE:
                    if current_file:
                        logger.warning('Trying to delete file without closing first!')
                        current_file.close()
                        current_file = None

                    try:
                        os.remove(full_path)
                    except OSError as e:
                        if not j.flags & TaskFlags.SILENT:
                            logger.error(f'Removing file failed: {e!r}')

                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue
                elif j.flags & TaskFlags.MAKE_EXECUTABLE:
                    if current_file:
                        logger.warning('Trying to chmod file without closing first!')
                        current_file.close()
                        current_file = None

                    try:
                        st = os.stat(full_path)
                        os.chmod(full_path, st.st_mode | 0o111)
                    except OSError as e:
                        if not j.flags & TaskFlags.SILENT:
                            logger.error(f'chmod\'ing file failed: {e!r}')

                    self.o_q.put(WriterTaskResult(success=True, **j.__dict__))
                    continue

                try:
                    if j.shared_memory:
                        shm_offset = j.shared_memory.offset + j.chunk_offset
                        shm_end = shm_offset + j.chunk_size
                        current_file.write(self.shm.buf[shm_offset:shm_end])
                    elif j.cache_file:
                        with open(os.path.join(self.cache_path, j.cache_file), 'rb') as f:
                            if j.chunk_offset:
                                f.seek(j.chunk_offset)
                            current_file.write(f.read(j.chunk_size))
                    elif j.old_file:
                        with open(os.path.join(self.base_path, j.old_file), 'rb') as f:
                            if j.chunk_offset:
                                f.seek(j.chunk_offset)
                            current_file.write(f.read(j.chunk_size))
                except Exception as e:
                    logger.warning(f'Something in writing a file failed: {e!r}')
                    self.o_q.put(WriterTaskResult(success=False, size=j.chunk_size, **j.__dict__))
                else:
                    self.o_q.put(WriterTaskResult(success=True, size=j.chunk_size, **j.__dict__))
            except Exception as e:
                logger.warning(f'Job {j.filename} failed with: {e!r}, fetching next one...')
                self.o_q.put(WriterTaskResult(success=False, **j.__dict__))

                try:
                    if current_file:
                        current_file.close()
                        current_file = None
                except Exception as e:
                    logger.error(f'Closing file after error failed: {e!r}')
            except KeyboardInterrupt:
                logger.warning('Immediate exit requested, quitting...')
                if current_file:
                    current_file.close()
                return
