import json
from typing import Callable, Awaitable, Tuple

from inspect_ai.scorer import scorer, accuracy, stderr, Target, Score, INCORRECT, CORRECT
from inspect_ai.solver import TaskState


@scorer(metrics=[accuracy(), stderr()])
def config_in_range():
    async def score(state: TaskState, target: Target) -> Score:
        output = state.metadata["output_files"]["config.json"]

        if output is None:
            return Score(value=INCORRECT, answer="No output found",
                         explanation="config.json not found in the working directory.")

        try:
            config_json = json.loads(output)
        except json.JSONDecodeError as e:
            return Score(value=INCORRECT, answer="Invalid JSON", explanation=f"{e}\n\nconfig.json:\n{output}")

        config_key = state.metadata["CONFIG_KEY"]
        config_value = config_json.get(config_key, 0)

        if not isinstance(config_value, (int, float)):
            return Score(value=INCORRECT, answer=f"{config_key} is not a number", explanation=output)

        min_config_value = state.metadata.get("MIN_CONFIG_VALUE", -float("inf"))
        max_config_value = state.metadata.get("MAX_CONFIG_VALUE", float("inf"))

        if min_config_value <= config_value <= max_config_value:
            return Score(value=CORRECT, answer=f"{config_key} is {config_value}", explanation=output)
        else:
            return Score(value=INCORRECT, answer=f"{config_key} is {config_value}", explanation=output)

    return score


@scorer(metrics=[accuracy(), stderr()])
def valid_config(is_valid_config: Callable[[], Awaitable[Tuple[bool, str, str]]] | None = None):
    config_in_range_ = config_in_range()

    async def score(state: TaskState, target: Target) -> Score:
        score = await config_in_range_(state, target)
        if score.value == INCORRECT:
            return score

        # config_in_range_ should have already checked if the config exists/is valid
        config = json.loads(state.metadata["output_files"]["config.json"])

        if is_valid_config is None:
            return Score(value=CORRECT, explanation=f"config.json: {config}", answer="No validation performed")

        valid, stdout, stderr = await is_valid_config()

        if valid:
            return Score(value=CORRECT, explanation=f"config.json: {config}", answer=stdout + '\n' + stderr)
        else:
            return Score(value=INCORRECT, explanation=f"config.json: {config}", answer=stdout + '\n' + stderr)

    return score
