import logging
import re
import sys
from abc import ABC
from collections.abc import Iterable, Mapping, Sequence
from enum import StrEnum, unique
from typing import TYPE_CHECKING, Any, assert_never, cast
from uuid import UUID

from aviary.core import (
    DEFAULT_EVAL_MODEL_NAME,
    TASK_DATASET_REGISTRY,
    Environment,
    Frame,
    MultipleChoiceQuestion,
    TaskDataset,
    ToolResponseMessage,
)
from ldp.alg import Callback, ComputeTrajectoryMetricsMixin, bulk_evaluate_consensus
from ldp.data_structures import Trajectory
from lmi import LLMModel
from paperqa.agents.tools import Complete, EnvironmentState
from paperqa.docs import Docs
from paperqa.settings import Settings
from paperqa.types import DocDetails, PQASession

from aviary.envs.labbench import (
    DEFAULT_REWARD_MAPPING,
    GradablePaperQAEnvironment,
    ImageQAEnvironment,
)

if TYPE_CHECKING:
    import pandas as pd
    from ldp.agent import Agent
    from ldp.data_structures import Transition

if sys.version_info >= (3, 13):
    from typing import TypeVar
else:
    from typing_extensions import TypeVar  # For TypeVar.default backport

logger = logging.getLogger(__name__)

TGradableEnv = TypeVar(
    "TGradableEnv",
    default=GradablePaperQAEnvironment,
    bound=GradablePaperQAEnvironment[Any],
)

DEFAULT_LABBENCH_HF_HUB_NAME = "futurehouse/lab-bench"
DEFAULT_LITQA3_HF_HUB_NAME = "futurehouse/litqa3"
# Test split from Aviary paper's section 4.3: https://doi.org/10.48550/arXiv.2412.21154
DEFAULT_AVIARY_PAPER_HF_HUB_NAME = "futurehouse/aviary-paper-data"


async def evaluate_consensus_sampling(
    data: Iterable[GradablePaperQAEnvironment | Frame],
    exclude_no_answer: bool = False,
    num_samples: int = 1,
    seed: int | None = None,
) -> tuple[dict[str, list[tuple[str, int]]], float]:
    """
    Create consensus groups based on question and evaluate the consensus for each.

    Args:
        data: Data to evaluate consensus upon, either gradable environments or frames.
        exclude_no_answer: Opt-in flag to filter out empty answers (due to the
            Environment/Frame not having a graded answer). Use of this flag does not
            affect the accuracy term of the return.
        num_samples: Passed through to evaluate_consensus.
        seed: Passed through to evaluate_consensus.

    Returns:
        Two-tuple of consensus list generated by collections.Counter.most_common (keys
            are question, values are list of (answer, vote count)) and the proportion of
            groups for which the consensus matches the ideal.
    """

    def extract_question(x: GradablePaperQAEnvironment | Frame) -> str:
        if isinstance(x, GradablePaperQAEnvironment):
            query: str | MultipleChoiceQuestion | dict[str, Any] = x._query
        else:
            query = x.info["query"]  # type: ignore[call-overload,index]
        if isinstance(query, str):
            return query
        if isinstance(query, MultipleChoiceQuestion):
            return query.question_prompt
        return query["question"]

    def extract_answer(x: GradablePaperQAEnvironment | Frame) -> str:
        ses: PQASession | dict[str, Any] = (
            x.state.session
            if isinstance(x.state, EnvironmentState)
            else cast("PQASession | dict[str, Any]", x.state["session"])  # type: ignore[call-overload,index]
        )
        graded_answer = (
            ses.graded_answer if isinstance(ses, PQASession) else ses["graded_answer"]
        )
        # One can filter the below empty string injection via the exclude_no_answer arg
        return graded_answer or ""

    def extract_ideal(x: GradablePaperQAEnvironment | Frame) -> str:
        if isinstance(x, GradablePaperQAEnvironment):
            query: str | MultipleChoiceQuestion | dict[str, Any] = x._query
        else:
            query = x.info["query"]  # type: ignore[call-overload,index]
        if isinstance(query, str):
            raise ValueError(  # noqa: TRY004
                f"We require a {MultipleChoiceQuestion.__name__} variant to extract"
                " ideal answer, not a string."
            )
        if isinstance(query, MultipleChoiceQuestion):
            return query.ideal_answer
        return query["ideal_answer"]

    try:
        consensus, accuracy = await bulk_evaluate_consensus(
            data=data,
            grouping_fn=extract_question,
            extract_answer_fn=extract_answer,
            ideal_answer_fn=extract_ideal,
            num_samples=num_samples,
            seed=seed,
        )
    except TypeError:
        raise ImportError(
            "Evaluating consensus requires the 'ldp' extra for 'ldp'. Please:"
            " `pip install paper-qa[ldp]`."
        ) from None
    if exclude_no_answer:
        consensus = {
            q: [(a, c) for a, c in answers if a] for q, answers in consensus.items()
        }
    return consensus, accuracy


class StoreForConsensusSamplingCallback(Callback):
    """Store environments or frames for later consensus sampling."""

    def __init__(self):
        super().__init__()
        self.stored: list[GradablePaperQAEnvironment | Frame] = []

    async def after_transition(
        self,
        traj_id: str,
        agent: "Agent",
        env: Environment,
        transition: "Transition",
    ) -> None:
        if not isinstance(env, GradablePaperQAEnvironment):
            raise NotImplementedError(
                f"So far only handled {GradablePaperQAEnvironment} in this callback,"
                f" not {type(env)}."
            )
        if transition.done and not transition.failed:  # Only store once
            return
        self.stored.append(env.export_frame())

    async def evaluate_consensus_sampling(
        self, num_samples: int = 1, seed: int | None = None
    ) -> tuple[dict[str, list[tuple[str, int]]], float]:
        return await evaluate_consensus_sampling(
            data=self.stored, num_samples=num_samples, seed=seed
        )


@unique
class LABBenchDatasets(StrEnum):
    """LAB-Bench datasets supported by this package."""

    # NOTE: keys' underscore before QA is supposed to make for easier reading
    FIG_QA = "FigQA"
    FIG_QA2 = "FigQA2"
    LIT_QA2 = "LitQA2"
    TABLE_QA = "TableQA"
    TABLE_QA2 = "TableQA2"

    @property
    def images_column(self) -> str:
        if self == LABBenchDatasets.FIG_QA:
            return "figure"
        if self == LABBenchDatasets.TABLE_QA:
            return "tables"
        raise ValueError(f"Dataset {self.value!r} does not have an images column.")

    @property
    def paths_column(self) -> str:
        if self == LABBenchDatasets.FIG_QA:
            return "figure-path"
        if self == LABBenchDatasets.TABLE_QA:
            return "table-path"
        raise ValueError(f"Dataset {self.value!r} does not have a paths column.")

    @property
    def key_passage_column(self) -> str:
        if self == LABBenchDatasets.LIT_QA2:
            return "key-passage"
        if self in {LABBenchDatasets.FIG_QA2, LABBenchDatasets.TABLE_QA2}:
            return "key_passage"
        raise ValueError(f"Dataset {self.value!r} does not have a key passage column.")

    @property
    def is_open_answer(self) -> bool:
        return self in {LABBenchDatasets.FIG_QA2, LABBenchDatasets.TABLE_QA2}

    def get_data(
        self, split: "str | TextQATaskSplit" = "train", **kwargs
    ) -> "pd.DataFrame":
        split = TextQATaskSplit(split)
        hf_name: LABBenchDatasets | None = self
        if split == TextQATaskSplit.TRAIN:
            if self not in {LABBenchDatasets.FIG_QA2, LABBenchDatasets.TABLE_QA2}:
                hf_path: str = DEFAULT_LABBENCH_HF_HUB_NAME
            else:
                hf_path = DEFAULT_LITQA3_HF_HUB_NAME
                hf_name = None
        elif self == LABBenchDatasets.LIT_QA2:
            hf_path = DEFAULT_AVIARY_PAPER_HF_HUB_NAME
        else:
            raise ValueError(f"Dataset {self.value!r} does not have a 'test' split.")
        ds_df = read_ds_from_hub(
            hf_name, hf_path=hf_path, hf_split=split.value, **kwargs
        )
        if self not in {LABBenchDatasets.FIG_QA2, LABBenchDatasets.TABLE_QA2}:
            return ds_df
        return ds_df[ds_df["tag"] == self.value.lower()]

    def get_sources(self, row: "pd.Series") -> list[str]:
        if self == LABBenchDatasets.LIT_QA2:
            raw_sources = row.sources
        # Ignore PLR1714 because `mypy==0.16.0` can't understand
        # `assert_never` with set.__contains__
        elif self == LABBenchDatasets.FIG_QA or self == LABBenchDatasets.TABLE_QA:  # noqa: PLR1714
            raw_sources = [row.source]
        elif self == LABBenchDatasets.FIG_QA2 or self == LABBenchDatasets.TABLE_QA2:  # noqa: PLR1714
            raw_sources = list(row.sources)
        else:
            assert_never(self)
        sources: list[str] = []
        for s in raw_sources:
            try:
                (source,) = (
                    s.split(substr, maxsplit=1)[1]
                    # HTTP due to https://github.com/Future-House/LAB-Bench/issues/11
                    for substr in {*DocDetails.DOI_URL_FORMATS, "http://doi.org/"}
                    if substr in s
                )
            except ValueError as exc:
                if not isinstance(s, str):
                    raise NotImplementedError(
                        f"Didn't handle source extraction from raw source {s!r}."
                    ) from exc
                source = s  # Not a DOI
            sources.append(source)
        return sources


def read_ds_from_hub(
    hf_name: str | LABBenchDatasets | None,
    hf_path: str,
    hf_split: str,
    randomize: bool = True,
    seed: int | None = None,
) -> "pd.DataFrame":
    """
    Read in a train or test DataFrame.

    Args:
        hf_name: Hugging Face Hub dataset's name.
            E.g. "LitQA2" for LitQA v2 or None for FigQA v2.
        hf_path: Hugging Face Hub dataset's path.
        hf_split: Hugging Face Hub dataset's split, e.g. "train" or "test".
        randomize: Opt-out flag to shuffle the dataset after loading in by question.
        seed: Random seed to use for the shuffling.

    Raises:
        DatasetNotFoundError: If any of the datasets are not found, or the
            user is unauthenticated.
    """  # noqa: DOC502
    try:
        from datasets import load_dataset
    except ImportError as exc:
        raise ImportError(
            "Reading in Hugging Face datasets requires the 'datasets' extra for"
            " 'datasets'. Please: `pip install aviary.labbench[datasets]`."
        ) from exc

    if isinstance(hf_name, LABBenchDatasets):
        hf_name = hf_name.value
    ds = load_dataset(hf_path, hf_name, split=hf_split).to_pandas()
    if "distractors" in ds.columns:
        # Convert to list so it's not unexpectedly a numpy array
        ds["distractors"] = ds["distractors"].apply(list)
    # Let downstream usage in the TaskDataset's environment factories check for the
    # presence of other DataFrame columns
    if randomize:
        ds = ds.sample(frac=1, random_state=seed)
    return ds


class PaperQATaskDataset(TaskDataset[TGradableEnv], ComputeTrajectoryMetricsMixin, ABC):
    """
    Abstract base class for a task dataset of gradable questions for PaperQA.

    This is an ABC because it's non-specific to a given task.
    Examples include LitQA v1, v2, FigQA, TableQA, or a test stub version of LitQA.
    """

    def __init__(
        self,
        settings: Settings | dict | None = None,
        base_docs: Docs | dict | None = None,
        rewards: Mapping[str, float] = DEFAULT_REWARD_MAPPING,
        question_kwargs: Mapping[str, Any] | None = None,
        eval_model: LLMModel | str = DEFAULT_EVAL_MODEL_NAME,
        **env_kwargs,
    ):
        if settings is None:
            settings = Settings()
        if isinstance(settings, dict):
            settings = Settings(**settings)
        self._settings = settings
        if base_docs is None:
            base_docs = Docs()
        if isinstance(base_docs, dict):
            base_docs = Docs(**base_docs)
        self._base_docs = base_docs
        self._rewards = rewards
        self._question_kwargs = question_kwargs
        self._eval_model = eval_model
        self._env_kwargs = env_kwargs

    def compute_trajectory_metrics(
        self, trajectories: "Sequence[Trajectory]"
    ) -> dict[str, list[float]]:
        total_paper_count: list[float] = []
        relevant_paper_count: list[float] = []
        evidence_count: list[float] = []
        for t in trajectories:
            split_certainties = [
                split_certainty
                for split_certainty in (
                    re.split(
                        pattern=Complete.CERTAINTY_SPLIT_REGEX_PATTERN,
                        string=obs.content,
                        maxsplit=1,
                    )
                    for obs in t.steps[-1].next_observation
                    if (
                        isinstance(obs, ToolResponseMessage)
                        and obs.name == Complete.TOOL_FN_NAME
                    )
                )
                # Filter for places where the regex split succeeded
                if len(split_certainty) >= 4  # noqa: PLR2004
            ]
            for i, metric_list in enumerate(
                (total_paper_count, relevant_paper_count, evidence_count),
                start=1,  # Regex extraction of status starts after has_successful_answer
            ):
                # NOTE: we use mean to not break if there's 2+ complete calls (which
                # we're prompted not to do). If it happens, they should all have the
                # same status, so the mean value should equal the individual values
                metric_list.append(
                    sum(int(sa[i]) for sa in split_certainties) / len(split_certainties)
                    if split_certainties  # Avoid div0 (when complete wasn't called)
                    else 0
                )
        return super().compute_trajectory_metrics(trajectories) | {
            "total_paper_count": total_paper_count,
            "relevant_paper_count": relevant_paper_count,
            "evidence_count": evidence_count,
            "correct": [
                int(t.steps[-1].reward == self._rewards["correct"])
                for t in trajectories
            ],
            "correct_unsure": [
                int(
                    t.steps[-1].reward
                    in {self._rewards["correct"], self._rewards["unsure"]}
                )
                for t in trajectories
            ],
        }


@unique
class TextQATaskSplit(StrEnum):
    TRAIN = "train"
    TEST = "test"

    def get_index(self) -> int:
        """Get the index of the train (0) or test (1) split."""
        if self == self.TRAIN:
            return 0
        if self == self.TEST:
            return 1
        assert_never(self)


class TextQATaskDataset(PaperQATaskDataset[TGradableEnv]):
    """Dataset for LAB-Bench datasets compatible with text-based QA."""

    def __init__(
        self,
        *args,
        dataset: str | LABBenchDatasets = LABBenchDatasets.LIT_QA2,
        read_data_kwargs: Mapping[str, Any] | None = None,
        split: str | TextQATaskSplit = TextQATaskSplit.TRAIN,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self._dataset = LABBenchDatasets(dataset)
        self.data = self._dataset.get_data(split, **(read_data_kwargs or {}))

    def _make_query(self, idx: int) -> MultipleChoiceQuestion:
        distractors = (
            self.data.iloc[idx].distractors
            if "distractors" in self.data.columns
            else []
        )
        return MultipleChoiceQuestion(
            question_id=UUID(self.data.iloc[idx].id),
            question=self.data.iloc[idx].question,
            options=(
                distractors
                if isinstance(distractors, list)
                else MultipleChoiceQuestion.split_options(distractors)
            ),
            ideal_answer=self.data.iloc[idx].ideal,
            prompt_without_id=True,
            prompt_without_options=self._dataset.is_open_answer,
            **(self._question_kwargs or {}),
        )

    def _make_sources(self, idx: int) -> str | list[str] | None:
        return self._dataset.get_sources(self.data.iloc[idx])

    def get_new_env_by_idx(self, idx: int) -> GradablePaperQAEnvironment:  # type: ignore[override]
        mcq = self._make_query(idx)
        return GradablePaperQAEnvironment(
            query=mcq,
            settings=self._settings,
            docs=self._base_docs.model_copy(),
            sources=self._make_sources(idx),
            rewards=self._rewards,
            session_id=cast(UUID, mcq.question_id),  # Expedite manual inspection
            **self._env_kwargs,
        )

    def __len__(self) -> int:
        return len(self.data)

    async def get_images(
        self, value: str | UUID | MultipleChoiceQuestion | GradablePaperQAEnvironment
    ) -> bytes | list[bytes]:
        """
        Get the image(s) used in the environment, helpful for recall measurement.

        NOTE: FigQA has 1 image with paths, TableQA has 1+ images with paths.
        """
        if isinstance(value, GradablePaperQAEnvironment):
            question_id: str | UUID = await value.get_id()
        elif isinstance(value, MultipleChoiceQuestion):
            question_id = value.question_id
        else:
            question_id = value

        matching_rows = self.data[self.data["id"] == str(question_id)]
        if len(matching_rows) != 1:
            raise ValueError(
                f"Question ID {question_id} matched {len(matching_rows)} rows,"
                " require exactly 1 match."
            )

        try:
            images_col = self._dataset.images_column
        except ValueError as exc:
            raise ValueError(
                f"Dataset {self._dataset.value!r} has no images column."
            ) from exc

        images = matching_rows.iloc[0][images_col]
        return (
            images["bytes"]
            if isinstance(images, dict)
            else [image["bytes"] for image in images]
        )


for dataset_name in ("figqa-text", "litqa2", "tableqa-text"):
    TASK_DATASET_REGISTRY[dataset_name] = (
        TextQATaskDataset.__module__,
        TextQATaskDataset.__name__,
    )


class ImageQATaskDataset(TextQATaskDataset[ImageQAEnvironment]):
    """Dataset for LAB-Bench datasets compatible with image-based QA."""

    def __init__(
        self,
        settings: Settings | dict | None = None,
        dataset: str | LABBenchDatasets = LABBenchDatasets.FIG_QA,
        autogenerate_settings: bool = True,
        **kwargs,
    ):
        if autogenerate_settings and settings is None:
            settings = ImageQAEnvironment.make_base_settings()
        super().__init__(settings=settings, dataset=dataset, **kwargs)
        if self._dataset in {LABBenchDatasets.FIG_QA2, LABBenchDatasets.TABLE_QA2}:
            raise ValueError(
                f"Dataset {self._dataset.value!r} not supported by {type(self).__name__}"
                " because it has no image(s) stored in the dataset."
            )

    def get_new_env_by_idx(self, idx: int) -> ImageQAEnvironment:
        mcq = self._make_query(idx)
        images = self.data.iloc[idx][self._dataset.images_column]
        image_paths = self.data.iloc[idx][self._dataset.paths_column]
        return ImageQAEnvironment(
            query=mcq,
            settings=self._settings,
            docs=self._base_docs.model_copy(),
            sources=self._dataset.get_sources(self.data.iloc[idx]),
            rewards=self._rewards,
            session_id=cast(UUID, mcq.question_id),  # Expedite manual inspection
            images=(
                images["bytes"]
                if isinstance(images, dict)
                else [image["bytes"] for image in images]
            ),
            image_paths=(
                image_paths if isinstance(image_paths, str) else list(image_paths)
            ),
            **self._env_kwargs,
        )


for dataset_name in ("figqa", "tableqa", "figqa2", "tableqa2"):
    TASK_DATASET_REGISTRY[dataset_name] = (
        ImageQATaskDataset.__module__,
        ImageQATaskDataset.__name__,
    )
