import time
import json

from typing import Tuple, Sequence, Any, Iterable, Dict, MutableSequence, MutableMapping, Union, overload
from coba.backports import Literal

from coba.random import random
from coba.pipes import Pipes, Source, HttpSource, Drop, ArffReader, ListSource, Structure
from coba.contexts import CobaContext, CobaContext
from coba.exceptions import CobaException

from coba.environments.simulated.supervised import SupervisedSimulation
from coba.pipes.readers import SparseWithMeta

class OpenmlSource(Source[Iterable[Tuple[Union[MutableSequence, MutableMapping],Any]]]):
    """Load a source (local if cached) from openml.
    
    This is primarily used by OpenmlSimulation to create Environments for Experiments.
    """

    @overload
    def __init__(self, *, data_id:int, label_type:Literal["C","R"]=None, cat_as_str:bool=False):
        """Instantiate an OpenmlSource.
        
        Args:
            data_id: The data id uniquely identifying the dataset on openml (i.e., openml.org/d/{id})
            label_type: Indicates if a regression or classification label should be used on the dataset.
            cat_as_str: Indicates if categorical features should be encoded as a string rather than one hot encoded. 
        """
        ...

    @overload
    def __init__(self, *, task_id:int, cat_as_str:bool=False):
        """Instantiate an OpenmlSource.

        Args:
            task_id: The openml task id which identifies the dataset to use from openml along with its label
            cat_as_str: Indicates if categorical features should be encoded as a string rather than one hot encoded. 
        """
        ...

    def __init__(self, **kwargs):
        """Instantiate an OpenmlSource."""

        self._data_id    = kwargs.get('data_id',None)
        self._task_id    = kwargs.get('task_id',None)
        self._label_type = kwargs.get('label_type',None)
        self._cat_as_str = kwargs.get('cat_as_str',False)
        self._cache_keys = {}

    @property
    def params(self) -> Dict[str,Any]:
        """Parameters describing the openml source."""
        if not self._task_id:
            return  { "openml": self._data_id, "cat_as_str": self._cat_as_str }
        else:
            return  { "openml": f"T{self._task_id}", "cat_as_str": self._cat_as_str }

    def read(self) -> Iterable[Tuple[Any, Any]]:
        """Read and parse the openml source."""
        try:

            data_id    = self._data_id
            label_type = self._label_type
            task_id    = self._task_id
            target     = None

            if task_id is not None:
                self._cache_keys['task'] = f"openml_{task_id:0>6}_task"
                task   = self._get_task_description(task_id)
                source = ([i for i in task["input"] if i.get('name',None) == 'source_data'] + [None])[0]
                
                if source:
                    data_id = source["data_set"]["data_set_id"]
                    target  = source["data_set"]["target_feature"]
                else:
                    raise CobaException(f"Openml task {task_id} does not appear to have an associated data source.")

            self._cache_keys['descr'] = f"openml_{data_id:0>6}_descr"
            self._cache_keys['feats'] = f"openml_{data_id:0>6}_feats"
            self._cache_keys['csv'  ] = f"openml_{data_id:0>6}_csv"
            self._cache_keys['arff' ] = f"openml_{data_id:0>6}_arff"
            self._cache_keys['tasks'] = f"openml_{data_id:0>6}_tasks"

            dataset_description  = self._get_dataset_description(data_id)

            if dataset_description['status'] == 'deactivated':
                raise CobaException(f"Openml {data_id} has been deactivated. This is often due to flags on the data.")

            is_ignore = lambda r: (
                r['is_ignore'        ] == 'true' or
                r['is_row_identifier'] == 'true' or
                r['data_type'        ] not in ['numeric', 'nominal']
            )

            ignore = [ self._clean_name(f['name']) for f in self._get_feature_descriptions(data_id) if is_ignore(f)]

            if not target and label_type is not None:
                target = self._clean_name(self._get_target_for_problem_type(data_id, label_type))
            
            if not target and label_type is None:
                target = self._clean_name(dataset_description["default_target_attribute"])

            if target in ignore: 
                ignore.pop(ignore.index(target))

            def row_has_missing_values(row):
                row_values = row._values.values() if isinstance(row,SparseWithMeta) else row._values
                return "?" in row_values or "" in row_values

            source    = ListSource(self._get_dataset_lines(dataset_description["file_id"], None))
            reader    = ArffReader(cat_as_str=self._cat_as_str)
            drop      = Drop(drop_cols=ignore, drop_row=row_has_missing_values)
            structure = Structure([None, target])

            return Pipes.join(source, reader, drop, structure).read()

        except KeyboardInterrupt:
            #we don't want to clear the cache in the case of a KeyboardInterrupt
            raise

        except CobaException:
            #we don't want to clear the cache if it is an error we know about (the original raise should clear if needed)
            raise

        except Exception:
            #if something unexpected went wrong clear the cache just in case it was corrupted somehow
            self._clear_cache()
            raise

    def _clean_name(self, name: str) -> str:
        return name.strip().strip('\'"').replace('\\','')

    def _get_data(self, url:str, key:str, checksum:str=None) -> Iterable[str]:

        # This can't be done in a streaming manner unless there is a persistent cacher.
        # Because we don't require cacher this means this may not be possible with streaming. 
        # if checksum is not None and md5(bites).hexdigest() != checksum:
        #     #if the cache has become corrupted we need to clear it
        #     CobaContext.cacher.rmv(key)
        #     message = (
        #         f"The response from {url} did not match the given checksum {checksum}. This could be the result "
        #         "of network errors or the file becoming corrupted. Please consider downloading the file again. "
        #         "If the error persists you may want to manually download and reference the file.")
        #     raise CobaException(message) from None
        try:
            for b in CobaContext.cacher.get_put(key, lambda: self._http_request(url)):
                yield b.decode('utf-8')
        except Exception:
            self._clear_cache()
            raise

    def _http_request(self, url:str) -> Iterable[bytes]:
        api_key = CobaContext.api_keys['openml']
        srcsema = CobaContext.store.get("srcsema")

        # we only allow three paralellel request, another attempt at being more "considerate".
        if srcsema:srcsema.acquire() 

        # An attempt to be considerate of how often we hit their REST api. 
        # They don't publish any rate-limiting guidelines so this is just a guess.
        
        if srcsema: time.sleep(2*random())

        try:
            with HttpSource(url + (f'?api_key={api_key}' if api_key else '')).read() as response:

                if response.status_code == 412:
                    if 'please provide api key' in response.text:
                        message = (
                            "Openml has requested an API Key to access openml's rest API. A key can be obtained by creating "
                            "an openml account at openml.org. Once a key has been obtained it should be placed within "
                            "~/.coba as { \"api_keys\" : { \"openml\" : \"<your key here>\", } }.")
                        raise CobaException(message)

                    if 'authentication failed' in response.text:
                        message = (
                            "The API Key you provided no longer seems to be valid. You may need to create a new one by "
                            "logging into your openml account and regenerating a key. After regenerating the new key "
                            "should be placed in ~/.coba as { \"api_keys\" : { \"openml\" : \"<your key here>\", } }.")
                        raise CobaException(message)

                if response.status_code == 404:
                    raise CobaException("We're sorry but we were unable to find the requested dataset on openml.")

                if response.status_code != 200:
                    raise CobaException(f"An error was returned by openml: {response.text}")

                # NOTE: These two checks need to be gated with a status code failure 
                # NOTE: otherwise this will cause the data to be downloaded all at once
                # NOTE: unfortunately I don't know the appropriate status code so commenting out for now
                # if "Usually due to high server load" in response.text:
                #     message = (
                #         "Openml has experienced an error that they believe is the result of high server loads. "
                #         "Openml recommends that you try again in a few seconds. Additionally, if not already "
                #         "done, consider setting up a DiskCache in a coba config file to reduce the number of "
                #         "openml calls in the future.")
                #     raise CobaException(message) from None

                # if '' == response.text:
                #     raise CobaException("Openml experienced an unexpected error. Please try requesting the data again.") from None

                for b in response.iter_lines(decode_unicode=False):
                    yield b
        finally:
            if srcsema: srcsema.release()

    def _get_dataset_description(self, data_id:int) -> Dict[str,Any]:

        description_txt = " ".join(self._get_data(f'https://www.openml.org/api/v1/json/data/{data_id}', self._cache_keys['descr']))
        description_obj = json.loads(description_txt)["data_set_description"]

        return description_obj

    def _get_feature_descriptions(self, data_id:int) -> Sequence[Dict[str,Any]]:

        types_txt = " ".join(self._get_data(f'https://www.openml.org/api/v1/json/data/features/{data_id}', self._cache_keys['feats']))
        types_obj = json.loads(types_txt)["data_features"]["feature"]

        return types_obj

    def _get_task_descriptions(self, data_id) -> Sequence[Dict[str,Any]]:
        
        try:
            text  = " ".join(self._get_data(f'https://www.openml.org/api/v1/json/task/list/data_id/{data_id}', self._cache_keys['tasks']))
            return json.loads(text).get("tasks",{}).get("task",[])
        except CobaException:
            return []

    def _get_task_description(self, task_id) -> Dict[str,Any]:
        
        description_txt = " ".join(self._get_data(f'https://www.openml.org/api/v1/json/task/{task_id}', self._cache_keys['task']))
        description_obj = json.loads(description_txt)['task']
        
        return description_obj

    def _get_dataset_lines(self, file_id:str, md5_checksum:str) -> Iterable[str]:

            arff_url = f"https://www.openml.org/data/v1/download/{file_id}"

            csv_key  = self._cache_keys['csv']
            arff_key = self._cache_keys['arff']

            if csv_key in CobaContext.cacher: CobaContext.cacher.rmv(csv_key)
            return self._get_data(arff_url, arff_key, md5_checksum)

    def _get_target_for_problem_type(self, data_id:int, label_type:str):

        task_type_id = 1 if label_type == "C" else 2

        for task in self._get_task_descriptions(data_id):
            if task["task_type_id"] == task_type_id:
                for input in task['input']:
                    if input['name'] == 'target_feature':
                        return input['value'] # just take the first one

        raise CobaException(f"Openml {data_id} does not appear to be a {label_type} dataset")

    def _clear_cache(self) -> None:
        for key in self._cache_keys.values():
            CobaContext.cacher.release(key) #to make sure we don't get stuck in a race condition
            CobaContext.cacher.rmv(key)

class OpenmlSimulation(SupervisedSimulation):
    """A supervised simulation created from an openml dataset.

    Download a dataset from openml.org and create a SupervisedSimulation.
    """

    @overload
    def __init__(self, data_id: int, label_type: Literal["C","R"] = None, cat_as_str: bool = False, take: int = None):
        """Instantiate an OpenmlSimulation.
        
        Args:
            data_id: The data id uniquely identifying the dataset on openml (i.e., openml.org/d/{id})
            label_type: Indicates if a regression or classification label should be used on the dataset.
            cat_as_str: Indicates if categorical features should be encoded as a string rather than one hot encoded.
            take: The number of interactions we'd like the simulation to have (these will be selected at random).
        """
        ...

    @overload
    def __init__(self, *, task_id: int, cat_as_str: bool = False, take: int=None):
        """Instantiate an OpenmlSimulation.

        Args:
            task_id: The openml task id which identifies the dataset to use from openml along with its label
            cat_as_str: Indicates if categorical features should be encoded as a string rather than one hot encoded.
            take: The number of interactions we'd like the simulation to have (these will be selected at random).
        """
        ...

    def __init__(self, *args, **kwargs) -> None:
        """Instantiate an OpenmlSimulation."""
        kwargs.update(zip(['data_id','label_type','cat_as_str','take'], args))
        super().__init__(OpenmlSource(**kwargs), None, kwargs.get('label_type',None), kwargs.get('take',None))

    @property
    def params(self) -> Dict[str, Any]:
        return {**super().params, "type": "OpenmlSimulation" }

    def __str__(self) -> str:
        return f"Openml(id={self.params['openml']}, cat_as_str={self.params['cat_as_str']})"
