# Copyright 2024 University of Calgary
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import requests
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from .classes import FileListingResponse, FileDownloadResult, Dataset
from ..exceptions import SRSAPIError, SRSDownloadError
from .._util import show_warning


def __download_url(
    url,
    prefix,
    output_base_path,
    headers,
    timeout,
    overwrite=False,
    pbar=None,
    pbar_iterator_nfiles=False,
):
    # set output filename
    output_filename = Path(output_base_path) / Path(url.removeprefix(prefix + "/"))
    if (overwrite is False and os.path.exists(output_filename)):
        if (pbar is not None):
            if (pbar_iterator_nfiles is True):
                pbar.update()
            else:
                pbar.update(os.path.getsize(output_filename))
        return {"filename": output_filename, "bytes_downloaded": 0}

    # create destination directory
    try:
        os.makedirs(os.path.dirname(output_filename), exist_ok=True)
    except Exception:  # pragma: nocover
        # NOTE: sometimes when making directories in parallel there are race conditions. We put
        # in a catch here and carry on if there are ever issues.
        pass

    # retrieve file and save to disk
    r = requests.get(url, headers=headers, timeout=timeout)
    if (r.status_code == 200):
        this_bytes = len(r.content)
        with open(output_filename, 'wb') as fp:
            fp.write(r.content)

        # advance progress
        if (pbar is not None):
            if (pbar_iterator_nfiles is True):
                pbar.update()
            else:
                pbar.update(this_bytes)
    else:
        if (pbar is not None):
            if (pbar_iterator_nfiles is True):
                pbar.update()
            else:
                pbar.update(0)  # pragma: no cover
        raise SRSDownloadError("HTTP error %d when downloading '%s'" % (r.status_code, url))

    # return filename
    return {"filename": output_filename, "bytes_downloaded": this_bytes}


def __download_urls(srs_obj,
                    file_listing_obj,
                    n_parallel,
                    overwrite,
                    progress_bar_disable,
                    progress_bar_ncols,
                    progress_bar_ascii,
                    progress_bar_desc,
                    timeout,
                    progress_bar_format_numurls_nobytes=False):

    def __do_parallel_work(pbar=None, pbar_iterator_nfiles=False):

        def download_task(i):
            return __download_url(
                file_listing_obj.urls[i],
                path_prefix,
                output_path,
                srs_obj.api_headers,
                timeout,
                overwrite=overwrite,
                pbar=pbar,
                pbar_iterator_nfiles=pbar_iterator_nfiles,
            )

        with ThreadPoolExecutor(max_workers=n_parallel) as executor:
            futures = [executor.submit(download_task, i) for i in range(len(file_listing_obj.urls))]
            job_data = [future.result() for future in futures]

        return job_data

    # set output path
    output_path = srs_obj.download_output_root_path / Path(file_listing_obj.dataset.name)

    # check if there's files to download
    if (file_listing_obj.count == 0):
        show_warning("No data found to download", stacklevel=5)

        # cast into return obj
        download_obj = FileDownloadResult(
            filenames=[],
            count=0,
            dataset=file_listing_obj.dataset,
            total_bytes=0,
            output_root_path=output_path,
        )
    else:
        # set timeout
        if (timeout is None):
            timeout = srs_obj.api_timeout

        # set path prefix
        path_prefix = file_listing_obj.path_prefix

        # set progress bar description text
        desc_str = "Downloading %s files" % (file_listing_obj.dataset.name)
        if (progress_bar_desc is not None):
            desc_str = progress_bar_desc

        # download the urls
        parallel_data = []
        if (progress_bar_disable is True):
            parallel_data = __do_parallel_work()
        else:
            if (progress_bar_format_numurls_nobytes is True):
                # total bytes could be inaccurate, progress bar should use count of urls
                # and iterator of 'files' instead of bytes.
                with srs_obj._tqdm(total=len(file_listing_obj.urls), desc=desc_str, unit="files", ncols=progress_bar_ncols,
                                   ascii=progress_bar_ascii) as pbar:
                    parallel_data = __do_parallel_work(pbar=pbar, pbar_iterator_nfiles=True)
            else:
                # total bytes is accurate, show MB/s
                with srs_obj._tqdm(total=file_listing_obj.total_bytes,
                                   desc=desc_str,
                                   ncols=progress_bar_ncols,
                                   ascii=progress_bar_ascii,
                                   unit="B",
                                   unit_scale=True) as pbar:
                    parallel_data = __do_parallel_work(pbar=pbar)

        # cast into return obj
        filenames_list = []
        total_bytes_downloaded = 0
        for p in parallel_data:
            filenames_list.append(p["filename"])  # type: ignore
            total_bytes_downloaded += p["bytes_downloaded"]  # type: ignore
        download_obj = FileDownloadResult(
            filenames=filenames_list,
            count=len(parallel_data),
            dataset=file_listing_obj.dataset,
            total_bytes=total_bytes_downloaded,
            output_root_path=output_path,
        )

    # return
    return download_obj


def get_urls(srs_obj, dataset_name, start, end, site_uid, device_uid, timeout, warning_stack_level=3):
    # set timeout
    if (timeout is None):
        timeout = srs_obj.api_timeout

    # set up API file listing request
    params = {
        "name": dataset_name,
        "start": start,
        "end": end,
        "include_total_bytes": True,
    }
    if (site_uid is not None):
        params["site_uid"] = site_uid
    if (device_uid is not None):
        params["device_uid"] = device_uid

    # check warnings about site and device being supplied
    if (site_uid is not None and "calibration" in dataset_name.lower()):
        show_warning("The site_uid filter was used when retrieving URLs for the dataset %s. Please note "
                     "that the site_uid field is not used when filtering in this dataset. Consider removing "
                     "it from your function call." % (dataset_name),
                     stacklevel=warning_stack_level)
    if (device_uid is not None and "skymap" in dataset_name.lower()):
        show_warning("The device_uid filter was used when retrieving URLs for the dataset %s. Please note "
                     "that the device_uid field is not used when filtering in this dataset. Consider removing "
                     "it from your function call." % (dataset_name),
                     stacklevel=warning_stack_level)

    # make API request
    url = "%s/api/v1/data_distribution/urls" % (srs_obj.api_base_url)
    try:
        r = requests.get(url, params=params, headers=srs_obj.api_headers, timeout=timeout)
    except Exception as e:  # pragma: nocover
        raise SRSAPIError("Unexpected API error: %s" % (str(e))) from e
    if (r.status_code != 200):  # pragma: nocover
        try:
            res = r.json()
            msg = res["detail"]
        except Exception:
            msg = r.content
        raise SRSAPIError("API error code %d: %s" % (r.status_code, msg))
    res = r.json()

    # get list of file reading supported datasets
    file_reading_supported_datasets = srs_obj.data.list_supported_read_datasets()

    # cast response into FileDownloadResult object
    file_listing_obj = FileListingResponse(**res)

    # cast dataset part of response
    file_reading_supported = True if res["dataset"]["name"] in file_reading_supported_datasets else False
    file_listing_obj.dataset = Dataset(**res["dataset"], file_reading_supported=file_reading_supported)

    # return
    return file_listing_obj


def download_generic(srs_obj, dataset_name, start, end, site_uid, device_uid, n_parallel, overwrite, progress_bar_disable, progress_bar_ncols,
                     progress_bar_ascii, progress_bar_desc, timeout):
    # get file listing
    file_listing_obj = get_urls(srs_obj, dataset_name, start, end, site_uid, device_uid, timeout, warning_stack_level=4)

    # download the urls
    download_obj = __download_urls(
        srs_obj,
        file_listing_obj,
        n_parallel,
        overwrite,
        progress_bar_disable,
        progress_bar_ncols,
        progress_bar_ascii,
        progress_bar_desc,
        timeout,
    )

    # return
    return download_obj


def download_using_urls(srs_obj, file_listing_obj, n_parallel, overwrite, progress_bar_disable, progress_bar_ncols, progress_bar_ascii,
                        progress_bar_desc, timeout):
    # download the urls
    download_obj = __download_urls(srs_obj,
                                   file_listing_obj,
                                   n_parallel,
                                   overwrite,
                                   progress_bar_disable,
                                   progress_bar_ncols,
                                   progress_bar_ascii,
                                   progress_bar_desc,
                                   timeout,
                                   progress_bar_format_numurls_nobytes=True)

    # return
    return download_obj
