import time
import logging

from Bio import Entrez
from functools import partial
from typing import Callable, Any, Union, Iterable
from os import PathLike

import multiprocessing as mp
import pandas as pd
import itertools as it

from . import concurrency
from . import dbutils


def uid_from_esearch(accession: str, db: str) -> list[list[str]]:
    """
    uses Bio.Entrez.esearch to request the database internal UID of the given accession
    see  https://www.ncbi.nlm.nih.gov/books/NBK25497/table/chapter2.T._entrez_unique_identifiers_ui/?report=objectonly
    for available databases
    
    :param accession:  accession to retrieve UID for
    :param db:         name of the Entrez database to retrieve the UID from
    
    :return:           list of lists of accession, uid and database name the uid was retrieved from
    """
    response_handle = Entrez.esearch(db = db, term = accession)
    uid_list = Entrez.read(response_handle)['IdList']
    return [[accession, uid, db] for uid in uid_list]


def map_accessions_to_uids(
    accessions: Iterable[pd.Series], 
    db: str, 
    outfilename: Union[PathLike, str], 
    filelock: Union[mp.Manager().Lock, None] = None
) -> None:
    """
    takes a list of pandas.Series object generated by pandas.DataFrame.iterrows and a path to an outputfile
    and uses uid_from_esearch to map the contained accessions to the SRA internal UID. If the function is used
    concurrently a Lock object must be passed via the filelock argument
    
    :param accessions:     list of pandas.Series objects containing the keys srx_accession, biosample_accession and geo_accession
    :param db:             string denoting the NCBI database to retrieve the UIDs from
    :param outfilename:    path to the file to write the retrieved UID mappings to
    :param filelock:       A Lock object used to savely write to the outfile in case of concurrent usage
    
    :return:               None
    """
    logging.info('starting mapping process')
    write_to_file = (
        concurrency.write_to_file_multiprocess if filelock 
        else concurrency.write_to_file_singleprocess
    )
    
    uid_list = []
    start = accessions[0][0]
    for i, row in accessions:
        srx = row.srx_accession
        samn = row.biosample_accession
        gsm = row.geo_accession
        
        if not any([srx, samn]):
            accession = gsm

        else:
            accession = srx if srx else samn
            
        mapped_uids = dbutils.retry(
            uid_from_esearch,
            accession, 
            db
        )
        
        uid_list.extend(mapped_uids)
        
        if not (i + 1) % 1000:
            write_to_file(
                outfilename,
                uid_list,
                filelock
            )
            uid_list = []
            logging.info(f'Written uids for accessions {start} to {i}')
    
    
    logging.info(f'finishing up. writing {len(uid_list)} remaining uid mappings')

    # write remaining uids if there are any
    if uid_list:
        write_to_file(
            outfilename,
            uid_list,
            filelock
        )
    

def data_from_esummary(
    web_env_info: dict[str, str], 
    db: str, 
    parse_function: Callable, 
    data_parsers: dict[str, Any], 
    n_uids: int, 
    chunksize: int = 5000
) -> pd.DataFrame:
    """
    retrieves all the data from eSummary associated with the UIDs posted to the Entrez history server

    :param web_env_info:        dictionary containing the keys 'QueryKey' and 'WebEnv' as returned by Entrez.epost
    :param db:                  string denoting the database the supplied UIDs are from
    :param parse_function:      function used to parse the esummary response
    :param data_parsers:        dictionary containing keys and whatever you parser needs to parse the response 
                                keys will be the columns of the returned dataframe
    :param n_uids:              number of UIDs posted to the history server
    :param chunksize:           size of the chunks of summaries to retrieve

    :return:                    pandas.DataFrame with columns equal to data_parsers keys and the parsed data as values
    """
    esummary = partial(
        Entrez.esummary,
        db = db,
        WebEnv = web_env_info['WebEnv'],
        query_key = web_env_info['QueryKey']
    )

    def esummaries_from_history(retstart, retmax):
        response_handle = dbutils.retry(
            esummary,
            retstart = retstart, 
            retmax = retmax
        )
        return Entrez.read(response_handle)
    
    data = {
        key.lower(): [] for key in data_parsers.keys()
    }
    
    for retmin, retmax in dbutils.get_chunklimits(n_uids, chunksize):
        time.sleep(5) # avoid too many requests error
        response = dbutils.retry(
            esummaries_from_history,
            retstart = retmin,
            retmax = retmax
        )
        parse_function(
            response,
            data,
            data_parsers
        )
        
    data = pd.DataFrame().from_dict(
        data,
        orient = 'columns'
    )
    
    # somehow we get duplicates, suspect the post call
    # because I did it twice so we might end up with 
    # duplicated uids in the webenv but not sure
    # anyway dropping duplicates here
    return data.drop_duplicates().reset_index(drop = True)


def uids_to_summaries(uid_list, db, parse_function, data_parsers):
    """
    retrieves summaries for uids from eSummary and returns a pandas.DataFrame with columns = keys of data_parsers
    Values in these columns are determined by parse_function

    :param uids:                iterable of uids to retrieve eSummaries for
    :param db:                  database to retrieve the summaries from
    :param parse_function:      function used to parse the summaries
    :param data_parsers:        parser dictionary forwarded to parse function
    :param chunksize:           size of the chunks of uids posted to the Entrez history server

    :return:                    pandas.DataFrame containing the summary parsing results
    """
    web_env_info = dbutils.create_webenv(uid_list, db)

    logging.debug(web_env_info)
    
    accessions = data_from_esummary(
        web_env_info,
        db,
        parse_function,
        data_parsers,
        len(uid_list)
    )

    return accessions


def summaries_from_uids(
    uids: Iterable[Union[int, str]], 
    db: str,
    parse_function: Callable, 
    data_parsers: dict[str, Any],
    chunksize: int = 50000
) -> pd.DataFrame:
    """
    retrieves summaries for uids from eSummary and returns a pandas.DataFrame with columns = keys of data_parsers
    Values in these columns are determined by parse_function

    :param uids:                iterable of uids to retrieve eSummaries for
    :param db:                  database to retrieve the summaries from
    :param parse_function:      function used to parse the summaries
    :param data_parsers:        parser dictionary forwarded to parse function
    :param chunksize:           size of the chunks of uids posted to the Entrez history server

    :return:                    pandas.DataFrame containing the summary parsing results
    """
    result_frames = dbutils.process_in_chunks(
        uids,
        uids_to_summaries,
        chunksize = chunksize,
        db = db,
        parse_function = parse_function,
        data_parsers = data_parsers,
    )

    return pd.concat(result_frames)
