import json
import logging
import random
import re
import string
import warnings
from typing import Any, NamedTuple

from inspect_ai.dataset import Sample
from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import Generate, Solver, TaskState, solver

logger = logging.getLogger(__name__)

DEFAULT_MULTICHOICE_TEMPLATE = """{question}

{choices}

Answer with the letter of the chosen answer in JSON: {{"answer": "<letter>"}}."""


class ShuffledChoices(NamedTuple):
    choices: list[str]
    """The shuffled choices"""

    shuffle_map: list[int]
    """Shuffle map containing the permuted order of choice indices.  That is,
    `shuffle_map[old_idx] == new_idx`."""


def remap_choice_letter(letter: str, shuffle_map: list[int]) -> str:
    """Remap a letter (A-Z) to its corresponding letter in the shuffled order.

    Args:
        letter: The letter to remap.
        shuffle_map: The shuffle map containing the permuatation of choice
                     indexes to use for remapping.  That is,
                     `shuffle_map[old_idx] == new_idx`.
    """
    assert letter in string.ascii_uppercase and len(letter) == 1
    old_idx = ord(letter) - ord("A")
    return string.ascii_uppercase[shuffle_map[old_idx]]


def shuffle_choices(choices: list[str], rng: random.Random) -> ShuffledChoices:
    """Shuffle the choices and return a mapping from old to new indices.

    Args:
        choices: The choices to shuffle.
        rng: Random number generator to use for shuffling.
    """
    perm = list(range(len(choices)))
    rng.shuffle(perm)
    return ShuffledChoices(
        choices=[choices[perm.index(new_idx)] for new_idx in range(len(choices))],
        shuffle_map=perm,
    )


def format_multichoice_as_textin(
    sample: Sample,
    template: str = DEFAULT_MULTICHOICE_TEMPLATE,
    shuffle: bool = False,
    rng: random.Random | None = None,
) -> Sample:
    """Converts a sample using Inspect's `sample.choices` to one that shows
    all choices in the input text and (by default) instructs the model to
    output its answer as text in a JSON format.

    *Don't* use Inspect's `dataset.shuffle_choices` after calling this, as that
    will update `sample.choices` but not the order of choices in the input
    text.

    Args:
        sample: The sample to convert.  Will be copied and returned.
        template: The template for the new input text.  Should contain
            `{question}` and `{choices}` placeholders and any instructions
            about the task and how to format the response.
        shuffle: Whether to shuffle the order of the choices in the output.
            This uses `shuffle_choices`.  Note that it can only be used if
            `target` is a single letter or list of letters; for more complex
            targets, use `shuffle_choices` directly and `remap_choice_letter`
            to convert any corresponding target values.
        rng: Random number generator for shuffling (required if shuffle=True)
    """
    sample = sample.model_copy(deep=True)
    assert sample.choices is not None
    assert sample.target is not None

    if shuffle:
        if not rng:
            raise ValueError("Must provide rng if shuffle is True.")
        shuffled_choices, shuffle_map = shuffle_choices(sample.choices, rng)
        sample.choices = shuffled_choices
        if isinstance(sample.target, str):
            sample.target = remap_choice_letter(sample.target, shuffle_map)
        elif isinstance(sample.target, list):
            sample.target = [remap_choice_letter(t, shuffle_map) for t in sample.target]
        else:
            raise ValueError(
                "Cannot shuffle choices for targets of type "
                f"{type(sample.target)}.  Only str and list[str] targets are "
                "supported."
            )

    if any("\n" in choice for choice in sample.choices):
        warnings.warn(
            "Newline found in sample.choices; this may cause ambiguity in the formatted choice list"
        )

    choices_text = "\n".join(
        f"{chr(ord('A') + idx)}. {choice_text}"
        for idx, choice_text in enumerate(sample.choices)
    )
    sample.input = template.format(
        question=sample.input,
        choices=choices_text,
    )
    return sample


async def mark_multichoice_answer(
    state: TaskState,
    override_marked_state=False,
) -> TaskState:
    """Parse JSON `{"answer": "<letter>"}` or labbench
    `[ANSWER]<letter>[/ANSWER]` or single-letter-output answer format and mark
    `state.choices` accordingly.

    Args:
        state: The task state.
        override_marked_state: Whether to override the marked state of the
            choices.  If this is False and *any* state.choices option is marked
            as True or False, this will not mark any choices.  If True, all
            existing state.choices marks will be overridden by the parsed value
            from text (all will be marked False if no answer is found in the
            text).
    """

    if not override_marked_state and any(c.correct is not None for c in state.choices):
        # Choices already marked; don't override
        return state

    chosen_idxs = set()

    # Parse state.output.completion to get raw answer text, and mark appropriately
    raw_output = state.output.completion

    # Single letter output
    if len(raw_output) == 1 and raw_output in string.ascii_uppercase:
        chosen_idxs.add(ord(raw_output) - ord("A"))

    if not chosen_idxs:
        # Try JSON
        # answer can be a single letter or (if multiple answer selections allowed) a list of letters
        m = re.search(
            r"\"answer\":\s*(?P<answer>\[?\s*(\"[A-Za-z]+\",\s*)*(\"[A-Za-z]+\")\s*\]?)",
            raw_output,
        )
        if m is not None:
            answer = json.loads(m.group("answer"))
            if isinstance(answer, list):
                for ans in answer:
                    chosen_idxs.add(ord(ans) - ord("A"))
            else:
                chosen_idxs.add(ord(answer) - ord("A"))

    if not chosen_idxs:
        # Try LAB-bench's answer format
        #
        # From chem-bench (MIT-licensed):
        # https://github.com/lamalab-org/chem-bench/blob/b5dc61583c867f598715e29c0e9fe4b68698f8dd/src/chembench/utils.py
        #
        # Copied here because it would be wasteful to introduce the dependency for
        # a single regex (indirectly grabs gigabytes of nvidia-cuda libs). LAB-bench
        # uses this to parse answers in
        # https://github.com/Future-House/LAB-Bench/blob/77c14efc2192a1b3d14f695e120c668eef72cd82/labbench/zero_shot.py.
        #
        # Parses [ANSWER]<letter>[/ANSWER] or <ANS><letter></ANS> or <ANSWER><letter></ANSWER>
        MCQ_REGEX_TEMPLATE_1 = r"(?:\[ANSWER\]|\[ANS\]|<ANS>|<ANSWER>|<ans>)\s*([A-Z](?:,\s*[A-Z])*)\s*(?:\..*?|,.*?)?(?:\[/?ANSWER\]|\[/ANS\]|</ANS>|</ANSWER>|</ans>)"

        try:
            letter_choices = set(
                re.findall(MCQ_REGEX_TEMPLATE_1, state.output.completion, re.DOTALL)
            )
            chosen_idxs = {
                string.ascii_uppercase.index(letter) for letter in letter_choices
            }
        except (ValueError, IndexError):
            # No choice found
            chosen_idxs = set()

    # Mark
    for idx in range(len(state.choices)):
        state.choices.mark_choice(idx, bool(idx in chosen_idxs))

    return state


@scorer(metrics=[accuracy(), stderr()])
def mark_choices_and_score(inner_scorer: Scorer) -> Scorer:
    """Scorer that marks the choices based on the output and then scores the
    output using the provided inner_scorer.

    Args:
        inner_scorer: The scorer to use to score the output after marking the
            choices.
    """

    async def score(state: TaskState, target: Target) -> Score:
        state = await mark_multichoice_answer(state)
        return await inner_scorer(state, target)

    return score


@solver
def not_implemented_solver() -> Solver:
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        raise NotImplementedError(
            "This is a placeholder solver.  Please specify a real solver using --solver from the cli."
        )

    return solve


def extract_json_from_response(response: str) -> dict[str, Any] | None:
    """Extract JSON from a response string.  Tries to be flexible with format,
    allowing other text around the json.  Returns None if no valid JSON."""
    json_start = response.find("{")
    json_end = response.rfind("}") + 1
    if json_start == -1 or json_end == -1:
        return None

    try:
        return json.loads(response[json_start:json_end])
    except json.JSONDecodeError:
        try:
            return json.loads(response[json_start + 1 : json_end - 1])
        except json.JSONDecodeError:
            logger.warning(
                f"Could not decode JSON from response: {response[json_start:json_end]}"
            )
        return None
