import json
from datetime import datetime
from io import BytesIO
import pandas as pd
from cognite import config
from cognite._utils import InputError, get_aggregate_func_return_name, to_camel_case, to_snake_case
from cognite.v05 import files
from cognite.v05 import timeseries as time_series_v05
from cognite.v06 import time_series as time_series_v06


class TimeSeries:
    "Object for specifying a specific time series from the TimeSeries API when using a data spec.\n\n    Args:\n        id (int): id of time series to retrieve\n        aggregates (List[str]): Local aggregate functions to apply\n        missing_data_strategy(str): Missing data strategy to apply\n        label (str): name of the column in the resulting data frame when passed to data transfer service.\n\n    Examples:\n        When you supply a label the resulting data frames produced by data transfer service\n        will use the label as column names.\n    "

    def __init__(self, id, aggregates=None, missing_data_strategy=None, label=None):
        self.id = id
        self.aggregates = aggregates
        self.missing_data_strategy = missing_data_strategy
        self.label = label
        self._name = None


class TimeSeriesDataSpec:
    'Object for specifying data from the TimeSeries API when using a data spec.\n\n    Args:\n        time_series (List[data_transfer_service.TimeSeries]):  Time series\n        aggregates (List[str]): List of aggregate functions\n        granularity (str): Granularity of aggregates\n        missing_data_strategy (str): Missing data strategy to apply, can be "linearInterpolation" or "ffill"\n        start (Union[str, int, datetime]): Start time\n        end (Union[str, int, datetime]): end time\n        label (str): Label for this data spec\n\n    Examples:\n        When you specify a label, data transfer service will exhibit the following behaviour::\n\n            ts_data_spec = TimeSeriesDataSpec(..., label="some_label")\n            data_spec = DataSpec([ts_data_spec], ...)\n            dts = DataTransferService(data_spec)\n            dataframes = dts.get_dataframes\n            my_df = dataframes["some_label"]\n    '

    def __init__(
        self, time_series, aggregates, granularity, missing_data_strategy=None, start=None, end=None, label=None
    ):
        self.time_series = time_series
        self.aggregates = aggregates
        self.granularity = granularity
        self.missing_data_strategy = missing_data_strategy
        self.start = start
        self.end = end
        self.label = label or "default"


class FilesDataSpec:
    "Object for specifying data from the Files API when using a data spec.\n\n    Args:\n        file_ids (Dict[str, int]):    Dictionary of fileNames -> fileIds\n    "

    def __init__(self, file_ids):
        self.file_ids = file_ids


class DataSpec:
    "Object for specifying data when querying CDP.\n\n    Args:\n        time_series_data_specs (List[TimeSeriesDataSpec]):  Time Series data specs\n        files_data_spec (FilesDataSpec):                    Files data spec\n\n    Raises:\n        DataSpecValidationError: An error occurred while validating the data spec\n\n    "

    def __init__(self, time_series_data_specs=None, files_data_spec=None):
        self.time_series_data_specs = time_series_data_specs
        self.files_data_spec = files_data_spec
        self.__validate_time_series_data_specs()
        self.__validate_files_data_spec()

    def __validate_time_series_data_specs(self):
        if self.time_series_data_specs:
            labels = []
            for tsds in self.time_series_data_specs:
                if not isinstance(tsds, TimeSeriesDataSpec):
                    raise DataSpecValidationError("Time series data specs must be TimeSeriesDataSpec objects.")
                label = tsds.label or "default"
                if label in labels:
                    raise DataSpecValidationError("Unique labels for each dataspec must be used")
                labels.append(label)
                if not isinstance(tsds.time_series, list):
                    raise DataSpecValidationError("Time series must be a list of time series objects.")
                if not (len(tsds.time_series) > 0):
                    raise DataSpecValidationError("A time series data spec does not contain any time series")
                ts_labels = []
                for ts in tsds.time_series:
                    if not isinstance(ts, TimeSeries):
                        raise DataSpecValidationError("Time series must be a TimeSeries object")
                    if ts.label and (ts.label in ts_labels):
                        raise DataSpecValidationError("Time series labels must be unique")
                    ts_labels.append(ts.label)

    def __validate_files_data_spec(self):
        if self.files_data_spec:
            if not isinstance(self.files_data_spec, FilesDataSpec):
                raise DataSpecValidationError("files_data_spec must be a FilesDataSpec object")
            if not isinstance(self.files_data_spec.file_ids, dict):
                raise DataSpecValidationError("file_ids must be a dict of form {name: id}")
            for (name, id) in self.files_data_spec.file_ids.items():
                if not isinstance(name, str):
                    raise DataSpecValidationError("File names must be a strings")
                if not isinstance(id, int):
                    raise DataSpecValidationError("File ids must be integers")

    def to_JSON(self):
        return json.dumps(
            {to_camel_case(key): value for (key, value) in self.__dict__.items()}, cls=self.__DataSpecEncoder
        )

    @classmethod
    def from_JSON(cls, json_repr):
        json_repr = json.loads(json_repr)
        time_series_data_specs_json = json_repr.get("timeSeriesDataSpecs")
        files_data_spec_json = json_repr.get("filesDataSpec")
        if not (time_series_data_specs_json or files_data_spec_json):
            raise DataSpecValidationError("Not a valid data spec")
        time_series_data_specs = []
        if time_series_data_specs_json:
            for tsds_json in time_series_data_specs_json:
                time_series_json = tsds_json.get("timeSeries")
                if time_series_json:
                    tsds = TimeSeriesDataSpec(**{to_snake_case(key): value for (key, value) in tsds_json.items()})
                    tsds.time_series = [
                        TimeSeries(**{to_snake_case(key): value for (key, value) in ts.items()})
                        for ts in time_series_json
                    ]
                    time_series_data_specs.append(tsds)
        files_data_spec = None
        if files_data_spec_json:
            files_data_spec = FilesDataSpec(
                **{to_snake_case(key): value for (key, value) in files_data_spec_json.items()}
            )
        return DataSpec(time_series_data_specs, files_data_spec)

    class __DataSpecEncoder(json.JSONEncoder):
        def default(self, obj):
            if isinstance(obj, (TimeSeries, TimeSeriesDataSpec, FilesDataSpec)):
                new_dict = {}
                for (key, value) in obj.__dict__.items():
                    if value is not None:
                        new_dict[to_camel_case(key)] = value
                return new_dict
            elif obj is None:
                del obj
            return super(self, self).default(obj)


class DataSpecValidationError(Exception):
    pass


class DataTransferService:
    "Create a Data Transfer Service object.\n\n    Fetch timeseries from the api.\n    "

    def __init__(self, data_spec, project=None, api_key=None, cookies=None, num_of_processes=10):
        "\n        Args:\n            data_spec (data_transfer_service.DataSpec):   Data Spec.\n            project (str):          Project name.\n            api_key (str):          Api key.\n            cookies (dict):         Cookies.\n        "
        (config_api_key, config_project) = config.get_config_variables(api_key, project)
        if isinstance(data_spec, DataSpec):
            self.data_spec = data_spec
        elif isinstance(data_spec, str):
            self.data_spec = DataSpec.from_JSON(data_spec)
        else:
            raise InputError("DataTransferService accepts a DataSpec instance or a json represantion of it.")
        self.ts_data_specs = self.data_spec.time_series_data_specs
        self.files_data_spec = self.data_spec.files_data_spec
        self.api_key = api_key or config_api_key
        self.project = project or config_project
        self.cookies = cookies
        self.num_of_processes = num_of_processes

    def get_dataframes(self, drop_agg_suffix=True):
        "Return a dictionary of dataframes indexed by label - one per data spec.\n\n        Args:\n            drop_agg_suffix (bool): If a time series has only one aggregate, drop the `|<agg-func>` suffix on\n                                    those column names.\n        Returns:\n            Dict[str, pd.DataFrame]: A label-indexed dictionary of data frames.\n        "
        if len(self.ts_data_specs) == 0:
            return {}
        if self.ts_data_specs is None:
            raise InputError("Data spec does not contain any TimeSeriesDataSpecs")
        dataframes = {}
        for tsds in self.ts_data_specs:
            df = self.get_dataframe(tsds.label, drop_agg_suffix=drop_agg_suffix)
            dataframes[tsds.label] = df
        return dataframes

    def get_dataframe(self, label="default", drop_agg_suffix=True):
        if self.ts_data_specs is None:
            raise InputError("Data spec does not contain any TimeSeriesDataSpecs")
        tsds = None
        for ts_data_spec in self.ts_data_specs:
            if ts_data_spec.label == label:
                tsds = ts_data_spec
        if tsds:
            ts_list = []
            ts_res = time_series_v06.get_multiple_time_series_by_id(
                ids=list(set([ts.id for ts in tsds.time_series])), api_key=self.api_key, project=self.project
            )
            id_to_name = {ts["id"]: ts["name"] for ts in ts_res.to_json()}
            for ts in tsds.time_series:
                ts._name = id_to_name[ts.id]
                if not ts.label:
                    ts.label = id_to_name[ts.id]
                ts_dict = dict(
                    name=id_to_name[ts.id],
                    aggregates=(ts.aggregates or tsds.aggregates),
                    missingDataStrategy=ts.missing_data_strategy,
                )
                ts_list.append(ts_dict)
            df = time_series_v05.get_datapoints_frame(
                ts_list,
                tsds.aggregates,
                tsds.granularity,
                tsds.start,
                tsds.end,
                api_key=self.api_key,
                project=self.project,
                cookies=self.cookies,
                processes=self.num_of_processes,
            )
            df = self.__apply_missing_data_strategies(df, ts_list, tsds.missing_data_strategy)
            df = self.__convert_ts_names_to_labels(df, tsds, drop_agg_suffix)
            return df
        return None

    def get_file(self, name):
        "Return files by name as specified in the DataSpec\n\n        Args:\n            name (str): Name of file\n        "
        if not self.files_data_spec:
            raise InputError("Data spec does not contain a FilesDataSpec")
        if isinstance(self.files_data_spec, FilesDataSpec):
            id = self.files_data_spec.file_ids.get(name)
            if id:
                file_bytes = files.download_file(id, get_contents=True, api_key=self.api_key, project=self.project)
                return BytesIO(file_bytes)
            raise InputError("Invalid name")

    def __convert_ts_names_to_labels(self, df, tsds, drop_agg_suffix):
        name_to_label = {}
        for ts in tsds.time_series:
            for agg in ts.aggregates or tsds.aggregates:
                agg_return_name = get_aggregate_func_return_name(agg)
                if (len((ts.aggregates or tsds.aggregates)) > 1) or (not drop_agg_suffix):
                    name_to_label[((ts._name + "|") + agg_return_name)] = (ts.label + "|") + agg_return_name
                else:
                    name_to_label[((ts._name + "|") + agg_return_name)] = ts.label
        return df.rename(columns=name_to_label)

    def __apply_missing_data_strategies(self, df, ts_list, global_missing_data_strategy):
        "Applies missing data strategies to dataframe.\n\n        Local strategies have precedence over global strategy.\n        "
        new_df = df["timestamp"]
        for ts in ts_list:
            name = ts["name"]
            colnames = [colname for colname in df.columns.values if colname.startswith(name)]
            missing_data_strategy = ts.get("missingDataStrategy", global_missing_data_strategy)
            partial_df = df[colnames]
            if missing_data_strategy == "ffill":
                partial_df = df[colnames].fillna(method="pad")
            elif missing_data_strategy and missing_data_strategy.endswith("Interpolation"):
                method = missing_data_strategy[:(-13)].lower()
                partial_df = df[colnames].interpolate(method=method, axis=0)
            new_df = pd.concat([new_df, partial_df], axis=1)
            new_df = new_df.iloc[:, (~new_df.columns.duplicated())]
        return new_df
