from dataclasses import dataclass
import datetime
import json
import os
import sys
import tempfile
from typing import List
import uuid

import click

from askanna.core import client, exceptions
from askanna.core.dataclasses import Metric, MetricDataPair, MetricLabel
from askanna.core.utils import create_suuid, json_serializer

__all__ = ["track_metric", "track_metrics", "MetricCollector", "MetricGateway"]

"""
How to use metrics module:

from askanna.core.metrics import *

mc = MetricCollector()
mc.add(metric=Metric(metric=MetricDataPair(name="foo", value=1, dtype="int"), label=[]))

# get json:

mc.to_json()
"""


class MetricsQuerySet:
    metrics = []
    queries = []
    parent = None

    def __init__(self, metrics: list = [], queries: list = [], parent=None):
        self.metrics = metrics
        self.queries = queries
        self.parent = parent

    def __len__(self):
        return len(self.metrics)

    def __iter__(self):
        yield from self.metrics

    def get(self, name, **kwargs):
        ret = self.filter(name, **kwargs)
        if len(ret) == 1:
            return ret[0]
        elif len(ret) > 1:
            click.echo(
                f"Found more metrics matching name '{name}',"
                f" use .filter(name='{name}') to get all matches.",
                err=True,
            )
            return ret[0]
        else:
            click.echo(f"A metric with name '{name}' is not found")

    def filter(self, name, **kwargs):
        ret = list(
            filter(lambda x: x.get("metric", {}).get("name") == name, self.metrics)
        )
        return ret

    def to_json(self) -> str:
        return json.dumps(self.metrics, indent=1)

    def to_raw(self) -> list:
        return self.metrics


@dataclass
class MetricCollector:
    metrics = []
    run_suuid: str = None
    local_suuid: str = None
    metrics_folder: str = None
    metrics_file: str = None
    changed: bool = False

    def __post_init__(self):
        if not self.run_suuid:
            # generate local suuid
            local_uuid = uuid.uuid4()
            self.local_suuid = create_suuid(local_uuid)

        self.metrics_folder = self.get_metrics_filepath()
        self.metrics_file = os.path.join(self.metrics_folder, "metrics.json")

        if self.run_suuid:
            self.restore_session()

    def get_suuid(self):
        """
        Return the run_suuid if exist, otherwise the local_suuid
        """
        return self.run_suuid or self.local_suuid

    def get_metrics_filepath(self):
        tmpdir = tempfile.gettempdir()
        return os.path.join(tmpdir, "askanna", "run", self.get_suuid())

    def metric_to_type(self, metric: dict = None) -> MetricDataPair:
        return MetricDataPair(
            name=metric["name"], value=metric["value"], dtype=metric["type"]
        )

    def label_to_type(self, label_list: list = None) -> list:
        _ret = []
        for label in label_list:
            _ret.append(
                MetricLabel(
                    name=label["name"], value=label["value"], dtype=label["type"]
                )
            )
        return _ret

    def restore_session(self):
        """
        We can only restore a session when the self.run_uuid is set
        """
        if not self.run_suuid:
            return

        restored_metrics = []
        try:
            with open(self.metrics_file, "r") as f:
                stored_metrics = json.load(f)
                for stored_metric in stored_metrics:
                    restored_metrics.append(
                        Metric(
                            metric=self.metric_to_type(stored_metric.get("metric")),
                            label=self.label_to_type(stored_metric.get("label")),
                            run_suuid=stored_metric.get("run_suuid"),
                            created=datetime.datetime.fromisoformat(
                                stored_metric.get("created")
                            ),
                        )
                    )

                self.metrics = restored_metrics
                del restored_metrics
        except FileNotFoundError:
            pass

    def save_session(self):
        # create folder for this session, name is the run_suuid with prefix askanna-metrics

        os.makedirs(self.metrics_folder, exist_ok=True)
        with open(self.metrics_file, "w") as f:
            f.write(self.to_json())

    def has_metrics(self) -> bool:
        return len(self.metrics) > 0

    def add(self, metric: Metric) -> None:
        metric.run_suuid = self.get_suuid()
        self.metrics.append(metric)
        self.save_session()
        self.changed = True

    def to_dict(self, yes=True) -> List:
        return [m.to_dict() for m in self.metrics]

    def to_json(self, yes=True) -> str:
        """
        We serialize all metrics to json for reporting or sending to AskAnna
        """
        return json.dumps(self.to_dict(), default=json_serializer)

    def save(self, run_suuid: str = None, force: bool = False) -> None:
        """
        Connect to askanna backend to save the metrics
        We will only attempt to submit if there is data recorded
        """
        if not self.changed and not force:
            # We will not save when there is nothing changed (e.g. no new metric added)
            return
        if len(self.to_dict()) and not run_suuid and not self.run_suuid:
            click.echo(
                "The run SUUID is not set for this session.\n"
                + "AskAnna cannot submit the metrics to the platform.",
                err=True,
            )
            click.echo(
                f"Your metrics are saved locally in:\n{self.metrics_file}", err=True
            )
            sys.exit(1)
        self.save_session()
        mgw = MetricGateway()
        # the run_suuid takes higher prio then the one set in init
        short_uuid = run_suuid or self.run_suuid
        mgw.change(short_uuid=short_uuid, metrics=self.to_dict())


# Start logic for metric collection
mc = MetricCollector(run_suuid=os.getenv("JOBRUN_SUUID"))


def translate_dtype(typename) -> str:
    """
    Return the full name of the type, if not listed, return the typename from the input
    """
    return {
        "bool": "boolean",
        "str": "string",
        "int": "integer",
        "float": "float",
        "dict": "dictionary",
    }.get(typename, typename)


def validate_value(value) -> bool:
    """
    Validate whether the value set is supported
    """
    return type(value).__name__ in [
        "bool",
        "str",
        "int",
        "float",
        "time",
        "date",
        "datetime",
        "tag",  # single string marked as tag
        "dict",
    ]


def _labels_to_type(label: dict = None) -> List:
    # process labels
    labels = []

    # test label type, if it is a string, then convert it to tag
    if isinstance(label, str):
        labels.append(MetricLabel(name=label, value=None, dtype="tag"))
        label = None

    if label:
        for k, v in label.items():
            if v is None:
                labels.append(MetricLabel(name=k, value=None, dtype="tag"))
            else:
                labels.append(
                    MetricLabel(
                        name=k, value=v, dtype=translate_dtype(type(v).__name__)
                    )
                )
    return labels


def track_metric(name: str, value, label: dict = None) -> None:
    # store the metric
    if not validate_value(value):
        click.echo(
            f"AskAnna cannot store this datatype. Metric not stored for {name}, {value}, {label}."
        )
        return
    # add value to track queue
    datapair = MetricDataPair(
        name=name, value=value, dtype=translate_dtype(type(value).__name__)
    )
    labels = _labels_to_type(label)

    metric = Metric(metric=datapair, label=labels)
    mc.add(metric)


def track_metrics(metrics: dict, label: dict = None) -> None:
    """
    Transform many metrics to internal format
    """
    for name, value in metrics.items():
        track_metric(name, value, label)


class MetricGateway:
    def __init__(self, *args, **kwargs):
        self.client = client

    def get(self, run: str = None, job: str = None) -> MetricsQuerySet:
        """
        Depending on what we get as argument, we call the 'list' function
        with specific query params for a different endpoint
        """

        if not any([run, job]):
            # we try to get local metrics if this exists

            if mc.run_suuid and mc.has_metrics():
                return MetricsQuerySet(metrics=mc.to_dict())
            elif mc.run_suuid:
                # RUN_SUUID is set but we don't have metrics
                return MetricsQuerySet(metrics=[])
            else:
                click.echo(
                    "We cannot find a run metric file for the active run.", err=True
                )
        else:
            pre_query_params = {
                "run_suuid": run,
                "job_suuid": job,
            }
            # filter out the none's:
            query_params = dict(
                filter(lambda x: x[1] is not None, pre_query_params.items())
            )
            return self.list(query_params=query_params)

    def list(self, query_params: dict = None) -> MetricsQuerySet:
        """
        List endpoint order:
        - job
        - run
        """

        endpoints = {
            "job": lambda x: "{}job/{}/metrics/".format(self.client.config.remote, x),
            "run": lambda x: "{}runinfo/{}/metrics/".format(
                self.client.config.remote, x
            ),
        }
        if "run_suuid" in query_params.keys():
            url = endpoints.get("run")(query_params.get("run_suuid"))
        elif "job_suuid" in query_params.keys():
            url = endpoints.get("job")(query_params.get("job_suuid"))

        r = self.client.get(url, params=query_params)
        return MetricsQuerySet(metrics=r.json())

    def change(self, short_uuid, metrics):
        url = "{}{}/{}/{}/{}/".format(
            self.client.config.remote, "runinfo", short_uuid, "metrics", short_uuid
        )

        r = self.client.put(url, json={"metrics": metrics})

        if r.status_code != 200:
            raise exceptions.GetError(
                "{} - Something went wrong while updating metrics "
                ": {}".format(r.status_code, r.json())
            )


# code for exiting in ipython from https://stackoverflow.com/a/40222538
# exit_register runs at the end of ipython %run or the end of the python interpreter
try:

    def exit_register(fun, *args, **kwargs):
        """Decorator that registers at post_execute. After its execution it
        unregisters itself for subsequent runs."""

        def callback():
            fun()
            ip.events.unregister("post_execute", callback)

        ip.events.register("post_execute", callback)

    ip = get_ipython()
except NameError:
    from atexit import register as exit_register


@exit_register
def handle_exit_save_metrics():
    try:
        if mc.has_metrics():
            mc.save()
    except Exception as e:
        # do a hard exit
        click.echo(f"Exit error in AskAnna SDK. {e}", err=True)
        os._exit(1)
