import shutil
import tarfile
import zipfile
import IPython

from pathlib import Path
from urllib.parse import urlparse
from typing import List, Union, Optional, Iterable, Generator
from tqdm.auto import tqdm
from os.path import relpath


__all__ = [
    "InvalidURLException",
    "download_dataset",
    "download",
    "read_dataset",
    "read",
    "prepare_dataset",
    "prepare",
]

DEFAULT_CHUNK_SIZE = 8 << 10


class InvalidURLException(Exception):
    """
    Raised if URL is invalid.
    """

    def __init__(self, url):
        self.url = url
        self.message = f"'{self.url}' is not a valid URL."
        super().__init__(self.message)


def _is_jupyterlite() -> bool:
    """
    :returns: True if running in jupyterlite environment, False otherwise.
    """
    return IPython.sys.platform == "emscripten"


def _is_url_valid(url: str) -> bool:
    """
    :param url: URL which is checked for validity.
    :returns: True if url is valid URL, False otherwise.
    """
    try:
        result = urlparse(url)
        # Assume it's a valid URL if a URL scheme and netloc are successfully parsed
        return all([result.scheme, result.netloc])
    except Exception:
        # If urlparse chokes on something, assume it's an invalid URL
        return False


async def _get_chunks(url: str, chunk_size: int) -> Generator[bytes, None, None]:
    """
    Generator that yields consecutive chunks of bytes from URL 'url'
    :param url: The URL containing the data file to be read
    :param chunk_size: The size of each chunk (in no. of bytes).
    :returns: Generator yielding chunks of bytes from file at URL until done.
    :raise InvalidURLException: When URL is invalid.
    :raise Exception: When Exception encountered when reading from URL.
    """
    if not _is_url_valid(url):
        raise InvalidURLException(url)
    desc = f"Downloading {Path(urlparse(url).path).name}"
    if _is_jupyterlite():
        from js import fetch  # pyright: ignore
        from pyodide import JsException  # pyright: ignore

        try:
            response = await fetch(url)
            reader = response.body.getReader()
            pbar = tqdm(
                mininterval=1,
                desc=desc,
                total=int(response.headers.get("content-length", 0)),
            )
            while True:
                res = (await reader.read()).to_py()
                value, done = res["value"], res["done"]
                if done:
                    break
                value = value.tobytes()
                yield value
                pbar.update(len(value))
            pbar.close()
        except JsException:
            raise Exception(f"Failed to read dataset at '{url}'.") from None
    else:
        import requests  # pyright: ignore
        from requests.exceptions import ConnectionError  # pyright: ignore

        try:
            with requests.get(url, stream=True) as response:
                # If requests.get fails, it will return readable error
                if response.status_code >= 400:
                    raise Exception(
                        f"received status code {response.status_code} from '{url}'."
                    )
                pbar = tqdm(
                    miniters=1,
                    desc=desc,
                    total=int(response.headers.get("content-length", 0)),
                )
                for chunk in response.iter_content(chunk_size=chunk_size):
                    yield chunk
                    pbar.update(len(chunk))
                pbar.close()
        except ConnectionError:
            raise Exception(f"Failed to read dataset at '{url}'.") from None


def _rmrf(path: Path) -> None:
    if path.is_dir():
        shutil.rmtree(path)
    else:
        path.unlink()


def _verify_files_dont_exist(
    paths: Iterable[Path], remove_if_exist: bool = False
) -> None:
    """
    Verifies all paths in 'paths' don't exist.
    :param paths: A iterable of pathlib.Path s.
    :param remove_if_exist=False: Remove each file at each path in paths if they already exist.
    :returns: None
    :raises FileExistsError: On the first path found that already exists if remove_if_exist is False.
    """
    for path in paths:
        # Could be a broken symlink => path.exists() is False
        if path.exists() or path.is_symlink():
            if remove_if_exist:
                while path.is_symlink():
                    temp = path.readlink()
                    path.unlink(missing_ok=True)
                    path = temp
                if path.exists():
                    _rmrf(path)
            else:
                raise FileExistsError(f"Error: File '{path}' already exists.")


def _is_file_to_symlink(path: Path) -> bool:
    """
    :param path: path to check.
    :returns: True if file should be symlinked, False otherwise.
    """
    # Don't symlink "._" junk
    return not (path.name.startswith("._") or path.name in ["__MACOSX"])


async def download(
    url: str,
    path: Optional[str] = None,
    verbose: bool = True,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
) -> None:
    """
    Downloads file located at URL to path.

    >>> import skillsnetwork
    >>> path = "./my_file.txt"
    >>> await skillsnetwork.download("https://example.com/myfile", path)
    Saved as './my_file.txt'
    >>> with open(path, "r") as f:
    >>>     content = f.read()

    :param url: The URL where the file is located.
    :param path: The path to which the file at URL should be downloaded. Auto-generated from url by default.
    :param chunk_size: The number of bytes to read from url at a time.
    :raise FileNotFoundError: If path is invalid.
    :raise skillsnetwork.InvalidURLException: If URL is invalid.
    """
    filename = Path(urlparse(url).path).name
    if path is None:
        path = Path.cwd() / filename
    else:
        path = Path(path)
        if path.is_dir():
            path /= filename
    with open(path, "wb") as f:  # Will raise FileNotFoundError if invalid path
        async for chunk in _get_chunks(url, chunk_size):
            f.write(chunk)
    if verbose:
        print(f"Saved as '{relpath(path.resolve())}'")


async def read(url: str, chunk_size: int = DEFAULT_CHUNK_SIZE) -> bytes:
    """
    Reads file at URL into bytes

    >>> import skillsnetwork
    >>> content = await skillsnetwork.read("https://example.com/myfile") # Is bytes
    >>> content_str = content.decode()                                   # Is str

    :param url: The URL where the file is located.
    :param chunk_size: Number of bytes to read from url at a time.
    :returns: bytes containing file located at URL
    :raise FileNotFoundError: If path is invalid.
    :raise InvalidURLException: If URL is invalid.
    """
    return b"".join([chunk async for chunk in _get_chunks(url, chunk_size)])


async def prepare(
    url: str, path: Optional[str] = None, verbose: bool = True, overwrite: bool = False
) -> None:
    """
    Prepares a dataset for learners. Downloads a dataset from the given url,
    decompresses it if necessary. If not using jupyterlite, will extract to
    /tmp and and symlink it so it's available at the desired path.

    >>> import skillsnetwork
    >>> await skillsnetwork.prepare("https://cf-courses-data.s3.us.cloud-object-storage.appdomain.cloud/IBM-ML0187EN-SkillsNetwork/labs/module%203/images/images.tar.gz")
    Saved to '.'

    :param url: The URL to download the dataset from.
    :param path: The path the dataset will be available at. Current working directory by default.
    :param verbose=True: Prints saved path if True.
    :param overwrite=False: Overwrites any existing files at destination if they exist.
    :raise InvalidURLException: When URL is invalid.
    :raise FileExistsError: it raises this when a file to be symlinked already exists.
    :raise ValueError: When requested path is in /tmp, or cannot be saved to path.
    """

    filename = Path(urlparse(url).path).name
    path = Path.cwd() if path is None else Path(path)
    # Check if path contains /tmp
    if Path("/tmp") in path.parents:
        raise ValueError("path must not be in /tmp.")
    elif path.is_file():
        raise ValueError("Datasets must be prepared to directories, not files.")
    # Create the target path if it doesn't exist yet
    path.mkdir(exist_ok=True)

    # For avoiding collisions with any other files the user may have downloaded to /tmp/
    dname = f"skills-network-{hash(url)}"
    # The file to extract data to. If not jupyterlite, to be symlinked to as well
    extract_dir = path if _is_jupyterlite() else Path(f"/tmp/{dname}")
    # The file to download the (possibly) compressed data to
    tmp_download_file = Path(f"/tmp/{dname}-{filename}")
    # Download the dataset to tmp_download_file file
    # File will be overwritten if it already exists
    await download(url, tmp_download_file, verbose=False)

    # Delete extract_dir directory if it already exists
    if not _is_jupyterlite():
        if extract_dir.is_dir():
            shutil.rmtree(extract_dir)
        extract_dir.mkdir()

    try:
        if tarfile.is_tarfile(tmp_download_file):
            with tarfile.open(tmp_download_file) as tf:
                _verify_files_dont_exist(
                    [
                        path / child.name
                        for child in map(Path, tf.getnames())
                        if len(child.parents) == 1 and _is_file_to_symlink(child)
                    ],  # Only check if top-level fileobject
                    remove_if_exist=overwrite,
                )
                pbar = tqdm(iterable=tf.getmembers(), total=len(tf.getmembers()))
                pbar.set_description(f"Extracting {filename}")
                for member in pbar:
                    tf.extract(member=member, path=extract_dir)
            tmp_download_file.unlink()
        elif zipfile.is_zipfile(tmp_download_file):
            with zipfile.ZipFile(tmp_download_file) as zf:
                _verify_files_dont_exist(
                    [
                        path / child.name
                        for child in map(Path, zf.namelist())
                        if len(child.parents) == 1 and _is_file_to_symlink(child)
                    ],  # Only check if top-level fileobject
                    remove_if_exist=overwrite,
                )
                pbar = tqdm(iterable=zf.infolist(), total=len(zf.infolist()))
                pbar.set_description(f"Extracting {filename}")
                for member in pbar:
                    zf.extract(member=member, path=extract_dir)
            tmp_download_file.unlink()
        else:
            _verify_files_dont_exist([path / filename], remove_if_exist=overwrite)
            shutil.move(tmp_download_file, extract_dir / filename)
    except FileExistsError as e:
        raise FileExistsError(
            str(e)
            + "\nIf you want to overwrite any existing files, use prepare(..., overwrite=True)."
        ) from None

    # If in jupyterlite environment, the extract_dir = path, so the files are already there.
    if not _is_jupyterlite():
        # If not in jupyterlite environment, symlink top-level file objects in extract_dir
        for child in filter(_is_file_to_symlink, extract_dir.iterdir()):
            if (path / child.name).is_symlink() and overwrite:
                (path / child.name).unlink()
            (path / child.name).symlink_to(child, target_is_directory=child.is_dir())

    if verbose:
        print(f"Saved to '{relpath(path.resolve())}'")


def setup() -> None:
    if _is_jupyterlite():
        tqdm.monitor_interval = 0


setup()


# For backwards compatibility
download_dataset = download
read_dataset = read
prepare_dataset = prepare
