#  Copyright 2021 Collate
#  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.

"""
ColumnValuesToBeBetween validation implementation
"""
# pylint: disable=duplicate-code

import traceback
from datetime import datetime
from functools import singledispatch
from typing import Union

from pandas import DataFrame
from sqlalchemy import inspect
from sqlalchemy.sql.sqltypes import DATE, DATETIME, TIMESTAMP

from metadata.generated.schema.tests.basic import (
    TestCaseResult,
    TestCaseStatus,
    TestResultValue,
)
from metadata.generated.schema.tests.testCase import TestCase
from metadata.orm_profiler.metrics.registry import Metrics
from metadata.orm_profiler.profiler.runner import QueryRunner
from metadata.utils.column_base_model import fetch_column_obj
from metadata.utils.entity_link import get_decoded_column
from metadata.utils.logger import test_suite_logger
from metadata.utils.test_suite import get_test_case_param_value

logger = test_suite_logger()


def convert_timestamp(timestamp: str) -> Union[int, float]:
    """convert timestamp to int

    Args:
        timestamp (str):

    Retunrs:
        int
    """
    if len(timestamp) < 13:  # check for ms timestamp
        return int(timestamp)
    return float(timestamp) / 1000


def test_case_status_result(min_value_res, max_value_res, min_bound, max_bound):
    return (
        TestCaseStatus.Success
        if min_value_res >= min_bound and max_value_res <= max_bound
        else TestCaseStatus.Failed,
        f"Found min={min_value_res}, max={max_value_res} vs."
        + f" the expected min={min_bound}, max={max_bound}.",
    )


@singledispatch
def column_values_to_be_between(
    runner,
    test_case: TestCase,
    execution_date: Union[datetime, float],
):
    raise NotImplementedError


@column_values_to_be_between.register
def _(
    runner: QueryRunner,
    test_case: TestCase,
    execution_date: Union[datetime, float],
) -> TestCaseResult:
    """
    Validate Column Values metric
    :param test_case: ColumnValuesToBeBetween
    :param col_profile: should contain MIN & MAX metrics
    :param execution_date: Datetime when the tests ran
    :return: TestCaseResult with status and results
    """

    try:
        column_name = get_decoded_column(test_case.entityLink.__root__)
        col = next(
            (col for col in inspect(runner.table).c if col.name == column_name),
            None,
        )
        if col is None:
            raise ValueError(
                f"Cannot find the configured column {column_name} for test case {test_case.name}"
            )

        min_value_dict = dict(
            runner.dispatch_query_select_first(Metrics.MIN.value(col).fn())
        )
        min_value_res = min_value_dict.get(Metrics.MIN.name)
        max_value_dict = dict(
            runner.dispatch_query_select_first(Metrics.MAX.value(col).fn())
        )
        max_value_res = max_value_dict.get(Metrics.MAX.name)
        if max_value_res is None or min_value_res is None:
            raise ValueError(
                f"Query on column {column_name} for test case {test_case.name} returned None"
            )

    except Exception as exc:
        msg = (
            f"Error computing {test_case.name} for {runner.table.__tablename__}: {exc}"
        )
        logger.debug(traceback.format_exc())
        logger.warning(msg)
        return TestCaseResult(
            timestamp=execution_date,
            testCaseStatus=TestCaseStatus.Aborted,
            result=msg,
            testResultValue=[
                TestResultValue(name="min", value=None),
                TestResultValue(name="max", value=None),
            ],
        )

    if isinstance(col.type, (TIMESTAMP, DATE, DATETIME)):
        min_bound = get_test_case_param_value(
            test_case.parameterValues,  # type: ignore
            "minValue",
            datetime.fromtimestamp,
            default=datetime.min,
            pre_processor=convert_timestamp,
        )

        max_bound = get_test_case_param_value(
            test_case.parameterValues,  # type: ignore
            "maxValue",
            datetime.fromtimestamp,
            default=datetime.max,
            pre_processor=convert_timestamp,
        )
    else:
        min_bound = get_test_case_param_value(
            test_case.parameterValues,  # type: ignore
            "minValue",
            float,
            default=float("-inf"),
        )

        max_bound = get_test_case_param_value(
            test_case.parameterValues,  # type: ignore
            "maxValue",
            float,
            default=float("inf"),
        )

    status, result = test_case_status_result(
        min_value_res, max_value_res, min_bound, max_bound
    )

    return TestCaseResult(
        timestamp=execution_date,
        testCaseStatus=status,
        result=result,
        testResultValue=[
            TestResultValue(
                name="min",
                value=str(
                    min_value_res.timestamp()
                    if isinstance(min_value_res, datetime)
                    else min_value_res
                ),
            ),
            TestResultValue(
                name="max",
                value=str(
                    max_value_res.timestamp()
                    if isinstance(max_value_res, datetime)
                    else max_value_res
                ),
            ),
        ],
    )


@column_values_to_be_between.register
def _(
    runner: DataFrame,
    test_case: TestCase,
    execution_date: Union[datetime, float],
):
    """
    Validate Column Values metric
    :param test_case: ColumnValuesToBeBetween
    :param col_profile: should contain MIN & MAX metrics
    :param execution_date: Datetime when the tests ran
    :return: TestCaseResult with status and results
    """
    column_obj = fetch_column_obj(test_case.entityLink.__root__, runner)

    min_bound = get_test_case_param_value(
        test_case.parameterValues,  # type: ignore
        "minValue",
        float,
        default=float("-inf"),
    )

    max_bound = get_test_case_param_value(
        test_case.parameterValues,  # type: ignore
        "maxValue",
        float,
        default=float("inf"),
    )
    min_value_res = Metrics.MIN.value(column_obj).dl_fn(runner)

    max_value_res = Metrics.MAX.value(column_obj).dl_fn(runner)

    status, result = test_case_status_result(
        min_value_res, max_value_res, min_bound, max_bound
    )

    return TestCaseResult(
        timestamp=execution_date,
        testCaseStatus=status,
        result=result,
        testResultValue=[
            TestResultValue(name="min", value=str(min_value_res)),
            TestResultValue(name="max", value=str(max_value_res)),
        ],
    )
