import os
import time
import requests
import xml.etree.ElementTree as ET
import logging
import collections
import json
from tqdm import tqdm

try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

import pandas as pd

from bird_tool_utils import iterable_chunks

# Define these constants so that they can be referred to in other classes
# without index errors.
STUDY_ACCESSION_KEY = 'study_accession'
RUN_ACCESSION_KEY = 'run'
BASES_KEY = 'bases'
SAMPLE_NAME_KEY = 'sample_name'
NCBI_API_KEY_ENV = 'NCBI_API_KEY'
BIOPROJECT_ACCESSION_KEY = 'bioproject'

class SraMetadata:
    def add_api_key(self, other_params):
        if NCBI_API_KEY_ENV in os.environ:
            other_params['api_key'] = os.environ[NCBI_API_KEY_ENV]
        return other_params

    def _retry_request(self, description, func):
        '''Retry a reqests.post or requests.get 3 times, returning the request
        when OK, otherwise raising an Exception'''

        num_retries = 3
        sleep_time = 60
        def retrying(i, num_retries=3):
            if i < num_retries-1:
                logging.warning("Retrying request (retry {} of {})".format(i+1, num_retries-1))
        
        for i in range(num_retries):
            try:
                this_res = func()
                if not this_res.ok:
                    logging.warning("Request not OK when {}: {}: {}".format(description, this_res, this_res.text))
                    logging.warning("Sleeping for {} seconds before retrying".format(sleep_time))
                    time.sleep(60)
                    retrying(i)
                else:
                    return this_res
            except Exception as e:
                logging.warning("Exception raised when {}: {}".format(description, e))
                logging.warning("Sleeping for {} seconds before retrying".format(sleep_time))
                time.sleep(60)
                retrying(i)
        raise Exception("Failed to {} after {} attempts".format(description, num_retries))


    def fetch_runs_from_bioprojects(self, bioproject_accessions):
        retmax = 10000
        query_string = " OR ".join(["{}[BioProject]".format(bioproject_accession) for bioproject_accession in bioproject_accessions])
        logging.debug("Querying with string: {}".format(query_string))
        res = requests.get(
            url="https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
            params=self.add_api_key({
                "db": "sra",
                "term": query_string,
                "tool": "kingfisher",
                "email": "kingfisher@github.com",
                "retmax": retmax,
                "usehistory": "y",
                }),
            )
        if not res.ok:
            raise Exception("HTTP Failure when requesting search from bioproject: {}: {}".format(res, res.text))
        root = ET.fromstring(res.text)
        sra_ids = list([c.text for c in root.find('IdList')])
        if len(sra_ids) == retmax:
            logging.warning("Unexpectedly found the maximum number of results for this query, possibly some results will be missing")
        webenv = root.find('WebEnv').text

        # Now convert the IDs into runs
        metadata = self.efetch_metadata_from_ids(webenv, None, len(sra_ids))
        return metadata[RUN_ACCESSION_KEY].to_list()

    def efetch_metadata_from_ids(self, webenv, accessions, num_ids):
        data_frames = []

        retmax = num_ids+10
        logging.debug("Running efetch ..")
        res = self._retry_request(
            'efetch_from_ids',
            lambda: requests.get(
                url="https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
                params=self.add_api_key({
                    "db": "sra",
                    "tool": "kingfisher",
                    "email": "kingfisher@github.com",
                    "webenv": webenv,
                    "query_key": 1
                    }),
                ))
        if not res.ok:
            raise Exception("HTTP Failure when requesting efetch from IDs: {}: {}".format(res, res.text))

        root = ET.fromstring(res.text)

        def try_get(func):
            try:
                return func()
            except AttributeError:
                return ''
            except KeyError:
                return None

        if root.find("ERROR") is not None:
            logging.error("Error when fetching metadata: {}".format(root.find("ERROR").text))

        # Some samples such as SAMN13241871 are linked to multiple runs e.g. SRR10489833
        accessions_set = None if accessions is None else set(accessions)

        for pkg in root.findall('EXPERIMENT_PACKAGE'):
            d = collections.OrderedDict()
            d['experiment_accession'] = try_get(lambda: pkg.find('./EXPERIMENT').attrib['accession'])
            d['experiment_title'] = try_get(lambda: pkg.find('./EXPERIMENT/TITLE').text)
            l = pkg.find('./EXPERIMENT/DESIGN/LIBRARY_DESCRIPTOR')
            d['library_name'] = try_get(lambda: l.find('LIBRARY_NAME').text)
            d['library_strategy'] = try_get(lambda: l.find('LIBRARY_STRATEGY').text)
            d['library_source'] = try_get(lambda: l.find('LIBRARY_SOURCE').text)
            d['library_selection'] = try_get(lambda: l.find('LIBRARY_SELECTION').text)
            d['library_layout'] = try_get(lambda: l.find('LIBRARY_LAYOUT')[0].tag)
            d['platform'] = try_get(lambda: pkg.find('./EXPERIMENT/PLATFORM')[0].tag)
            d['model'] = try_get(lambda: pkg.find('./EXPERIMENT/PLATFORM/')[0].text)
            d['submitter'] = ''
            for k, v in pkg.find('./SUBMISSION').attrib.items():
                if k not in ('accession','alias'):
                    if d['submitter'] == '':
                        d['submitter'] = v
                    else:
                        d['submitter'] = "{}, {}".format(d['submitter'], v)
            d[STUDY_ACCESSION_KEY] = try_get(lambda: pkg.find('./STUDY').attrib['accession'])
            d[BIOPROJECT_ACCESSION_KEY] = try_get(lambda: pkg.find('./STUDY/IDENTIFIERS/EXTERNAL_ID[@namespace="BioProject"]').text)
            d['study_alias'] = try_get(lambda: pkg.find('./STUDY').attrib['alias'])
            d['study_centre_project_name'] = try_get(lambda: pkg.find('./STUDY/DESCRIPTOR/CENTER_PROJECT_NAME').text)
            d['organisation'] = try_get(lambda: pkg.find('./Organization/Name').text)
            d['organisation_department'] = try_get(lambda: pkg.find('./Organization/Address/Department').text)
            d['organisation_institution'] = try_get(lambda: pkg.find('./Organization/Address/Institution').text)
            d['organisation_street'] = try_get(lambda: pkg.find('./Organization/Address/Street').text)
            d['organisation_city'] = try_get(lambda: pkg.find('./Organization/Address/City').text)
            d['organisation_country'] = try_get(lambda: pkg.find('./Organization/Address/Country').text)
            first_name = try_get(lambda: pkg.find('./Organization/Contact/Name/First').text)
            last_name = try_get(lambda: pkg.find('./Organization/Contact/Name/Last').text)
            d['organisation_contact_name'] = "{} {}".format(first_name, last_name)
            d['organisation_contact_email'] = try_get(lambda: pkg.find('./Organization/Contact').attrib['email'])
            # Remove carriage return from sample description e.g. for SRR1263047
            d['sample_description'] = try_get(lambda: pkg.find('./SAMPLE/DESCRIPTION').text.replace('\r',''))
            d['sample_alias'] = try_get(lambda: pkg.find('./SAMPLE').attrib['alias'])
            d['sample_accession'] = try_get(lambda: pkg.find('./SAMPLE').attrib['accession'])
            d['biosample'] = try_get(lambda: pkg.find('./SAMPLE/IDENTIFIERS/EXTERNAL_ID[@namespace="BioSample"]').text)
            d['sample_title'] = try_get(lambda: pkg.find('./SAMPLE/TITLE').text)
            d['taxon_name'] = try_get(lambda: pkg.find('./SAMPLE/SAMPLE_NAME/SCIENTIFIC_NAME').text)
            sample_sample_name = None
            sample_title = None
            if pkg.find('./SAMPLE/SAMPLE_ATTRIBUTES'):
                for attr in pkg.find('./SAMPLE/SAMPLE_ATTRIBUTES'):
                    tag_el = attr.find('TAG')
                    value_el = attr.find('VALUE')

                    # Some samples like
                    # https://www.ncbi.nlm.nih.gov/sra?term=ERS1240061&report=FullXml
                    # have entries with a tag, but no value. Ignore these.
                    if tag_el is not None and value_el is not None:
                        tag = tag_el.text
                        value = value_el.text
                        if tag == 'Title':
                            sample_title = value
                        elif tag == 'sample name':
                            sample_sample_name = value
                        d[tag] = value
            if sample_sample_name is not None:
                d[SAMPLE_NAME_KEY] = sample_sample_name
            elif sample_title is not None:
                d[SAMPLE_NAME_KEY] = sample_title
            else:
                d[SAMPLE_NAME_KEY] = d['library_name'] #default, maybe there's always a title though?
            d['study_title'] = try_get(lambda: pkg.find('./STUDY/DESCRIPTOR/STUDY_TITLE').text)
            d['design_description'] = try_get(lambda: pkg.find('./EXPERIMENT/DESIGN/DESIGN_DESCRIPTION').text)
            d['study_abstract'] = try_get(lambda: pkg.find('./STUDY/DESCRIPTOR/STUDY_ABSTRACT').text)
            study_links_xrefs = try_get(lambda: pkg.find('./STUDY/STUDY_LINKS'))
            if study_links_xrefs is not None:
                # Convert db to lower case because otherwise have PUBMED and pubmed e.g. ERR1914274 and SRR9113719
                study_links = list([
                    {'db': x.find('DB').text.lower(), 'id': x.find('ID').text}
                    for x in study_links_xrefs.findall('STUDY_LINK/XREF_LINK')])
                # Record URL links like SRR7051324
                for x in study_links_xrefs.findall('STUDY_LINK/URL_LINK'):
                    study_links.append({
                        'label': x.find('LABEL').text,
                        'url': x.find('URL').text
                    })
                d['study_links'] = json.dumps(study_links)
            else:
                d['study_links'] = json.dumps([])
            
            # Account for the fact that multiple runs may be associated with
            # this sample
            d['number_of_runs_for_sample'] = len(pkg.findall('./RUN_SET/RUN'))

            for run in pkg.findall('./RUN_SET/RUN'):
                accession_here = run.attrib['accession']
                if accessions_set is None or accession_here in accessions_set:
                    d2 = d.copy()
                    d2['spots'] = try_get(lambda: int(run.attrib['total_spots']))
                    d2[BASES_KEY] = try_get(lambda: int(run.attrib['total_bases']))
                    d2['run_size'] = try_get(lambda: int(run.attrib['size']))
                    d2[RUN_ACCESSION_KEY] = try_get(lambda: run.attrib['accession'])
                    d2['published'] = try_get(lambda: run.attrib['published'])
                    stats = run.find('Statistics')
                    if stats is not None:
                        for (i, r) in enumerate(stats):
                            d2['read{}_length_average'.format(i+1)] = r.attrib['average']
                            d2['read{}_length_stdev'.format(i+1)] = r.attrib['stdev']
                    data_frames.append(d2)

        return pd.DataFrame(data_frames)

    def _print_xml(self, element, prefix):
        if prefix is None or prefix == '':
            p2 = ""
        else:
            p2 = "{}_".format(prefix)
        for k, v in element.attrib.items():
            print("\t".join(['{}{}'.format(p2,k),v]))
        if element.text:
            print("\t".join(['{}{}'.format(p2, element.tag), element.text]))
        for e in element:
            self.print_xml(e, '{}{}'.format(p2, e.tag))

    def efetch_sra_from_accessions(self, accessions):
        all_accessions = list(set(accessions))
        if len(all_accessions) == 0:
            return []

        metadata_chunks = []
        chunk_size = 500

        # Complicated calculation here.
        num_chunks = len(all_accessions)/chunk_size
        if num_chunks > int(num_chunks):
            num_chunks = int(num_chunks) + 1
        num_chunks = int(num_chunks)

        if len(all_accessions) > chunk_size:
            logging.info("Querying for {} accessions in {} chunks".format(len(all_accessions), num_chunks))
            accession_iter = tqdm(iterable_chunks(all_accessions, chunk_size), total=num_chunks)
        else:
            accession_iter = [all_accessions]

        for accessions_plus in accession_iter:
            accessions = [a for a in accessions_plus if a is not None]

            if num_chunks == 1:
                logging.info("Querying NCBI esearch for {} distinct accessions e.g. {}".format(
                    len(accessions), accessions[0]))
            sra_ids = []

            webenv = None
            request_term = ' OR '.join(["{}[accn]".format(acc) for acc in accessions])

            retmax = len(accessions)+10
            params=self.add_api_key({
                "db": "sra",
                "term": request_term,
                "tool": "kingfisher",
                "email": "kingfisher@github.com",
                "retmax": retmax,
                "usehistory": "y",
                })
            if webenv is None:
                params['WebEnv'] = webenv

            res = self._retry_request(
                "esearch from accessions", 
                lambda: requests.post(
                    url="https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
                    data=params))

            root = ET.fromstring(res.text)
            if webenv is None:
                webenv = root.find('WebEnv').text
            id_list_node = root.find('IdList')
            sra_ids = list(set([c.text for c in id_list_node]))

            if len(sra_ids) == 0:
                logging.warning("Unable to find any accessions, from the list: {}".format(accessions))
                return None

            if num_chunks == 1:
                logging.info("Querying NCBI efetch for {} distinct IDs e.g. {}".format(
                    len(sra_ids), sra_ids[0]))
            metadata = self.efetch_metadata_from_ids(webenv, accessions, len(sra_ids))

            # Ensure all hits are found, and trim results to just those that are real hits
            if RUN_ACCESSION_KEY not in metadata.columns:
                raise Exception("No metadata could be retrieved")

            if len(metadata) != len(accessions):
                found_runs = set(metadata[RUN_ACCESSION_KEY].to_list())
                not_found = list([a for a in accessions if a not in found_runs])
                logging.warning("Unable to find all accessions. The {} missing ones were: {}".format(
                    len(not_found), not_found
                ))
            metadata_chunks.append(metadata)

        metadata = pd.concat(metadata_chunks)
        metadata.sort_values([STUDY_ACCESSION_KEY,RUN_ACCESSION_KEY], inplace=True)

        return metadata
    
    def fetch_pubmed_ids_from_term(self, term):
        retmax = 10000
        res = requests.get(
            url="https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
            params=self.add_api_key({
                "db": "pubmed",
                "term": term,
                "tool": "kingfisher",
                "email": "kingfisher@github.com",
                "retmax": retmax,
                "usehistory": "y",
                }),
            )
        if not res.ok:
            raise Exception("HTTP Failure when requesting search from term: {}: {}".format(res, res.text))
        root = ET.fromstring(res.text)
        logging.debug("Root of response: {}".format(ET.tostring(root)))
        pubmed_ids = list([c.text for c in root.find('IdList')])
        if len(pubmed_ids) == retmax:
            logging.warning("Unexpectedly found the maximum number of results for this query, possibly some results will be missing")
        return pubmed_ids

    def fetch_citations_from_query_title(self, title):
        # https://www.ebi.ac.uk/europepmc/webservices/rest/search?query=Genome-centric%20view%20of%20carbon%20processing%20in%20thawing%20permafrost&format=json

        # Search for the title using the ENA rest API. Found it to be superior to the NCBI esearch e.g. the query 'Metagenomics of Urban Sewage Identifies an Extensively Shared Antibiotic Resistome in China' hits on the PubMed website, but not in the NCBI esearch - unsure why. Worked out of the box with the ENA rest API.

        res = requests.get(
            url="https://www.ebi.ac.uk/europepmc/webservices/rest/search",
            params={
                "query": title,
                "format": "json",
                },
            )
        if not res.ok:
            raise Exception("HTTP Failure when requesting search from term: {}: {}".format(res, res.text))
        root = res.json()
        logging.debug("Root of response: {}".format(root))

        # Return only those that have an exact title match
        citations = []
        for result in root['resultList']['result']:
            logging.debug("Title: {}".format(result['title']))
            if result['title'].lower() == title.lower() or result['title'].lower() == title.lower() + '.':
                citations.append(result)
        return citations

    def fetch_citations_from_query_bioproject(self, bioproject):

        res = requests.get(
            url="https://www.ebi.ac.uk/europepmc/webservices/rest/search",
            params={
                "query": bioproject,
                "format": "json",
                },
            )
        if not res.ok:
            raise Exception("HTTP Failure when requesting search from term: {}: {}".format(res, res.text))
        root = res.json()
        logging.debug("Root of response: {}".format(root))

        # Return all hits
        return root['resultList']['result']