import gzip
import os
import zipfile
from os.path import exists, getsize
from typing import Tuple, Union, TextIO, Optional, Dict, List
from urllib.error import URLError
from logzero import logger

import dask.dataframe as dd
import filetype
import rarfile
import requests
import sqlalchemy as sa
import validators
from astropy.utils import data
from requests.adapters import HTTPAdapter, Retry

import openomics


# @astropy.config.set_temp_cache(openomics.config["cache_dir"])
def get_pkg_data_filename(baseurl: str, filepath: str):
    """Downloads a remote file given the url, then caches it to the user's home
    folder.

    Args:
        baseurl: Url to the download path, excluding the file name
        filepath: The file path to download

    Returns:
        filename (str): A file path on the local file system corresponding to
        the data requested in data_name.
    """
    # Split data url and file name if the user provided a whole path in `file_resources`
    if validators.url(filepath):
        base, filepath = os.path.split(filepath)
        base = base + "/"
    else:
        base, filepath = baseurl, filepath

    try:
        # TODO doesn't yet save files to 'cache_dir' but to astropy's default cache dir
        # logger.debug(f"Fetching file from: {base}{filepath}, saving to {openomics.config['cache_dir']}")

        with data.conf.set_temp("dataurl", base), data.conf.set_temp("remote_timeout", 30):
            return data.get_pkg_data_filename(filepath, package="openomics.database", show_progress=True)

    except (URLError, ValueError) as e:
        raise Exception(f"Unable to download file at {os.path.join(base, filepath)}. "
                        f"Please try manually downloading the files and add path to `file_resources` arg. \n{e}")


def decompress_file(filepath: str, filename: str, file_ext: filetype.Type, write_uncompressed=False) \
    -> Tuple[Union[gzip.GzipFile, TextIO], str]:
    """
    Decompress the `filepath` corresponding to its `file_ext` compression type, then return the uncompressed data (or its path) and
    the `filename` without the `file_ext` suffix.

    Args:
        filepath (str): The file path to the data file
        filename (str): The filename of the data file
        file_ext (filetype.Type): The file extension of the data file
        write_uncompressed (bool): Whether to write the uncompressed file to disk

    Returns:
        uncompressed_file (): The uncompressed file path
        updated_filename (str): The filename without the `file_ext` suffix
    """
    data = filepath

    if file_ext is None:
        return data, filename

    elif file_ext.extension == "gz":
        data = gzip.open(filepath, "rt")

    elif file_ext.extension == "zip":
        with zipfile.ZipFile(filepath, "r") as zf:
            for subfile in zf.infolist():
                # Select first file with matching file extension
                subfile_name = os.path.splitext(subfile.filename)[-1]
                if subfile_name == os.path.splitext(filename.replace(".zip", ""))[-1]:
                    data = zf.open(subfile.filename, mode="r")

    elif file_ext.extension == "rar":
        with rarfile.RarFile(filepath, "r") as rf:

            for subfile in rf.infolist():
                # If the file extension matches
                subfile_name = os.path.splitext(subfile.filename)[-1]
                if subfile_name == os.path.splitext(filename.replace(".rar", ""))[-1]:
                    data = rf.open(subfile.filename, mode="r")

    else:
        logger.warn(f"WARNING: filepath_ext.extension {file_ext.extension} not supported.")
        return data, filename

    filename = get_uncompressed_filepath(filename)
    uncompressed_path = get_uncompressed_filepath(filepath)

    if write_uncompressed and not exists(uncompressed_path):
        with open(uncompressed_path, 'w', encoding='utf8') as f_out:
            logger.info(f"Writing uncompressed {filename} file to {uncompressed_path}")
            f_out.write(data.read())

    if exists(uncompressed_path) and getsize(uncompressed_path) > 0:
        data = uncompressed_path

    return data, filename


def get_uncompressed_filepath(filepath: str) -> str:
    """Return the uncompressed filepath by removing the file extension suffix.

    Args:
        filepath (str): File path to the compressed file

    Returns:
        uncompressed_path (str): File path to the uncompressed file
    """
    uncompressed_path = ''
    if filepath.endswith(".gz"):
        uncompressed_path = filepath.removesuffix(".gz")
    elif filepath.endswith(".zip"):
        uncompressed_path = filepath.removesuffix(".zip")
    elif filepath.endswith(".rar"):
        uncompressed_path = filepath.removesuffix(".rar")
    else:
        uncompressed_path = filepath + ".uncompressed"

    if uncompressed_path and filepath != uncompressed_path:
        return uncompressed_path
    else:
        return ''


def select_files_with_ext(file_resources: Dict[str, str], ext: str, contains: Optional[str] = None) -> Dict[str, str]:
    """Return a list of file paths with the specified file extension. Only string values are considered as file paths.

    Args:
        file_resources (dict): A dictionary of file names and their corresponding file paths
        ext (str): The file extension to filter the file names by
        contains (str): If not None, only return file paths that contain the specified string

    Returns:
        file_paths (dict): A dict of file names and corresponding paths with the specified file extension
    """
    subset_file_resources = {}
    for filename, filepath in file_resources.items():
        if not isinstance(filepath, str): continue
        if filename.endswith(ext) and (contains is None or contains in filename):
            subset_file_resources[filename] = filepath

    return subset_file_resources


def read_db(path, table, index_col):
    """
    Args:
        path:
        table:
        index_col:
    """
    engine = sa.create_engine(path)
    # conn = engine.connect()
    m = sa.MetaData()
    table = sa.Table(table, m, autoload=True, autoload_with=engine)

    # conn.execute("create table testtable (uid integer Primary Key, datetime NUM)")
    # conn.execute("insert into testtable values (1, '2017-08-03 01:11:31')")
    # print(conn.execute('PRAGMA table_info(testtable)').fetchall())
    # conn.close()

    uid, dt = list(table.columns)
    q = sa.select([dt.cast(sa.types.String)]).select_from(table)

    daskDF = dd.read_sql_table(table, path, index_col=index_col, parse_dates={'datetime': '%Y-%m-%d %H:%M:%S'})
    return daskDF


def retry(num=5):
    """retry connection.

    define max tries num if the backoff_factor is 0.1, then sleep() will
    sleep for [0.1s, 0.2s, 0.4s, ...] between retries. It will also force a
    retry if the status code returned is 500, 502, 503 or 504.

    Args:
        num:
    """
    s = requests.Session()
    retries = Retry(total=num, backoff_factor=0.1,
                    status_forcelist=[500, 502, 503, 504])
    s.mount('http://', HTTPAdapter(max_retries=retries))

    return s