import json
from typing import Any

from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import TaskState
from inspect_ai.util import sandbox

EVALUATE_COMMAND = ["python3", "-m", "desktop_env.cli", "evaluate"]


@scorer(metrics=[accuracy(), stderr()])
def osworld_scorer() -> Scorer:
    async def check(state: TaskState, target: Target) -> Score:
        sb = sandbox()

        evaluate_result = await sb.exec(EVALUATE_COMMAND)
        if not evaluate_result.success:
            raise RuntimeError(
                "evaluator failed", evaluate_result.returncode, evaluate_result.stderr
            )

        # infeasible samples are scored here since the scoring relies on the TaskState
        if evaluate_result.stdout.strip() == "infeasible":
            return _score_infeasible(state)
        else:
            try:
                result_data: dict[str, Any] = json.loads(evaluate_result.stdout)
                return Score(
                    value=result_data["score"], explanation=result_data["explanation"]
                )
            except ValueError as e:
                raise RuntimeError(
                    "Failed to parse score from evaluator output",
                    evaluate_result.stdout,
                ) from e

    return check


def _score_infeasible(state: TaskState) -> Score:
    if state.output.completion.startswith("FAIL"):
        explanation = state.output.completion[4:].strip()
        return Score(
            value=1,
            explanation=f"Properly determined to be infeasible - {explanation}"
            if explanation
            else "Properly determined to be infeasible",
        )
    else:
        return Score(
            value=0,
            explanation=f"Failed to determine infeasibility. Last message: {state.output.completion}",
        )
