"""
Wrapper around version 2 of the Energy Aspects API

"""

import os
import re
import json
from typing import List
from functools import reduce, lru_cache

import pandas as pd
import requests

base_url = "https://api.energyaspects.com/data/datasets/timeseries/"
timeseries_url = "https://api.energyaspects.com/data/timeseries/"
api_dataset_mapping_endpoint = "https://api.energyaspects.com/data/dataset_mappings/"

api_key = os.getenv('EA_API_KEY')

headers = {"Accept": "application/json"}

geo_map = {
    'WORLD': 'Global',
    'NA': 'Northern America',
    'EUR': 'Europe',
}


@lru_cache(maxsize=1)
def get_all_webfile_mappings() -> dict:
    """
    Gets webfile mappings from datasets mapping endpoint

    Returns:
        Dict response from datasets mapping endpoint
    """
    query = api_dataset_mapping_endpoint + "?api_key=" + api_key
    headers = {"accept": "application/json"}
    response = requests.get(query, headers=headers)
    all_webfile_mappings = json.loads(response.content)

    return all_webfile_mappings


def get_all_webfile_mappings_names():
    """
    Simplifies the input data to a dictionary with keys as the main categories
    and values as a list of names associated with those categories.

    :param data: The original dictionary containing detailed mappings.
    :return: A simplified dictionary with categories as keys and a list of names as values.
    """
    simplified = {}
    data = get_all_webfile_mappings()
    for key, value in data.items():
        simplified[key] = [item['name'] for item in value]
    return simplified


def get_dataset_ids_for_dataset(all_webfile_mappings: dict, dataset_name: str):
    """
    Retrieve the dataset IDs associated with a given dataset

    Args:
        all_webfile_mappings: Response from dataset mapping endpoint
        dataset_name: The name of the file to search for.

    Returns:
        A list of dataset IDs corresponding to the given file name.
    """
    dataset_ids = next(
        (dataset_mapping["dataset_ids"] for category in all_webfile_mappings.values() for dataset_mapping in category
         if dataset_mapping["name"] == dataset_name), None)
    return dataset_ids


def get_data_set_mapping(web_file_name: str, category: str) -> List[int]:
    """
    Gets the dataset ids from the API for a given webfile name and category
    Args:
        web_file_name: Name of webfile you want to extract
        category: Name of data category.
    Returns:
        Dataset ids you want to extract data for
    """
    query = api_dataset_mapping_endpoint + "?api_key=" + api_key
    headers = {"accept": "application/json"}
    response = requests.get(query, headers=headers)
    all_webfile_mappings = json.loads(response.content)

    try:
        categories_to_search = all_webfile_mappings[category]
    except KeyError:
        raise KeyError("Please provide valid Category")
    required_job = None
    # Iterate through the list of dictionaries
    for job in categories_to_search:
        if web_file_name == job['name']:
            required_job = job
            break  # Stop searching once the key is found
    if not required_job:
        raise ValueError("Please provide a valid job name")

    dataset_ids = required_job['dataset_ids']
    return dataset_ids


def get_timeseries_def():
    querystring = {
        "api_key": api_key,
    }
    response = requests.request("GET", base_url, headers=headers, params=querystring)

    alls = pd.read_json(response.text)
    alls = pd.merge(alls, alls['metadata'].apply(lambda x: pd.Series(x), 1), left_index=True, right_index=True)
    alls = alls.set_index('dataset_id')
    alls = alls.drop('metadata', axis=1)
    return alls


def extract_region(m):
    region = m['country_iso']  # by default use country
    if not m['country_iso']:
        if m['region']:  # if country but do have a region use it
            region = m['region']
        else:
            region = m['description']  # some have neither - use desc

    if m['country_iso'] == 'US' and m['sub_region']:  # for padds
        region = m['sub_region']

    return region


def generate_title(dataset_id: int, metadata: dict, format_title: tuple, as_tuple: bool = False):
    """
    Given metadata generate a title for use in data frames.
    eg region.aspect would give 'US.demand'
    :param dataset_id:
    :param metadata:
    :param format_title:
    :return:
    """
    m = metadata[dataset_id]
    title = '.'.join(format_title)
    for token in format_title:
        if token == 'extract_region':
            title = re.sub(token, extract_region(m), title)
        else:
            title = re.sub(token, m[token], title)

    if as_tuple:
        title = tuple(title.split('.'))
    return title


def extract_data(s, id_):
    r = pd.Series(s)
    r.index = pd.to_datetime(r.index)
    r.name = id_
    return r


def format_columns(df: pd.DataFrame, metadata: dict, column_format: dict):
    column_format_value = column_format['format']
    if isinstance(column_format_value, str):
        column_format_value = tuple(column_format_value)
    multiindex = True if 'multiindex' in column_format and column_format['multiindex'] else False
    cmap = {x: generate_title(x, metadata, format_title=column_format_value, as_tuple=multiindex) for x in df.columns}
    # if making it multi index:
    if multiindex:
        index = pd.MultiIndex.from_tuples([cmap[x] for x in df.columns], names=column_format['format'])
        df.columns = index
    else:
        df = df.rename(columns=cmap)

    for t in cmap:  # add this title back to metadata to link metadata to columns
        metadata[t]['title'] = cmap[t]

    return df, metadata


def read_timeseries(column_format: dict = None, **kwargs) -> pd.DataFrame():
    querystring = {"api_key": api_key}
    for kwarg in kwargs:
        if kwarg == 'dataset_id':
            querystring["dataset_id"] = [str(x) for x in kwargs[kwarg]]
        if kwarg in ['category', 'category_detail', 'frequency', 'aspect', 'aspect_detail']:
            # note category_detail can't be specified without category
            querystring[kwarg] = kwargs[kwarg]
        if kwarg in ['date_from', "date_to"]:
            querystring[kwarg] = kwargs[kwarg].strftime("%Y-%m-%d")
        if kwarg in ['geography']:
            querystring["geography"] = [kwargs[kwarg]] if isinstance(kwargs[kwarg], str) else [str(x) for x in
                                                                                               kwargs[kwarg]]

    if 'category_detail' in kwargs and 'category' not in kwargs:
        querystring['category'] = 'crude_oil' if kwargs['category_detail'] == 'crude' else 'oil_products'

    response = requests.request("GET", timeseries_url, headers=headers, params=querystring)
    if response.status_code == 200:
        df = pd.read_json(response.text)
        data = df.apply(lambda x: extract_data(x.data, x.dataset_id), 1)
        data.index = data.index.map(lambda x: df.loc[x].dataset_id)
        data = data.T

        meta = df[['metadata', 'dataset_id']].set_index('dataset_id').to_dict()['metadata']

        if column_format:
            data, meta = format_columns(data, meta, column_format)

        # keep the metadata in the attrs of df
        data.attrs['metadata'] = meta
        return data
    else:
        return (response.status_code, response.text)


def detailed_metadata(id_: int) -> pd.DataFrame:
    querystring = {"api_key": api_key, }

    response = requests.request("GET", '%s%s' % (base_url, id_), headers=headers, params=querystring)
    if response.status_code == 200:
        df = pd.read_json(response.text)
        df = df['metadata']
        df.name = id_
        data = pd.DataFrame(df)
        return data.T


def read_commod_balance(**kwargs) -> pd.DataFrame:
    if 'commod' in kwargs:
        kwargs['category_detail'] = kwargs['commod']
    if 'frequency' not in kwargs:
        kwargs['frequency'] = 'monthly'

    df = read_timeseries(**kwargs)
    return df


def filter_meta_data(df, column, value):
    """
    Given a EA dataframe returned by this API, filter the meta based on values in column param
    :param df:
    :param column:
    :return:
    """
    metadata = df.attrs['metadata']
    try:
        res = [metadata[x][value] for x in metadata if metadata[x]['title'] == column]
        res = res[0]
        return res
    except ValueError as ex:
        pass
    except TypeError as ex:
        pass
    except IndexError as ex:
        pass
    except KeyError as ex:
        pass


def crude_quality_streams(**kwargs) -> pd.DataFrame:
    if 'aspect' not in kwargs:
        kwargs['aspect'] = 'production'
    if 'frequency' not in kwargs:
        kwargs['frequency'] = 'yearly'

    if 'category_subtype' not in kwargs:
        kwargs['category_subtype'] = ['condensate',
                                      'heavy_medium_sour',
                                      'heavy_sour',
                                      'heavy_sweet',
                                      'light_medium_sour',
                                      'light_sour',
                                      'light_sweet',
                                      'medium_medium_sour',
                                      'medium_sour',
                                      'medium_sweet']

    if 'category' not in kwargs:
        kwargs['category'] = 'crude_oil'
    crude = read_timeseries(**kwargs)
    kwargs['category'] = 'condensate'
    condensate = read_timeseries(**kwargs)
    df = pd.concat([crude, condensate], axis=1)
    df.attrs['metadata'] = {**crude.attrs['metadata'], **condensate.attrs['metadata']}
    return df


def crude_quality_streams_by_country(**kwargs) -> pd.DataFrame:
    res = crude_quality_streams(**kwargs)
    m = res.attrs['metadata']
    dfs = []
    for cntry in set([m[x]['country_iso'] for x in m]):
        if pd.notnull(cntry) and cntry != '':
            ts = [int(x) for x in m if m[x]['country_iso'] == cntry]
            c = res[ts]
            c = c.rename(columns={x: m[x]['category_subtype'] for x in ts})
            if 'as_pct' in kwargs and kwargs['as_pct'] == True:
                c = c.div(c.sum(1), axis=0)
            c.columns = pd.MultiIndex.from_product([[cntry], c.columns])
            dfs.append(c)

    df = reduce(lambda left, right: pd.merge(left, right, left_index=True, right_index=True), dfs)
    return df
