import os
from typing import List, Optional, Union, Dict
import time
from dataclasses import dataclass

from deepeval.utils import (
    drop_and_copy,
    get_or_create_event_loop,
    should_ignore_errors,
    should_use_cache,
)
from deepeval.telemetry import capture_evaluation_count
from deepeval.metrics import BaseMetric
from deepeval.metrics.indicator import (
    measure_metrics_with_indicator,
)
from deepeval.test_case import LLMTestCase, ConversationalTestCase
from deepeval.constants import PYTEST_RUN_TEST_NAME
from deepeval.test_run import (
    test_run_manager,
    LLMApiTestCase,
    ConversationalApiTestCase,
    MetricMetadata,
)
from deepeval.utils import get_is_running_deepeval, set_indicator
from deepeval.test_run.cache import (
    test_run_cache_manager,
    Cache,
    CachedTestCase,
    CachedMetricData,
)


@dataclass
class TestResult:
    """Returned from run_test"""

    success: bool
    metrics: List[BaseMetric]
    input: str
    actual_output: str
    expected_output: str
    context: List[str]
    retrieval_context: List[str]


def create_metric_metadata(metric: BaseMetric) -> MetricMetadata:
    if metric.error is not None:
        return MetricMetadata(
            metric=metric.__name__,
            threshold=metric.threshold,
            score=None,
            reason=None,
            success=False,
            strictMode=metric.strict_mode,
            evaluationModel=metric.evaluation_model,
            error=metric.error,
            evaluationCost=metric.evaluation_cost,
        )
    else:
        return MetricMetadata(
            metric=metric.__name__,
            score=metric.score,
            threshold=metric.threshold,
            reason=metric.reason,
            success=metric.is_successful(),
            strictMode=metric.strict_mode,
            evaluationModel=metric.evaluation_model,
            error=None,
            evaluationCost=metric.evaluation_cost,
        )


def create_test_result(
    test_case: Union[LLMApiTestCase, ConversationalApiTestCase],
    metrics: List[BaseMetric],
) -> TestResult:
    if isinstance(test_case, ConversationalApiTestCase):
        tc = test_case.messages[len(test_case.messages) - 1]
    else:
        tc = test_case

    return TestResult(
        success=tc.success,
        metrics=metrics,
        input=tc.input,
        actual_output=tc.actual_output,
        expected_output=tc.expected_output,
        context=tc.context,
        retrieval_context=tc.retrieval_context,
    )


def create_api_test_case(
    test_case: Union[LLMTestCase, ConversationalTestCase],
    index: Optional[int] = None,
    is_message: Optional[bool] = False,
    additional_metadata: Optional[Dict] = None,
    comments: Optional[str] = None,
) -> Union[LLMApiTestCase, ConversationalApiTestCase]:
    if isinstance(test_case, LLMTestCase):
        if is_message:
            success = None
            name = f"message_{index}"
            order = index

            # Manually set the metadata and comments on conversational test case
            # to each individual message (test case)
            test_case.additional_metadata = additional_metadata
            test_case.comments = comments
        else:
            success = True
            name = os.getenv(PYTEST_RUN_TEST_NAME, f"test_case_{index}")
            order = test_case._dataset_rank

        return LLMApiTestCase(
            name=name,
            input=test_case.input,
            actualOutput=test_case.actual_output,
            expectedOutput=test_case.expected_output,
            context=test_case.context,
            retrievalContext=test_case.retrieval_context,
            success=success,
            metricsMetadata=None,
            runDuration=None,
            evaluationCost=None,
            order=order,
            additionalMetadata=test_case.additional_metadata,
            comments=test_case.comments,
        )
    elif isinstance(test_case, ConversationalTestCase):
        return ConversationalApiTestCase(
            name=os.getenv(
                PYTEST_RUN_TEST_NAME, f"conversational_test_case_{index}"
            ),
            success=True,
            metricsMetadata=None,
            runDuration=0,
            evaluationCost=None,
            order=test_case._dataset_rank,
            testCases=[
                create_api_test_case(
                    tc,
                    i,
                    True,
                    test_case.additional_metadata,
                    test_case.comments,
                )
                for i, tc in enumerate(test_case.messages)
            ],
        )


def execute_test_cases(
    test_cases: List[Union[LLMTestCase, ConversationalTestCase]],
    metrics: List[BaseMetric],
    ignore_errors: bool,
    use_cache: bool,
    save_to_disk: bool = False,
) -> List[TestResult]:
    test_results: List[TestResult] = []
    test_run_cache_manager.disable_write_cache = save_to_disk == False
    test_run_manager.save_to_disk = save_to_disk
    test_run = test_run_manager.get_test_run()
    for index, test_case in enumerate(test_cases):
        if isinstance(test_case, ConversationalTestCase):
            last_message_index = len(test_case.messages) - 1
        else:
            last_message_index = -1

        cached_test_case = None
        if use_cache and isinstance(
            test_case, LLMTestCase
        ):  # for now, don't use cache when it is a conversation
            cached_test_case = test_run_cache_manager.get_cached_test_case(
                test_case, test_run.hyperparameters
            )

        ##### Metric Calculation #####
        # this can be a converational api or llm api test case
        api_test_case = create_api_test_case(test_case, index)
        new_cached_test_case: CachedTestCase = CachedTestCase()
        test_start_time = time.perf_counter()

        for metric in metrics:
            metric_metadata = None
            # cached_tet_case will always be false for conversationals
            if cached_test_case is not None:
                cached_metric_data = Cache.get_metric_data(
                    metric, cached_test_case
                )
                if cached_metric_data:
                    metric_metadata = cached_metric_data.metric_metadata

            # metric_metadata will always be false for conversationals
            if metric_metadata is None:
                metric.async_mode = False  # Override metric async
                try:
                    # if conversational, manually extract the last one for now
                    # In the future, might add support for evaluating any message in a convo
                    if isinstance(test_case, ConversationalTestCase):
                        tc = test_case.messages[last_message_index]
                    else:
                        tc = test_case
                    metric.measure(tc)
                except Exception as e:
                    if ignore_errors:
                        metric.error = str(e)  # Override metric error
                        metric.success = False  # Override metric success
                    else:
                        raise
                metric_metadata = create_metric_metadata(metric)

            if isinstance(test_case, ConversationalTestCase):
                # the index can be dynamic in the future, just not now
                api_test_case.update(metric_metadata, last_message_index)
            else:
                api_test_case.update(metric_metadata)

            if metric.error is None and isinstance(
                test_case, LLMTestCase
            ):  # Only save to cache if metric didn't error and not conversational
                cache_metric_metadata = create_metric_metadata(metric)
                cache_metric_metadata.evaluation_cost = (
                    0  # Create copy and save 0 for cost
                )
                updated_cached_metric_data = CachedMetricData(
                    metric_metadata=cache_metric_metadata,
                    metric_configuration=Cache.create_metric_configuration(
                        metric
                    ),
                )
                new_cached_test_case.cached_metrics_data.append(
                    updated_cached_metric_data
                )

        test_end_time = time.perf_counter()
        run_duration = test_end_time - test_start_time
        api_test_case.run_duration = run_duration

        ### Update Test Run ###
        test_run_manager.update_test_run(api_test_case, test_case)

        ### Cache Test Run ###
        if isinstance(
            test_case, LLMTestCase
        ):  # only cache if not conversational
            test_run_cache_manager.cache_test_case(
                test_case,
                new_cached_test_case,
                test_run.hyperparameters,
            )
            test_run_cache_manager.cache_test_case(
                test_case,
                new_cached_test_case,
                test_run.hyperparameters,
                to_temp=True,
            )

        test_result = create_test_result(
            api_test_case, drop_and_copy(metrics, ["model", "embeddings"])
        )
        test_results.append(test_result)

    return test_results


async def a_execute_test_cases(
    test_cases: List[Union[LLMTestCase, ConversationalTestCase]],
    metrics: List[BaseMetric],
    ignore_errors: bool,
    use_cache: bool,
    save_to_disk: bool = False,
) -> List[TestResult]:
    test_results: List[TestResult] = []
    test_run_cache_manager.disable_write_cache = save_to_disk == False
    test_run_manager.save_to_disk = save_to_disk
    test_run = test_run_manager.get_test_run()
    for index, test_case in enumerate(test_cases):
        cached_test_case = None
        # only use cache when NOT conversational test case
        if use_cache and isinstance(test_case, LLMTestCase):
            cached_test_case = test_run_cache_manager.get_cached_test_case(
                test_case,
                test_run.hyperparameters,
            )

        ##### Metric Calculation #####
        # api test case can be conversational
        api_test_case = create_api_test_case(test_case, index)

        new_cached_test_case: CachedTestCase = CachedTestCase()
        test_start_time = time.perf_counter()
        await measure_metrics_with_indicator(
            metrics, test_case, cached_test_case, ignore_errors
        )
        for metric in metrics:
            metric_metadata = create_metric_metadata(metric)

            if isinstance(test_case, ConversationalTestCase):
                # index hardcoded as the last message for now
                api_test_case.update(
                    metric_metadata, len(test_case.messages) - 1
                )
            else:
                api_test_case.update(metric_metadata)

            if metric.error is None and isinstance(
                test_case, LLMTestCase
            ):  # Only save to cache if metric didn't error AND is not conversational
                cache_metric_metadata = create_metric_metadata(metric)
                cache_metric_metadata.evaluation_cost = (
                    0  # Create new copy and save 0 for cost
                )
                updated_cached_metric_data = CachedMetricData(
                    metric_metadata=cache_metric_metadata,
                    metric_configuration=Cache.create_metric_configuration(
                        metric
                    ),
                )
                new_cached_test_case.cached_metrics_data.append(
                    updated_cached_metric_data
                )

        test_end_time = time.perf_counter()
        run_duration = test_end_time - test_start_time
        api_test_case.run_duration = run_duration

        ### Update Test Run ###
        test_run_manager.update_test_run(api_test_case, test_case)

        ### Cache Test Run ###
        if isinstance(
            test_case, LLMTestCase
        ):  # only cache if not conversational
            test_run_cache_manager.cache_test_case(
                test_case,
                new_cached_test_case,
                test_run.hyperparameters,
            )
            test_run_cache_manager.cache_test_case(
                test_case,
                new_cached_test_case,
                test_run.hyperparameters,
                to_temp=True,
            )

        test_result = create_test_result(
            api_test_case, drop_and_copy(metrics, ["model", "embeddings"])
        )
        test_results.append(test_result)

    return test_results


def assert_test(
    test_case: Union[LLMTestCase, ConversationalTestCase],
    metrics: List[BaseMetric],
    run_async: bool = True,
):

    # TODO: keep this for now, blocking conversational metrics like KR
    for metric in metrics:
        if not isinstance(metric, BaseMetric):
            raise TypeError("Provided 'metric' must be of type 'BaseMetric'.")

    if run_async:
        loop = get_or_create_event_loop()
        test_result = loop.run_until_complete(
            a_execute_test_cases(
                [test_case],
                metrics,
                ignore_errors=should_ignore_errors(),
                use_cache=should_use_cache(),
                save_to_disk=get_is_running_deepeval(),
            )
        )[0]
    else:
        test_result = execute_test_cases(
            [test_case],
            metrics,
            ignore_errors=should_ignore_errors(),
            use_cache=should_use_cache(),
            save_to_disk=get_is_running_deepeval(),
        )[0]

    if not test_result.success:
        failed_metrics: List[BaseMetric] = []
        # even for conversations, test_result right now is just the
        # result for the last message
        for metric in test_result.metrics:
            if metric.error is not None:
                failed_metrics.append(metric)
            else:
                # This try block is for user defined custom metrics,
                # which might not handle the score == undefined case elegantly
                try:
                    if not metric.is_successful():
                        failed_metrics.append(metric)
                except:
                    failed_metrics.append(metric)

        failed_metrics_str = ", ".join(
            [
                f"{metric.__name__} (score: {metric.score}, threshold: {metric.threshold}, strict: {metric.strict_mode}, error: {metric.error})"
                for metric in failed_metrics
            ]
        )
        raise AssertionError(f"Metrics: {failed_metrics_str} failed.")


def evaluate(
    test_cases: List[Union[LLMTestCase, ConversationalTestCase]],
    metrics: List[BaseMetric],
    run_async: bool = True,
    show_indicator: bool = True,
    print_results: bool = True,
    write_cache: bool = True,
    use_cache: bool = False,
    ignore_errors: bool = False,
):
    set_indicator(show_indicator)

    # TODO: keep this for now, blocking conversational metrics like KR
    for metric in metrics:
        if not isinstance(metric, BaseMetric):
            raise TypeError("Provided 'metric' must be of type 'BaseMetric'.")

    test_run_manager.reset()
    start_time = time.perf_counter()
    if print_results:
        print("Evaluating test cases...")
    if run_async:
        loop = get_or_create_event_loop()
        test_results = loop.run_until_complete(
            a_execute_test_cases(
                test_cases,
                metrics,
                ignore_errors=ignore_errors,
                use_cache=use_cache,
                save_to_disk=write_cache,
            )
        )
    else:
        test_results = execute_test_cases(
            test_cases,
            metrics,
            ignore_errors=ignore_errors,
            use_cache=use_cache,
            save_to_disk=write_cache,
        )

    capture_evaluation_count()
    end_time = time.perf_counter()
    run_duration = end_time - start_time
    if print_results:
        for test_result in test_results:
            print_test_result(test_result)

        print("")
        print("-" * 70)

    test_run_manager.wrap_up_test_run(run_duration, display_table=False)
    return test_results


def print_test_result(test_result: TestResult):
    print("")
    print("=" * 70 + "\n")
    print("Metrics Summary\n")
    for metric in test_result.metrics:
        successful = True
        if metric.error is not None:
            successful = False
        else:
            # This try block is for user defined custom metrics,
            # which might not handle the score == undefined case elegantly
            try:
                if not metric.is_successful():
                    successful = False
            except:
                successful = False

        if not successful:
            print(
                f"  - ❌ {metric.__name__} (score: {metric.score}, threshold: {metric.threshold}, strict: {metric.strict_mode}, evaluation model: {metric.evaluation_model}, reason: {metric.reason}, error: {metric.error})"
            )
        else:
            print(
                f"  - ✅ {metric.__name__} (score: {metric.score}, threshold: {metric.threshold}, strict: {metric.strict_mode}, evaluation model: {metric.evaluation_model}, reason: {metric.reason}, error: {metric.error})"
            )
        if metric.score_breakdown:
            for metric_name, score in metric.score_breakdown.items():
                print(f"      - {metric_name} (score: {score})")

    print("")
    print("For test case:\n")
    print(f"  - input: {test_result.input}")
    print(f"  - actual output: {test_result.actual_output}")
    print(f"  - expected output: {test_result.expected_output}")
    print(f"  - context: {test_result.context}")
    print(f"  - retrieval context: {test_result.retrieval_context}")
