from json import JSONDecodeError
from typing import Any, List, Literal, Tuple, TypeVar

from .config import execute_configs
from .env import OSWorldDesktopEnv
from .evaluators import getters as getters_module
from .evaluators import metrics as metrics_module
from .execute import execute
from .logging import logger
from .osworld_types import (
    Evaluator,
    Getter,
    GetterCloudFileMulti,
    GetterCloudFileSingle,
    GetterVMCommandLineOrError,
)
from .util import RedirectStdFdsToDockerLogs, to_getter_list


def evaluate(
    env: OSWorldDesktopEnv,
    evaluator: Evaluator,
) -> Tuple[float | Literal["infeasible"], str | None]:
    # `infeasible` is special cased since it's handled in the host by the inspect scorer
    if evaluator.func == "infeasible":
        return "infeasible", None

    if evaluator.postconfig:
        execute_configs(env, evaluator.postconfig)

    metric_conjunction = evaluator.conj or "and"
    metrics = evaluator.func if isinstance(evaluator.func, list) else [evaluator.func]
    metric_options: list[dict[str, Any]] = (
        evaluator.options
        if isinstance(evaluator.options, list)
        else [{}] * len(metrics)
        if evaluator.options is None
        else [evaluator.options]
    )
    expected_getters = to_getter_list(evaluator.expected)
    result_getters = to_getter_list(evaluator.result)

    # if there are multiple metrics, there must be the same number of expected getters, result
    # getters and metrics options
    assert len(metrics) > 0 and (
        len(metrics) == 1
        or (
            len(metrics)
            == len(metric_options)
            == len(result_getters)
            == len(expected_getters)
        )
    )

    if len(metrics) == 1:
        return _call_metric(
            env,
            _safe_get(expected_getters, 0),
            _safe_get(result_getters, 0),
            metrics[0],
            _safe_get(metric_options, 0),
        )
    else:
        scores: list[float] = []
        explanations: list[str] = []
        for expected_getter, result_getter, metric, metric_option in zip(
            expected_getters, result_getters, metrics, metric_options
        ):
            score, explanation = _call_metric(
                env,
                expected_getter,
                result_getter,
                metric,
                metric_option,
            )
            scores.append(score)
            if explanation:
                explanations.append(explanation)
            if score == 0 and metric_conjunction == "and":
                return 0, "\n".join(explanations) if explanations else None
            elif score == 1 and metric_conjunction == "or":
                return 1, "\n".join(explanations) if explanations else None
        return sum(scores) / len(scores), "\n".join(
            explanations
        ) if explanations else None


def _call_metric(
    env: OSWorldDesktopEnv,
    expected_getter: Getter | None,
    result_getter: Getter | None,
    metric_name: str,
    metric_options: dict[str, Any] | None,
) -> Tuple[float, str | None]:
    expected_value = _call_getter(env, expected_getter) if expected_getter else None
    try:
        result_value = _call_getter(env, result_getter) if result_getter else None
    except FileNotFoundError as e:
        logger.warning(
            f"call to result getter raised FileNotFound({e.filename}). Returning metric of 0."
        )
        return 0, f"The file '{e.filename}' was not found."

    logger.info(f"calling metric func '{metric_name}' with '{metric_options}'")
    metric_func = getattr(metrics_module, metric_name)
    try:
        with RedirectStdFdsToDockerLogs():
            metric: float = (
                metric_func(result_value, expected_value, **metric_options)
                if expected_value is not None
                else metric_func(result_value, **metric_options)
            )
        logger.info(f"{metric_name} returned {metric}")
        return (
            metric,
            f'{metric_name}("{str(result_value)}", "{str(expected_value)}") returned {metric}',
        )
    except JSONDecodeError:
        return 0, f"{metric_name} encountered invalid JSON"


def _call_getter(env: OSWorldDesktopEnv, getter: Getter) -> Any:
    """
    Calls the getter function implemented by OSWorld code.

    Getters return type varies by function. `vm_file`, for example, return the path in the
    cache directory where the file was downloaded. `rule`, on the other hand, simply returns
    config["rules"].
    """
    if (override_result := _handle_getter_override(getter)) is not None:
        return override_result

    config = getter.model_dump(exclude={"type"})
    logger.info(f"calling getter 'get_{getter.type}' with {config}")
    with RedirectStdFdsToDockerLogs():
        result = (getattr(getters_module, f"get_{getter.type}"))(env, config)
    logger.debug(f"getter returned '{result}'")
    return result


def _handle_getter_override(getter: Getter) -> Any | None:
    """
    Handles the override for a given getter.

    A small number of OSWorld getters must be overridden because of differences in the Inspect
    environment as compared to the OSWorld native environment.
    """
    if getter.type == "vm_command_line" or getter.type == "vm_command_error":
        return _get_vm_command_line_or_error_override(getter)

    if getter.type == "cloud_file":
        return _get_cloud_file_override(getter)

    return None


def _get_cloud_file_override(
    getter: GetterCloudFileSingle | GetterCloudFileMulti,
) -> str | list[str]:
    """cloud_file is overridden since we don't want the image to use the network the files are retrieved by the host and placed into the container up front"""
    if getter.multi is True:
        # gives is inscrutable
        # e.g. why would you have multiple path/dest's but not have a gives. or more generally, why
        # do you need gives at all? boggling.
        gives = getter.gives or [0]
        if len(gives) == 1:
            return f"/tmp/osworld/cache/{getter.dest[gives[0]]}"
        return [
            f"/tmp/osworld/cache/{d}"
            for d in [dest for idx, dest in enumerate(getter.dest) if idx in gives]
        ]
    else:
        return f"/tmp/osworld/cache/{getter.dest}"


def _get_vm_command_line_or_error_override(
    getter: GetterVMCommandLineOrError,
) -> str:
    """vm_command_line and vm_command_error are overridden since the OSWorld getters are implemented in terms of a POST the the VM's http server - which we don't have"""
    res = execute(getter.command, getter.shell)
    if getter.type == "vm_command_line":
        if res.status == "error":
            raise RuntimeError(
                f"Failed to execute vm_command_line '{getter.command}: {res.message}'"
            )
        return res.output
    else:
        if res.status == "success":
            raise RuntimeError(
                f"Success result from vm_command_error '{getter.command}: {res.output}'"
            )
        return res.message


T = TypeVar("T")


def _safe_get(li: List[T], index: int) -> T | None:
    try:
        return li[index]
    except IndexError:
        return None
