#!/usr/bin/python
#
# Copyright 2018-2021 Polyaxon, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json

from collections import defaultdict, namedtuple
from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple

from polyaxon.exceptions import PQLException
from traceml.processors.units_processors import to_cpu_value, to_memory_bytes


class QueryOpSpec(namedtuple("QueryOpSpec", "op negation params")):
    def items(self):
        return self._asdict().items()


def parse_operation_value(value: Any):
    try:
        return json.loads(value)
    except ValueError:
        return value


def parse_negation_operation(operation: str) -> Tuple[bool, str]:
    """Parse the negation modifier in an operation."""
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    negation = False
    if _operation[0] == "~":
        negation = True
        _operation = _operation[1:]

    return negation, _operation.strip()


def parse_nil_operation(operation: str) -> Optional["QueryOpSpec"]:
    """Parse the nil operator in an operation."""
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))

    # Check negation
    negation, _operation = parse_negation_operation(_operation)

    # Check nil
    if _operation == "nil":
        return QueryOpSpec("nil", negation, None)

    return None


def parse_comparison_operation(operation: str) -> Tuple[Optional[str], str]:
    """Parse the comparison operator in an operation."""
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    # Check inclusion comparison
    if _operation[:2] in ("<=", "=<"):
        return "<=", parse_operation_value(_operation[2:].strip())

    if _operation[:2] in (">=", "=>"):
        return ">=", parse_operation_value(_operation[2:].strip())

    # Non inclusive
    if _operation[:1] in (">", "<"):
        return _operation[:1], parse_operation_value(_operation[1:].strip())

    return None, parse_operation_value(_operation)


def parse_datetime_operation(operation: str) -> "QueryOpSpec":
    """Parse datetime operations.

    A datetime operation can be one of the following:

     * single value: start_date:2014-10-10, start_date:>2014-10-10, start_date:>=2014-10-10
     * negation single value: start_date:~2014-10-10
     * interval: start_date:2010-10-10 10:10 .. 2012-10-10
     * negation interval: start_date:~2010-10-10 10:10 .. 2012-10-10

    This parser does not allow `|`
    """
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    # Check not allowed ops
    if "|" in _operation:
        raise PQLException(
            "`|` is not allowed for datetime conditions. "
            "Operation: {}".format(operation)
        )

    # Check nil operation
    spec = parse_nil_operation(_operation)
    if spec:
        return spec

    # Check negation
    negation, _operation = parse_negation_operation(_operation)

    # Check range operator
    if ".." in _operation:
        op = ".."
        params = _operation.split("..")
        params = [parse_operation_value(param.strip()) for param in params if param]
        if len(params) != 2:
            raise PQLException(
                "Expression is not valid, ranges requires only 2 params, "
                "Operation: {}".format(operation)
            )
        return QueryOpSpec(op, negation, params)

    # Check comparison operators
    op, _operation = parse_comparison_operation(_operation)
    if not op:
        # Now the operation must be an equality param
        op = "="

    if not _operation:
        raise PQLException(
            "Expression is not valid, it must be formatted as "
            "name:operation, "
            "Operation: {}".format(operation)
        )
    return QueryOpSpec(op, negation, _operation)


def parse_scalar_operation(operation: str, loader: Callable = None) -> "QueryOpSpec":
    """Parse scalar operations.

    A scalar operation can be one of the following:

     * single value: start_date:12, metric1:>0.9, metric1:>=-0.12
     * negation single value: metric1:~1112, metric1:~<1112 equivalent to metric1:>=1112

    This parser does not allow `|` and `..`.
    """
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    # Check not allowed ops
    if "|" in _operation:
        raise PQLException(
            "`|` is not allowed for scalar conditions. "
            "Operation: {}".format(operation)
        )
    if ".." in _operation:
        raise PQLException(
            "`..` is not allowed for scalar conditions. "
            "Operation: {}".format(operation)
        )

    # Check nil operation
    spec = parse_nil_operation(_operation)
    if spec:
        return spec

    # Check negation
    negation, _operation = parse_negation_operation(_operation)

    # Check comparison operators
    op, _operation = parse_comparison_operation(_operation)
    if not op:
        # Now the operation must be an equality param
        op = "="

    if loader:
        _operation = loader(_operation)
    # Check that params are scalar (int, float)
    try:
        int(_operation)
    except (ValueError, TypeError):
        try:
            float(_operation)
        except (ValueError, TypeError):
            raise PQLException(
                "Scalar operation requires int or float params, "
                "received {}.".format(operation)
            )
    return QueryOpSpec(op, negation, _operation)


def parse_memory_operation(operation: str) -> "QueryOpSpec":
    return parse_scalar_operation(operation, loader=to_memory_bytes)


def parse_cpu_operation(operation: str) -> "QueryOpSpec":
    return parse_scalar_operation(operation, loader=to_cpu_value)


def parse_value_operation(operation: str) -> "QueryOpSpec":
    """Parse value operations.

    A value operation can be one of the following:

     * single value: tag1:foo
     * negation single value: tag1:~foo
     * multiple values: tag1:foo|bar|moo
     * negation multiple values: tag1:~foo|bar|moo

    This parser does not allow `..`.
    """
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    # Check range not allowed
    if ".." in _operation:
        raise PQLException(
            "`..` is not allowed for value conditions. "
            "Operation: {}".format(operation)
        )

    # Check nil operation
    spec = parse_nil_operation(_operation)
    if spec:
        return spec

    # Check negation
    negation, _operation = parse_negation_operation(_operation)

    # Early return
    if _operation is not None and not isinstance(_operation, str):
        return QueryOpSpec("=", negation, _operation)

    # Check in operator
    if "|" in _operation:
        op = "|"
        params = _operation.split("|")
        params = [
            parse_operation_value(param.strip()) for param in params if param.strip()
        ]
        if len(params) <= 1:
            raise PQLException(
                "`{}` is not allowed for value conditions, "
                "Operation: {}".format(op, operation)
            )
        return QueryOpSpec(op, negation, params)

    # Check comparison operators
    op, _operation = parse_comparison_operation(_operation)
    if not op:
        # Now the operation must be an equality param
        op = "="

    if _operation is None:
        raise PQLException(
            "Expression is not valid, it must be formatted as "
            "name:operation, "
            "Operation: {}".format(operation)
        )
    return QueryOpSpec(op, negation, _operation)


def parse_search_operation(operation: str) -> "QueryOpSpec":
    """Parse search operations.

    A search operation can be one of the following:

     * single value: tag1:foo
     * negation single value: tag1:~foo
     * multiple values: tag1:foo|bar|moo
     * negation multiple values: tag1:~foo|bar|moo
     * like value: name:%foo or name:foo% or name:%foo%
     * not like value: name:~%foo or name:~foo% or name:~%foo%

    This parser does not allow `..`, '>', '<', '>=', and '<='.
    """
    _operation = operation.strip()
    if not _operation:
        raise PQLException("Operation is not valid: {}".format(operation))
    # Check range not allowed
    if ".." in _operation:
        raise PQLException(
            "`..` is not allowed for search conditions. "
            "Operation: {}".format(operation)
        )

    # Check nil operation
    spec = parse_nil_operation(_operation)
    if spec:
        return spec

    # Check negation
    negation, _operation = parse_negation_operation(_operation)

    # Check comparison not allowed
    op, _operation = parse_comparison_operation(_operation)
    if op:
        raise PQLException(
            "`{}` is not allowed for search conditions, "
            "Operation: {}".format(op, operation)
        )

    # Early return
    if _operation is not None and not isinstance(_operation, str):
        return QueryOpSpec("=", negation, _operation)

    # Check in operator
    if "|" in _operation:
        op = "|"
        params = _operation.split("|")
        params = [
            parse_operation_value(param.strip()) for param in params if param.strip()
        ]
        if len(params) <= 1:
            raise PQLException(
                "`{}` is not allowed for search conditions, "
                "Operation: {}".format(op, operation)
            )
        return QueryOpSpec(op, negation, params)

    # Check like operator
    start_like = _operation.endswith("%")
    end_like = _operation.startswith("%")
    if start_like and end_like:
        op = "%%"
        params = _operation.split("%")
        params = [param.strip() for param in params if param.strip()]
        if len(params) != 1:
            raise PQLException(
                "`{}` is not allowed for search conditions, "