import json
from random import Random
from typing import Generator

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, json_dataset
from inspect_ai.model import ModelOutput, get_model
from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import TaskState
from pydantic import BaseModel

from astabench.evals.demo.arithmetic.task import INPUT_TEMPLATE, TASK_INSTRUCTIONS
from astabench.evals.utils import not_implemented_solver


class Rubric(BaseModel):
    category: str
    description: str
    weight: float


RUBRIC_GRADING_TEMPLATE = (
    "Grade the following response using the provided rubric:"
    "\n\n# Task:\n\n"
    "{prompt}"
    "\n\n# Model answer:\n\n"
    "{model_answer}"
    "\n\n# Rubric instructions:\n\n"
    "{grading_rubric}"
    "\n\n Grade the answer based on the rubric by outputting a score between 0 and 1 (float):"
)


def json_to_sample_style_rubric(obj: dict[str, str]) -> Sample:
    return Sample(
        id=obj["id"],
        input=INPUT_TEMPLATE.format(
            task_instructions=TASK_INSTRUCTIONS,
            question=f"{obj['question']}\n\n{obj['question_stylistic_instructions']}",
        ),
        target=json.dumps(
            {
                "answer": obj["answer"],
                "rubrics": obj["rubrics"],
            }
        ),
    )


def check_rubric(scorer_model: str, rubric_category: str) -> Scorer:
    """Scorer to check the solution based on custom rubrics."""

    # If this were a real task, we'd load an actual grading model here
    # model = get_model(scorer_model)
    # Instead, we'll use a mock model that generates random numbers
    def infinite_random_generator() -> Generator[ModelOutput, None, None]:
        rng = Random(0)
        while True:
            yield ModelOutput.from_content(
                model="mockllm/model", content=str(rng.random())
            )

    model = get_model("mockllm/model", custom_outputs=infinite_random_generator())

    async def score(state: TaskState, target: Target) -> Score:
        raw_output = state.output.completion

        rubrics = json.loads(target.text)["rubrics"]
        rubric_objects = []
        for rubric in rubrics:
            rubric_objects.append(Rubric.model_validate(rubric))

        weighted_scores = []
        for rubric in rubric_objects:
            if rubric.category == rubric_category:
                output = await model.generate(
                    input=RUBRIC_GRADING_TEMPLATE.format(
                        prompt=state.input,
                        model_answer=raw_output,
                        grading_rubric=rubric.description,
                    )
                )
                content = output.choices[0].message.content
                if isinstance(content, str):
                    try:
                        score = float(content)
                    except ValueError:
                        raise ValueError(
                            f"Invalid score '{content}' for rubric category '{rubric.category}'"
                        )
                else:
                    raise ValueError(
                        f"Invalid score '{content}' for rubric category '{rubric.category}'"
                    )
                weighted_scores.append(score * rubric.weight)
        if len(weighted_scores) > 0:
            s = sum(weighted_scores)
        else:
            raise ValueError("No rubric found for the specified category")

        return Score(
            value=s,
            answer=raw_output,
            explanation=f"Weighted scores for rubric category '{rubric_category}': {weighted_scores} (total score = {s})",
        )

    return score


@scorer(metrics=[accuracy(), stderr()])
def check_rubric_style(scorer_model: str) -> Scorer:
    """Scorer that checks just rubrics with category "style"."""
    return check_rubric(scorer_model, "style")


@task
def arithmetic_with_style_rubric(
    scorer_model: str = "openai/gpt-4o-2024-08-06",
) -> Task:
    return Task(
        dataset=json_dataset("data.json", sample_fields=json_to_sample_style_rubric),
        # Solver that just tells the user to specify a real solver via cli
        solver=not_implemented_solver(),
        scorer=[check_rubric_style(scorer_model)],
    )
