import asyncio, requests_html_playwright, more_itertools, datetime, time, sys
from boldigger3.id_engine import parse_fasta
from bs4 import BeautifulSoup as BSoup
from requests import Response
from tqdm.asyncio import tqdm
import pandas as pd
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from requests.exceptions import RetryError
from requests.exceptions import ChunkedEncodingError, ConnectionError


# function to collect the process ids from the hdf storage
def collect_process_ids(hdf_name_results: str) -> list:
    """Function to collect all process ids from the downloaded data.

    Args:
        hdf_name_results (str): Path to the hdf storage that's generated by the data download

    Returns:
        tuple: Returns a list of unique process ids.
    """

    # read the process ids from hdf storage
    unsorted_results = pd.read_hdf(
        hdf_name_results,
        key="results_unsorted",
        columns=["process_id"],
        where="process_id != ''",
        iterator=True,
        chunksize=1000000,
    )

    # define the unique process ids variable
    unique_process_ids = set()

    # remove duplicates and empty strings from process ids
    for chunk in unsorted_results:
        unique_process_ids = unique_process_ids.union(
            set(chunk["process_id"].to_list())
        )

    # return the ids
    return unique_process_ids


# function to check if the additional data or parts of it have already been downloaded
def check_already_downloaded(hdf_name_results: str, unique_process_ids: list) -> list:
    """Function that checks if the requested download has already been performed

    Args:
        hdf_name_results (str): Path to the hdf data storage.
        unique_process_ids (list): List of unique process ids that are requested for data download.

    Returns:
        list: Returns a list of the unique process ids that have not been downloaded yet.
    """
    try:
        additional_data = pd.read_hdf(
            hdf_name_results,
            key="additional_data",
            columns=["process_id"],
            iterator=True,
            chunksize=1000000,
        )

        # collect the already downloaded ids
        downloaded_ids = set()

        # loop over all chunks of downloaded ids
        for chunk in additional_data:
            downloaded_ids = downloaded_ids.union(set(chunk["process_id"].to_list()))

        # filter unique ids
        unique_process_ids = [
            unique_id
            for unique_id in unique_process_ids
            if unique_id not in downloaded_ids
        ]

        return unique_process_ids
    except KeyError:
        return unique_process_ids


# funtion to parse a record page
def parse_record_page(html_response: object, url: str) -> list:
    """This function parses a page for any given record on bold

    Args:
        html_response (object): Requests response object to parse
        url (string): url as string to parse the process id

    Returns:
        list: A line of data representing the record data
    """
    # extract the process id from the url
    process_id = url.split("/")[-1]

    # get the status from the status code, all retries exceeded will be routed to 500
    status_code_to_status = {
        200: "public",
        404: "private",
        500: "unavailable",
    }
    status = status_code_to_status[html_response.status_code]

    # testing
    if html_response.status_code == 200:
        # transform html into beautifulsoup to parse out the data
        soup = BSoup(html_response.text, "html5lib")

        # find the specimen table
        specimen = soup.find("h1", string="Specimen").find_next("table")
        # extract the sex and lifestage
        sex = specimen.find("th", string="Sex:").find_next("td").text
        lifestage = specimen.find("th", string="Life Stage:").find_next("td").text
        # extract institution storing
        identifiers = soup.find("h1", string="Identifiers").find_next("table")
        institution_storing = (
            identifiers.find("th", string="Deposited In:").find_next("td").text
        )
        # extract the country
        collection = soup.find("h1", string="Collection").find_next("table")
        country_ocean = (
            collection.find("th", string="Country/Ocean:").find_next("td").text
        )
        # extract identifier
        attribution = soup.find("h1", string="Attribution").find_next("table")
        identifier = (
            attribution.find("th", string="Specimen Identification:")
            .find_next("td")
            .text
        )
        # extract identification method
        taxonomy = soup.find("h1", string="Taxonomy").find_next("table")
        id_method = (
            taxonomy.find("th", string="Identification Method:").find_next("td").text
        )
        # record page
        record_page = "https://portal.boldsystems.org/record/{}".format(process_id)
        # return the resulting line
        return [
            process_id,
            status,
            sex,
            lifestage,
            institution_storing,
            country_ocean,
            identifier,
            id_method,
            record_page,
        ]
    else:
        return [process_id, status] + [""] * 7


# async function to perform the request
async def as_request(url, as_session) -> list:
    """Function to perform async requests and parse the data.

    Args:
        url (_type_): url to request
        as_session (_type_): async session to perform the request with

    Returns:
        list: Returns the parsed response as list.
    """
    # request the url from BOLD
    try:
        response = await as_session.get(url, timeout=3)
    except RetryError:
        response = Response()
        response.status_code = 500
    except (ChunkedEncodingError, ConnectionError):
        response = Response()
        response.status_code = 500

    # parse the response here
    response = parse_record_page(response, url)

    # return the response to the caller to append it to hdf
    return response


# function to limit concurrency of the code
async def limit_concurrency(url, as_session, semaphore) -> object:
    """Function to limit the concurrency of requests.

    Args:
        url (_type_): url to request.
        as_session (_type_): session to use for the requests
        semaphore (_type_): semaphore object

    Returns:
        function: as request with semaphore added
    """
    async with semaphore:
        return await as_request(url, as_session)


# function to launch the async session including retry strategy
async def as_session(download_urls, semaphore) -> list:
    """Function to launch the async session and perform the data download.

    Args:
        download_urls (_type_): All urls to download.
        semaphore (_type_): semaphore to limit the concurrency

    Returns:
        list: returns a list of list with the gathered responses of the batch downloaded
    """
    # create the session with correct headers and a retry strategy
    as_session = requests_html_playwright.AsyncHTMLSession()
    as_session.headers.update(
        {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/98.0.4758.82 Safari/537.36"
        }
    )

    # only retry if the page failed to load
    retry_strategy = Retry(
        total=5,
        status_forcelist=[500, 502, 504],
        backoff_factor=2,
    )
    adapter = HTTPAdapter(max_retries=retry_strategy)
    as_session.mount("https://", adapter)

    tasks = (limit_concurrency(url, as_session, semaphore) for url in download_urls)

    return await tqdm.gather(*tasks, desc="Downloading additional data")


# function to transform the response to a pandas dataframe
def response_to_dataframe(parsed_response: list) -> object:
    """Function that transforms a list of rows to dataframe for easy saving

    Args:
        parsed_response (list): List of lists that is  returned from parsing BOLD

    Returns:
        object: Returns a pandas dataframe that can be easily saved in hdf format.
    """
    # generate dataframe
    parsed_response = pd.DataFrame(
        parsed_response,
        columns=[
            "process_id",
            "status",
            "sex",
            "lifestage",
            "institution_storing",
            "country_or_ocean",
            "identifier",
            "id_method",
            "record_page",
        ],
    )

    return parsed_response


# function to add the data to the hdf storage
def add_to_hdf(hdf_name_results: str, response: object) -> None:
    """Function to add the downloaded additional data to the hdf storage.

    Args:
        hdf_name_results (str): Path to the hdf storage.
        response (object): Dataframe to be appended to the hdf storage.
    """
    # define the item sizes for the hdf storage
    item_sizes = {
        "process_id": 30,
        "status": 11,
        "sex": 8,
        "lifestage": 80,
        "institution_storing": 150,
        "country_or_ocean": 80,
        "identifier": 80,
        "id_method": 400,
        "record_page": 70,
    }

    # save results in hdf storage
    with pd.HDFStore(
        hdf_name_results, mode="a", complib="blosc:blosclz", complevel=9
    ) as hdf_output:
        hdf_output.append(
            value=response,
            key="additional_data",
            format="t",
            data_columns=True,
            min_itemsize=item_sizes,
            complib="blosc:blosclz",
            complevel=9,
        )


# main functio to run the additional data download
def main(fasta_path: str) -> None:
    """Main function to run the additional data download. Downloads additional data
    for all public process ids, saves the to hdf first and finally add them to the
    results.

    Args:
        hdf_name_results (str): Path to the fasta file to be identified
    """
    # user output
    tqdm.write(
        "{}: Collecting process ids.".format(
            datetime.datetime.now().strftime("%H:%M:%S")
        )
    )

    # read the input fasta
    fasta_dict, fasta_name, project_directory = parse_fasta(fasta_path)

    # generate a new for the hdf storage to store the downloaded data
    hdf_name_results = project_directory.joinpath(
        "{}_result_storage.h5.lz".format(fasta_name)
    )

    # collect all process ids first
    unique_process_ids = collect_process_ids(hdf_name_results)

    # check already downloaded data
    unique_process_ids = check_already_downloaded(hdf_name_results, unique_process_ids)

    # calculate the number of batches
    batch_count = len(list(more_itertools.chunked(unique_process_ids, 10000)))
    batch_counter = 1

    # only start the download loop if there are ids to downlaod
    if unique_process_ids:
        # user output
        tqdm.write(
            "{}: Divided {} unique process ids into {} batch(es) for download.".format(
                datetime.datetime.now().strftime("%H:%M:%S"),
                len(unique_process_ids),
                batch_count,
            )
        )

        urls = [
            "https://portal.boldsystems.org/record/{}".format(idx)
            for idx in unique_process_ids
        ]

        for chunk in more_itertools.chunked(urls, 10000):
            tqdm.write(
                "{}: Downloading batch {} of {}.".format(
                    datetime.datetime.now().strftime("%H:%M:%S"),
                    batch_counter,
                    batch_count,
                )
            )
            # requests the additional data asynchronously in batches of 10000 urls
            # limit the concurrent requests to 50
            semaphore = asyncio.Semaphore(50)

            # gather the responses
            response = asyncio.run(as_session(chunk, semaphore))

            # transform to dataframe for saving
            response = response_to_dataframe(response)

            # saving the data
            add_to_hdf(hdf_name_results, response)

            # user output
            tqdm.write(
                "{}: Saving batch {} of {}.".format(
                    datetime.datetime.now().strftime("%H:%M:%S"),
                    batch_counter,
                    batch_count,
                )
            )

            # increase batch counter for user output
            batch_counter += 1
    else:
        # user output
        tqdm.write(
            "{}: Additional data has already been completly downloaded.".format(
                datetime.datetime.now().strftime("%H:%M:%S"),
                batch_counter,
                batch_count,
            )
        )
