import asyncio
from typing import Awaitable, Dict, List, Union

from maxim.evaluators.base_evaluator import BaseEvaluator

from ..apis.maxim_apis import MaximAPI
from ..models.dataset import ManualData
from ..models.evaluator import (
    LocalEvaluationResult,
    LocalEvaluatorResultParameter,
    LocalEvaluatorReturn,
)
from ..models.test_run import YieldedOutput, YieldedOutputMeta, YieldedOutputTokenUsage


async def process_awaitable(awaitable: Awaitable):
    return await awaitable


async def run_local_evaluations(
    evaluators: List[BaseEvaluator],
    data_entry: ManualData,
    processed_data: LocalEvaluatorResultParameter,
) -> List[LocalEvaluationResult]:
    coroutines = [
        asyncio.to_thread(evaluator.guarded_evaluate, processed_data, data_entry)
        for evaluator in evaluators
    ]
    evaluator_results = await asyncio.gather(*coroutines)
    results: List[LocalEvaluationResult] = []
    for i, evaluator in enumerate(evaluators):
        if isinstance(evaluator, BaseEvaluator):
            try:
                combined_results = evaluator_results[i]
                for name, result in combined_results.items():
                    results.append(LocalEvaluationResult(
                        name=name,
                        pass_fail_criteria=evaluator.pass_fail_criteria[name],
                        result=result
                    ))
            except Exception as err:
                results.extend(
                    [
                        LocalEvaluationResult(
                            name=name,
                            pass_fail_criteria=evaluator.pass_fail_criteria[name],
                            result=LocalEvaluatorReturn(
                                score="Err",
                                reasoning=f"Error while running combined evaluator with names {evaluator.names}: {str(err)}",
                            ),
                        )
                        for name in evaluator.names
                    ]
                )
    return results


def workflow_id_output_function_closure(workflow_id: str, api_service: MaximAPI):
    async def output_function(
        data: Dict[str, Union[str, List[str], None]],
    ) -> YieldedOutput:
        result = api_service.execute_workflow_for_data(
            workflow_id=workflow_id, data_entry=data
        )

        output: str
        if result.output is None:
            output = ""
        else:
            output = result.output

        return YieldedOutput(
            data=output,
            retrieved_context_to_evaluate=result.context_to_evaluate,
            meta=YieldedOutputMeta(
                usage=YieldedOutputTokenUsage(
                    latency=result.latency,
                    total_tokens=0,
                    prompt_tokens=0,
                    completion_tokens=0,
                )
            ),
        )

    return output_function


def prompt_version_id_output_function_closure(
    prompt_version_id: str, input_data: str, api_service: MaximAPI
):
    async def output_function(
        data: Dict[str, Union[str, List[str], None]],
    ) -> YieldedOutput:
        result = api_service.execute_prompt_for_data(
            prompt_version_id=prompt_version_id, input=input_data, data_entry=data
        )

        output: str
        if result.output is None:
            output = ""
        else:
            output = result.output

        return YieldedOutput(
            data=output,
            retrieved_context_to_evaluate=result.context_to_evaluate,
            meta=YieldedOutputMeta(
                cost=result.cost if result.cost is not None else None,
                usage=result.usage if result.usage is not None else None,
            ),
        )

    return output_function
