"""PISCES Marvel utilities."""
from elasticsearch import Elasticsearch
import hashlib

from base.logging import zoox_logger

ES_CLUSTER = "https://search-sim-metrics-v6-rhfhvb54pmawksutaybix43i6q.us-west-1.es.amazonaws.com"
LOG = zoox_logger.ZooxLogger(__name__)


def alias(from_field, to_field):
    def rename(dic):
        try:
            dic[to_field] = dic.pop(from_field)
        except KeyError as e:
            LOG.error("Field access error %s", e)
            LOG.debug(json.dumps(dic))
            dic[to_field] = None
        return dic

    return rename


def get_marvel_run(marvel_uuid, max_entries=10000):
    es = Elasticsearch([ES_CLUSTER])

    suite_info = get_suite_info(es, marvel_uuid)

    query = {
        "_source": {
            "include": [
                "pipedream_id",
                "combined_chum_uri",
                "description",
                "bazel_target",
                "branch",
                "git_revision",
                "scenario_params",
                "passed",
                "result",
                "run_uuid",
            ]
        },
        "size": max_entries,
        "query": {
            "constant_score": {
                "filter": {
                    "bool": {"must": [{"term": {"run_uuid": marvel_uuid}}]}
                }
            }
        },
        "sort": [{"bazel_target": {"order": "desc"}}],
    }

    results = es.search(index="metric_runs", body=query)

    # TODO: handle sim failures properly, but for now, throw them away
    source = (x["_source"] for x in results["hits"]["hits"])
    has_combined_chum_uri = list(
        filter(lambda x: "combined_chum_uri" in x, source)
    )

    for record in has_combined_chum_uri:
        m = hashlib.sha256()
        m.update(record["bazel_target"].encode("utf8"))
        record["sdl_target"] = record["bazel_target"]
        for p in sorted(record["scenario_params"]):
            m.update(p.encode("utf8"))
            record["sdl_target"] += " " + p

        record["scenario_hash"] = m.hexdigest()
        # directory / label to store data in
        record["label"] = suite_info.get("labels", [marvel_uuid])[0]

    f = alias("combined_chum_uri", "chum_uri")
    hits = list(map(f, has_combined_chum_uri))
    return hits


def get_suite_info(es, marvel_uuid, max_entries=10000):
    query = {
        "size": max_entries,
        "query": {
            "constant_score": {
                "filter": {
                    "bool": {"must": [{"term": {"run_uuid": marvel_uuid}}]}
                }
            }
        },
    }

    results = es.search(index="suite_runs", body=query)

    source = list(x["_source"] for x in results["hits"]["hits"])
    assert len(source) == 1, f"Expected 1 result, got: {source}"
    return source[0]