import os
import shutil
from typing import Any, AsyncIterator, Dict, Optional, TypeVar, Union

from .boost import BoostExecutor, EagerAsyncIterator, consume
from .globals import config
from .listing import DirEntry, scantree
from .path import (
    AzurePath,
    BasePath,
    CloudPath,
    GooglePath,
    LocalPath,
    exists,
    isfile,
    pathdispatch,
    url_format,
)
from .read import read_single, read_stream, read_stream_unordered
from .request import Request, azurify_request, googlify_request
from .write import write_single, write_stream, write_stream_unordered

# ==============================
# copyfile
# ==============================


@pathdispatch
async def copyfile(
    src: Union[BasePath, str],
    dst: Union[BasePath, str],
    executor: BoostExecutor,
    overwrite: bool = False,
    size: Optional[int] = None,
) -> None:
    """Copy a file named ``src` to a file named ``dst``.

    We specialise based on the the types of the paths. For instance, if we're copying within the
    same cloud, we perform a remote copy operation without downloading the contents locally.

    :param src: The file to copy from.
    :param dst: The file to copy to.
    :param executor: An executor.
    :param overwrite: If False, raises if the path already exists.
    :param size: If specified, will save a network call.

    """
    raise ValueError(f"Unsupported path: {src}")


@copyfile.register  # type: ignore
async def _cloudpath_copyfile(
    src: CloudPath,
    dst: Union[BasePath, str],
    executor: BoostExecutor,
    overwrite: bool = False,
    size: Optional[int] = None,
) -> None:
    if isinstance(dst, str):
        dst = BasePath.from_str(dst)

    if isinstance(dst, LocalPath):
        if size is not None and size <= config.chunk_size:
            # skip a network request for small files
            await write_single(dst, await read_single(src), overwrite=overwrite)
            return

        stream = await read_stream(src, executor, size=size)
        await write_stream(dst, stream, executor, overwrite=overwrite)
        return
    if isinstance(dst, CloudPath):
        if type(src) is type(dst):
            await cloudcopyfile(src, dst, overwrite=overwrite)
            return

        if size is not None and size <= config.chunk_size:
            # skip a network request for small files
            await write_single(dst, await read_single(src), overwrite=overwrite)
            return

        if isinstance(dst, GooglePath):
            # google doesn't support unordered writes
            stream = await read_stream(src, executor, size=size)
            await write_stream(dst, stream, executor, overwrite=overwrite)
            return

        unordered_stream = await read_stream_unordered(src, executor, size=size)
        await write_stream_unordered(dst, unordered_stream, executor, overwrite=overwrite)
        return
    raise ValueError(f"Unsupported path: {dst}")


@copyfile.register  # type: ignore
async def _localpath_copyfile(
    src: LocalPath,
    dst: Union[BasePath, str],
    executor: BoostExecutor,
    overwrite: bool = False,
    size: Optional[int] = None,
) -> None:
    if isinstance(dst, str):
        dst = BasePath.from_str(dst)

    if isinstance(dst, LocalPath):
        if not overwrite:
            if await exists(dst):
                raise FileExistsError(dst)
        os.makedirs(dst.parent, exist_ok=True)
        shutil.copyfile(src, dst)
        return
    if isinstance(dst, CloudPath):
        if size is not None and size <= config.chunk_size:
            # skip a network request for small files
            await write_single(dst, await read_single(src), overwrite=overwrite)
            return

        stream = await read_stream(src, executor, size=size)
        await write_stream(dst, stream, executor, overwrite=overwrite)
        return
    raise ValueError(f"Unsupported path: {dst}")


# ==============================
# cloudcopyfile
# ==============================

CloudPathT = TypeVar("CloudPathT", bound=CloudPath)


@pathdispatch
async def cloudcopyfile(src: CloudPathT, dst: CloudPathT, overwrite: bool = False) -> None:
    """Copy a file named ``src` to a file named ``dst`` within the same cloud.

    Performs a remote copy operation without downloading the contents locally. Blocks until the
    copy completes.

    :param src: The file to copy from.
    :param dst: The file to copy to.
    :param overwrite: If False, raises if the path already exists.

    """
    raise ValueError(f"Unsupported path: {src}")


@cloudcopyfile.register  # type: ignore
async def _azure_cloudcopyfile(src: AzurePath, dst: AzurePath, overwrite: bool = False) -> None:
    assert isinstance(dst, AzurePath)
    if not overwrite:
        if await exists(dst):
            raise FileExistsError(dst)
    request = await azurify_request(
        Request(
            method="PUT",
            url=dst.format_url("https://{account}.blob.core.windows.net/{container}/{blob}"),
            headers={
                "x-ms-copy-source": src.format_url(
                    "https://{account}.blob.core.windows.net/{container}/{blob}"
                )
            },
            success_codes=(202,),
            failure_exceptions={404: FileNotFoundError(src)},
        )
    )

    async with request.execute() as resp:
        copy_id = resp.headers["x-ms-copy-id"]
        copy_status = resp.headers["x-ms-copy-status"]

    # wait for potentially async copy operation to finish
    # https://docs.microsoft.com/en-us/rest/api/storageservices/get-blob
    # pending, success, aborted, failed
    while copy_status == "pending":
        request = await azurify_request(
            Request(
                method="GET",
                url=dst.format_url("https://{account}.blob.core.windows.net/{container}/{blob}"),
            )
        )
        async with request.execute() as resp:
            if resp.headers["x-ms-copy-id"] != copy_id:
                raise RuntimeError("Copy id mismatch")
            copy_status = resp.headers["x-ms-copy-status"]
        if copy_status != "success":
            raise RuntimeError(f"Invalid copy status: '{copy_status}'")


@cloudcopyfile.register  # type: ignore
async def _google_cloudcopyfile(src: GooglePath, dst: GooglePath, overwrite: bool = False) -> None:
    assert isinstance(dst, GooglePath)
    if not overwrite:
        if await exists(dst):
            raise FileExistsError(dst)
    params: Dict[str, Any] = {}
    while True:
        request = await googlify_request(
            Request(
                method="POST",
                url=url_format(
                    (
                        "https://storage.googleapis.com/storage/v1/b/{src_bucket}/o/"
                        "{src_blob}/rewriteTo/b/{dst_bucket}/o/{dst_blob}"
                    ),
                    src_bucket=src.bucket,
                    src_blob=src.blob,
                    dst_bucket=dst.bucket,
                    dst_blob=dst.blob,
                ),
                params=params,
                failure_exceptions={404: FileNotFoundError(src)},
            )
        )
        async with request.execute() as resp:
            result = await resp.json()
            if result["done"]:
                return
            params["rewriteToken"] = result["rewriteToken"]


# ==============================
# copytree
# ==============================


async def copytree_iterator(
    src: BasePath, dst: Union[BasePath, str], executor: BoostExecutor
) -> AsyncIterator[BasePath]:
    """Copies the tree rooted at ``src`` to ``dst``.

    Yields the paths as they are copied.

    :param src: The root of the tree to copy from.
    :param dst: The root of the tree to copy to.
    :param executor: An executor.

    """
    if isinstance(dst, str):
        dst = BasePath.from_str(dst)

    if isinstance(src, LocalPath):
        src = LocalPath(os.path.abspath(src.path))

    dirpath = src.ensure_directory_like()
    subpath_exists = False

    async def copy_wrapper(entry: DirEntry) -> BasePath:
        nonlocal subpath_exists
        subpath_exists = True

        assert isinstance(dst, BasePath)
        if entry.is_dir:
            return entry.path
        size = entry.stat.size if entry.stat else None
        await copyfile(
            entry.path, dst / (entry.path.relative_to(src)), executor, size=size, overwrite=True
        )
        return entry.path

    async for path in executor.map_unordered(copy_wrapper, EagerAsyncIterator(scantree(dirpath))):
        yield path

    # If we find nothing, then run some checks so we throw the appropriate error.
    # Doing this means we avoid extra requests in the happy path.
    if not subpath_exists:
        if not await isfile(src):
            raise FileNotFoundError(src)
        raise NotADirectoryError(src)


@pathdispatch
async def copytree(
    src: Union[BasePath, str], dst: Union[BasePath, str], executor: BoostExecutor
) -> None:
    """Copies the tree rooted at ``src`` to ``dst``.

    :param src: The root of the tree to copy from.
    :param dst: The root of the tree to copy to.
    :param executor: An executor.

    """
    raise ValueError(f"Unsupported path: {src}")


@copytree.register  # type: ignore
async def _cloud_copytree(
    src: CloudPath, dst: Union[BasePath, str], executor: BoostExecutor
) -> None:
    await consume(copytree_iterator(src, dst, executor))


@copytree.register  # type: ignore
async def _local_copytree(
    src: LocalPath, dst: Union[BasePath, str], executor: BoostExecutor
) -> None:
    if isinstance(dst, str):
        dst = BasePath.from_str(dst)

    if isinstance(dst, CloudPath):
        await consume(copytree_iterator(src, dst, executor))
        return
    assert isinstance(dst, LocalPath)
    shutil.copytree(src, dst)
