from datetime import datetime
import logging
from typing import Dict, List, Optional

import simpleflake
import ujson

from aporia.api.input_validation import (
    is_valid_ids,
    is_valid_predict_param_features,
    is_valid_predict_param_list,
)
from aporia.api.types import FeatureInput, FeatureValue, PredictionIdentifier
from aporia.api.utils import (
    convert_list_members_to_lists,
    convert_to_list,
    convert_to_list_of_dicts,
)
from aporia.consts import LOGGER_NAME
from aporia.graphql_client import GraphQLClient
from aporia.model_context import ModelContext

logger = logging.getLogger(LOGGER_NAME)


async def log_predict_batch(
    graphql_client: GraphQLClient,
    model_id: str,
    model_version: str,
    environment: str,
    x: List[FeatureInput],
    y: List[List[float]],
    aporia_ids: List[int],
    confidence: Optional[List[List[float]]] = None,
    occurred_at: Optional[datetime] = None,
    extra_inputs: Optional[List[Dict[str, FeatureValue]]] = None,
    extra_outputs: Optional[List[Dict[str, FeatureValue]]] = None,
    user_ids: Optional[List[PredictionIdentifier]] = None,
):
    """Reports a batch of predictions.

    Args:
        graphql_client (GraphQLClient): GraphQL client
        model_id (str): Model ID
        model_version (str): Model version
        environment (str): Environment in which aporia is running.
        x (List[FeatureInput]): Feature values
        y (List[List[float]]): Prediction result
        aporia_ids (List[int]): List of unique identifiers generated by Aporia for each prediction
        confidence (List[List[float]], optional): Prediction confidence. Defaults to None.
        occurred_at (datetime, optional): Prediction timestamp. Defaults to None.
        extra_inputs (List[Dict[str, FeatureValue]], optional): Extra inputs. Defaults to None.
        extra_outputs (List[Dict[str, FeatureValue]], optional): Extra outputs. Defaults to None.
        user_ids (List[PredictionIdentifier], optional): List of unique identifiers provided by the user for
            each prediction. Defaults to None.
    """
    query = """
        mutation LogPredict(
            $modelId: String!,
            $modelVersion: String!,
            $x: JSONString!,
            $yPred: [[Float]]!,
            $confidence: [[Float]],
            $occurredAt: String,
            $environment: String!,
            $aporiaIds: [BigInt]!,
            $extraInputs: [[ExtraInputValue]],
            $extraOutputs: [[ExtraOutputValue]],
            $userIds: [PredictionIdentifier]
        ) {
            logPredictions(
                modelId: $modelId,
                modelVersion: $modelVersion,
                x: $x,
                yPred: $yPred,
                confidence: $confidence,
                occurredAt: $occurredAt,
                environment: $environment,
                aporiaIds: $aporiaIds,
                extraInputs: $extraInputs,
                extraOutputs: $extraOutputs,
                userIds: $userIds
            ) {
                warnings
            }
        }
    """

    variables = {
        "modelId": model_id,
        "modelVersion": model_version,
        "x": ujson.dumps(x),
        "yPred": y,
        "confidence": confidence,
        "occurredAt": None if occurred_at is None else occurred_at.isoformat(),
        "environment": environment,
        "aporiaIds": aporia_ids,
        "extraInputs": _build_extra_io_values(extra_inputs),
        "extraOutputs": _build_extra_io_values(extra_outputs),
        "userIds": user_ids,
    }

    result = await graphql_client.query_with_retries(query, variables)
    for warning in result["logPredictions"]["warnings"]:
        logger.warning(warning)


def _prepare_predict_batch_input(
    context: ModelContext,
    x: List[FeatureInput],
    y: List[List[float]],
    confidence: Optional[List[List[float]]] = None,
    extra_inputs: Optional[List[Dict[str, FeatureValue]]] = None,
    extra_outputs: Optional[List[Dict[str, FeatureValue]]] = None,
) -> tuple:
    x = convert_to_list_of_dicts(x)
    y = convert_list_members_to_lists(convert_to_list(y))
    if confidence is not None:
        confidence = convert_list_members_to_lists(convert_to_list(confidence))

    extra_inputs = _prepare_batch_input_extras(extra_inputs, context.extra_inputs, len(y))
    extra_outputs = _prepare_batch_input_extras(extra_outputs, context.extra_outputs, len(y))

    return x, y, confidence, extra_inputs, extra_outputs


def _prepare_batch_input_extras(
    extras: Optional[List[Dict[str, FeatureValue]]],
    context_extras: Optional[Dict[str, FeatureValue]],
    expected_length: int,
) -> Optional[List[Dict[str, FeatureValue]]]:
    if context_extras is not None:
        if extras is not None:
            # We want the extra_input/output parameters to override the values in the
            # context, so they must be second in the dict merge.
            return [{**context_extras, **extra} for extra in extras]
        else:
            return [context_extras for _ in range(expected_length)]

    return extras


def _build_extra_io_values(
    data: Optional[List[Dict[str, FeatureValue]]] = None
) -> Optional[List[List[dict]]]:
    if data is None:
        return None

    result = []
    for data_point in data:
        result.append([{"name": name, "value": value} for name, value in data_point.items()])

    return result


def is_valid_predict_input(
    x: List[FeatureInput],
    y: List[List[float]],
    confidence: Optional[List[List[float]]] = None,
    extra_inputs: Optional[List[Dict[str, FeatureValue]]] = None,
    extra_outputs: Optional[List[Dict[str, FeatureValue]]] = None,
    ids: Optional[List[PredictionIdentifier]] = None,
) -> bool:
    """Checks if log_predict_batch input is valid.

    Args:
        x (List[FeatureInput]): Feature values
        y (List[List[float]]): Prediction result
        confidence (List[List[float]], optional): Prediction confidence. Defaults to None.
        extra_inputs (List[Dict[str, FeatureValue]], optional): Extra inputs. Defaults to None.
        extra_outputs (List[Dict[str, FeatureValue]], optional): Extra outputs. Defaults to None.
        ids (List[PredictionIdentifier], optional): An unique identifier string of the
            prediction. Defaults to None.

    Returns:
        bool: True if all of the parameters are valid. False otherwise
    """
    if not is_valid_predict_param_features(x):
        logger.warning(
            "Invalid input format for x parameter - expected a non-empty list of dictionaries"
        )
        return False

    if not is_valid_predict_param_list(y):
        logger.warning(
            "Invalid input format for y parameter - expected a non-empty 2 dimensional list"
        )
        return False

    if len(x) != len(y):
        logger.warning("Invalid input: x and y should have identical length")
        return False

    if confidence is not None:
        if not is_valid_predict_param_list(confidence):
            logger.warning(
                "Invalid input format for confidence parameter - expected a "
                "non-empty 2 dimensional list"
            )
            return False

        if len(y) != len(confidence):
            logger.warning("Invalid input: y and confidence should have identical length")
            return False

    if extra_inputs is not None:
        if not _is_valid_extra_io(data=extra_inputs, expected_length=len(x)):
            logger.warning(
                "Invalid input: extra_inputs must be a list of values equal in length to the "
                "number of predictions, such that each element in the list is a dict "
                "containing the extra inputs for a single prediction"
            )
            return False

    if extra_outputs is not None:
        if not _is_valid_extra_io(data=extra_outputs, expected_length=len(x)):
            logger.warning(
                "Invalid input: extra_outputs must be a list of values equal in length to the "
                "number of predictions, such that each element in the list is a dict "
                "containing the extra outputs for a single prediction"
            )
            return False

    if ids is not None:
        if not is_valid_ids(data=ids, expected_length=len(x)):
            logger.warning(
                "Invalid input: ids must be a list of strings or ints equal "
                "in length to the number of predictions"
            )
            return False

    return True


def _is_valid_extra_io(data: List[Dict[str, FeatureValue]], expected_length: int) -> bool:
    if not isinstance(data, list):
        return False

    if len(data) != expected_length:
        return False

    if not all(isinstance(data_point, dict) for data_point in data):
        return False

    return True


def generate_aporia_ids(size: int) -> List[int]:
    """Generates a list of unique identifiers.

    Args:
        size (int): Size of the generated list

    Returns:
        List[int]: A list of unique identifiers in the specified size
    """
    ids = set()  # type: ignore
    # We want to make sure there are no duplications
    while len(ids) < size:
        ids.add(simpleflake.simpleflake())
    return list(ids)
