import json
import logging
import os
import re
from typing import Any, Dict, List

import datasets
from huggingface_hub import hf_hub_download
from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.model import Model, ResponseSchema, get_model
from inspect_ai.scorer import (
    Score,
    Scorer,
    ScoreReducer,
    Target,
    mean,
    multi_scorer,
    score_reducer,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState, use_tools
from inspect_ai.util import json_schema
from pydantic import BaseModel

from astabench.constants import ASTA_BENCH_DATASET_REPO, ASTA_BENCH_DATASET_REVISION
from astabench.evals.sqa.citation_eval import JUST_HAS_A_TITLE, CitationEval
from astabench.evals.sqa.precision_eval import PrecisionEval
from astabench.evals.sqa.rubric import (
    RubricCorpusQaGenericMetric,
    extract_json_from_response,
)
from astabench.evals.sqa.split_sentences import add_missing_citations
from astabench.evals.utils import not_implemented_solver
from astabench.tools import (
    make_asta_s2_mcp_tools,
    make_native_search_tools,
    report_editor,
)
from astabench.types.sqa import SQAResponseWithTable

logger = logging.getLogger(__name__)

any_json_schema = ResponseSchema(
    name="json",
    json_schema=json_schema({"type": "object", "additionalProperties": True}),
)


class ResultStringList(BaseModel):
    result: List[str]


result_string_list_schema = ResponseSchema(
    name="result_string_list",
    json_schema=json_schema(ResultStringList),
)


class CitationGroups(BaseModel):
    texts: List[str]
    inline_citations: List[List[str]]


citation_group_schema = ResponseSchema(
    name="citation_group", json_schema=json_schema(CitationGroups)
)

METRIC_LABELS = [
    "ingredient_recall",
    "answer_precision",
    "citation_precision",
    "citation_recall",
]


def json_to_sample(obj: dict[str, str], excerpt_prompt: bool = True) -> Sample:
    task_instructions = "Generate a report answering the following research question. Be sure to include inline citations for each claim. Your response should be divided into sections, each followed by a References section. {excerpt_prompt}\n\nQuestion:"
    if excerpt_prompt:
        task_instructions = task_instructions.format(
            excerpt_prompt="Each entry in the References section should quote the excerpts used from the cited source."
        )
        prompt = """Generate a report answering the following research question. Be sure to include inline citations for each claim. Return your result as valid JSON with a single key `sections` which is a list of sections, each having keys `title`, `text`, and `citations`. Each entry in `citations` should have a JSON list of `snippets` extracted from the reference document and an `id`, each of which appears exactly in the text. Each `id` should be an inline citation as it appears in the text (with wrapping parentheses or square brackets if appropriate). Each citation should have a `title` if one is available. Any additional information about the citation should go under `metadata`. Do not create a References section.\n\nHere is an example `section` to help you with formatting:

        {
          "title": "Background",
          "text": "Convolutional neural networks (CNNs) have achieved state-of-the-art results in image classification [1][2].",
          "citations": [
            {
              "id": "[1]",
              "snippets": [
                "CNNs have become the standard for many visual tasks."
              ],
              "title": "ImageNet Classification with Deep Convolutional Neural Networks",
              "metadata": {
                "authors": "Krizhevsky, A. et al.",
                "year": 2012,
                "arxiv": "1207.0580"
              }
            },
            {
              "id": "[2]",
              "snippets": [
                "Significant improvements in image recognition have been observed with CNNs."
              ],
              "title": "Very Deep Convolutional Networks for Large-Scale Image Recognition",
              "metadata": {
                "authors": "Simonyan, K. & Zisserman, A.",
                "year": 2014,
                "arxiv": "1409.1556"
              }
            }
          ]
        }

        Question: """
    else:
        task_instructions = task_instructions.format(
            excerpt_prompt="Each entry in the References section should ONLY have the title of the cited source."
        )
        prompt = task_instructions
    prompt_without_formatting_instructions = task_instructions + obj["question"]

    return Sample(
        id=obj["case_id"],
        input=prompt + obj["question"],
        target=json.dumps(obj, indent=2),
        metadata={
            "annotator": obj["annotator"],
            "agreement": obj.get("agreement", False),
            "initial_prompt": obj["question"],
            "case_id": obj["case_id"],
            "prompt_without_formatting_instructions": prompt_without_formatting_instructions,
        },
    )


def format_citation(cit):
    snippets = [
        " ".join(snippet.split()) for snippet in cit.snippets if snippet.strip()
    ]
    return f"{cit.id}:\n{'. '.join(snippets)}"


def format_report(response):
    response = SQAResponseWithTable(**response)
    report = ""
    for section in response.sections:
        if section.text:
            report += f"{section.text}\n\n"
        if section.table:
            report += f"{section.table.text}\n\n"
    return report.strip()


async def sent_tokenize(scorer_model, text: str) -> List[str]:
    return await add_missing_citations(scorer_model, text)


def normalize_agent_response_dict(response_dict: dict[str, Any]) -> dict[str, Any]:
    """We allow `{"response": {"sections": ...}}` or `{"sections": ...}`, so
    this function is for normalizing that format to always be `{"sections":
    ...}`"""
    if not response_dict.get("sections") and response_dict.get("response"):
        return response_dict["response"]
    return response_dict


@scorer(metrics={"score": [mean(), stderr()], "ingredient_recall": [mean(), stderr()]})
def score_sqa(
    model: str | Model,
    simplified_eval: bool = False,
    assess_jointly: bool = False,
    temperature: float = 0.5,
    top_p: float = 0.95,
) -> Scorer:
    grader_model = get_model(model)

    async def score(state: TaskState, target: Target) -> Score:
        raw_output = state.output.completion
        parsed_output = extract_json_from_response(raw_output)
        if simplified_eval:
            default_value = {"ingredient_recall": 0.0}
        else:
            default_value = {"score": 0.0, "ingredient_recall": 0.0}
        default_output = Score(value=default_value, answer=raw_output)
        if (
            not parsed_output
            or normalize_agent_response_dict(parsed_output).get("sections") is None
        ):
            # you get 0 if you cant return proper json containing a response
            return default_output

        parsed_output = normalize_agent_response_dict(parsed_output)

        answer = format_report(parsed_output)
        metric_config = json.loads(target[0])
        metric = RubricCorpusQaGenericMetric(
            metric_config, grader_model, temperature=temperature, top_p=top_p
        )
        if simplified_eval:
            scores, metadata = await metric.score_output_simplified(
                answer, assess_jointly
            )
        else:
            scores, metadata = await metric.score_output(answer)

        return Score(value=scores, answer=raw_output, metadata=metadata)

    return score


def clean_sentence(sentence: str) -> str:
    pattern = r'<Paper [^>]*paperTitle="\W*([ _a-zA-Z0-9,.;]+)\W*"[^>]*/?>\s*</Paper>'
    sentence = re.sub(pattern, r"(\1)", sentence)
    pattern = r'<Model name="Anthropic" version="[^"]+">'
    return re.sub(pattern, "", sentence)


@scorer(
    metrics={
        "citation_precision": [mean(), stderr()],
        "citation_recall": [mean(), stderr()],
        "f1": [mean(), stderr()],
    }
)
def score_citation(
    model: str | Model,
    is_retrieverless: bool = False,
    sentence_wise_cit_eval: bool = True,
    all_at_once: bool = False,
    temperature: float = 0.5,
    top_p: float = 0.95,
) -> Scorer:
    grader_model = get_model(model)

    async def score(state: TaskState, target: Target) -> Score:
        def filter_citation(citation: Dict[str, Any], cite_text: str) -> bool:
            """A check to determine if citation snippets are present and usable.
            If the snippets are:
            1. present, and
            2. are not exactly present in the section text, and
            3. either there is not title OR if a title is present, are not same as the title

            then return True.
            """
            cite_text_alpha = re.sub(r"[^a-zA-Z]", "", cite_text).lower()
            snippets_alpha = [
                re.sub(r"[^a-zA-Z]", "", s).lower()
                for s in citation.get("snippets", [])
            ]
            return (
                citation.get("snippets")
                and (not any(s in cite_text_alpha for s in snippets_alpha))
                and (
                    not (
                        citation.get("title")
                        and any(
                            re.sub(r"[^a-zA-Z]", "", citation["title"]).lower() == s
                            for s in snippets_alpha
                        )
                    )
                )
            )

        """
        The output format for the task needs to be fixed for the citation eval
        This implementation expects a json with the response text as a list of sections
        and each section having a corresponding list of citation with snippet excerpts, where
        the id of the citation excerpt should be present in the response text.
        eg. {"sections": [{"text": <response text with citations>, "citations": [{"id": <citation_id>, "snippets": <citation_snippets>}]}]}
        """
        raw_output = state.output.completion
        json_output = extract_json_from_response(raw_output)
        default_output = Score(
            value={"citation_recall": 0.0, "citation_precision": 0.0, "f1": 0.0},
            answer=raw_output,
        )
        if not json_output:
            # you get 0 if you cant return proper json
            return default_output
        bad_snippet = (
            "Please click on the paper title to read the abstract on Semantic Scholar."
        )
        if sentence_wise_cit_eval:
            metric = CitationEval(grader_model, temperature=temperature, top_p=top_p)
            claim_evidence_pairs = []
            json_output = normalize_agent_response_dict(json_output)
            for section in json_output["sections"]:
                sec_text = section["text"]
                citations = [
                    {
                        "id": c["id"],
                        "snippets": (
                            "... ".join(c["snippets"])
                            if isinstance(c["snippets"], list)
                            else c["snippets"]
                        ),
                    }
                    for c in section["citations"]
                    if bad_snippet not in c["snippets"]
                ]
                clean_text = clean_sentence(sec_text)
                sec_sentences = await sent_tokenize(model, clean_text)
                for sent in sec_sentences:
                    sent_citations = []
                    sent_cit_ids = []
                    for citation in citations:
                        if citation["id"] in sent:
                            sent_citations.append(citation["snippets"])
                            sent_cit_ids.append(citation["id"])
                    claim_evidence_pairs.append((sent, sent_citations, sent_cit_ids))
            scores = await metric.evaluate(claim_evidence_pairs)
            return Score(value=dict(scores), answer=raw_output)
        elif all_at_once:
            metric = CitationEval(grader_model, temperature=temperature, top_p=top_p)
            citation_group_cit_pairs = []
            json_output = normalize_agent_response_dict(json_output)
            for section in json_output["sections"]:
                sec_iter = (
                    [section]
                    if not section.get("table")
                    else [section, section["table"]]
                )
                # If the section has a table, we also include the table in the evaluation
                for curr_sec in sec_iter:
                    sec_text = curr_sec["text"]
                    citations = [
                        (
                            {
                                "id": c["id"],
                                "snippets": (
                                    "... ".join(c["snippets"])
                                    if isinstance(c["snippets"], list)
                                    else c["snippets"]
                                ),
                            }
                            if (
                                filter_citation(c, sec_text)
                                and bad_snippet not in c["snippets"]
                            )
                            else {  # if the citation has a title and is missing snippets:
                                "id": c["id"],
                                "snippets": "{}{}".format(JUST_HAS_A_TITLE, c["title"]),
                            }
                        )
                        for c in curr_sec["citations"]
                    ]
                    clean_text = clean_sentence(sec_text)
                    citation_group_cit_pairs.append((clean_text, citations))
            scores, metadata = await metric.evaluate_citation_groups(
                citation_group_cit_pairs, is_retrieverless=is_retrieverless
            )
            return Score(value=dict(scores), answer=raw_output, metadata=metadata)
        else:
            raise ValueError("The citation evaluation mode is not supported")

    return score


@scorer(metrics={"answer_precision": [mean(), stderr()]})
def score_precision(model: str | Model) -> Scorer:
    grader_model = get_model(model)

    async def score(state: TaskState, target: Target) -> Score:
        raw_output = state.output.completion
        parsed_output = extract_json_from_response(raw_output)
        default_output = Score(value={"answer_precision": 0.0}, answer=raw_output)
        if (
            not parsed_output
            or normalize_agent_response_dict(parsed_output).get("sections") is None
        ):
            # you get 0 if you cant return proper json containing a report
            return default_output
        parsed_output = normalize_agent_response_dict(parsed_output)
        answer = format_report(parsed_output)
        query = state.metadata["initial_prompt"]
        metric = PrecisionEval(grader_model)
        scores, meta = await metric.score_output(query, answer)

        return Score(
            value=scores,
            answer=raw_output,
            metadata=meta,
        )

    return score


@score_reducer(name="weighted_mean")
def weighted_mean_score(weights: Dict[str, float]) -> ScoreReducer:
    def reduce(scores: list[Score]) -> Score:
        global_avg = 0.0
        scores_map = {k: v for score in scores for k, v in score.value.items()}
        for mkey in METRIC_LABELS:
            global_avg += scores_map[mkey] * weights[mkey]
        scores_map["global_avg"] = global_avg
        return Score(
            value=scores_map,
            metadata={"score_meta": [score.metadata for score in scores]},
        )

    return reduce


@scorer(metrics={k: [mean(), stderr()] for k in (["global_avg"] + METRIC_LABELS)})
def score_all(
    scorer_model: str | Model = "google/gemini-2.5-flash",
    is_retrieverless: bool = False,
    simplified_eval: bool = True,
    assess_jointly: bool = True,
    sentence_wise_cit_eval: bool = False,
    all_at_once: bool = True,
    temperature: float = 0.5,
    top_p: float = 0.95,
) -> Scorer:
    scorers = [
        score_sqa(
            scorer_model,
            simplified_eval=simplified_eval,
            assess_jointly=assess_jointly,
            temperature=temperature,
            top_p=top_p,
        ),
        score_precision(scorer_model),
        score_citation(
            scorer_model,
            sentence_wise_cit_eval=sentence_wise_cit_eval,
            is_retrieverless=is_retrieverless,
            all_at_once=all_at_once,
            temperature=temperature,
            top_p=top_p,
        ),
    ]
    return multi_scorer(
        scorers,
        reducer=weighted_mean_score({k: 1 / len(METRIC_LABELS) for k in METRIC_LABELS}),
    )


@scorer(metrics={k: [mean(), stderr()] for k in (["global_avg"] + METRIC_LABELS)})
def aggregate_scorer(scorers: List[Scorer], weights: Dict[str, float]) -> Scorer:
    return multi_scorer(scorers, reducer=weighted_mean_score(weights))


def load_rubrics_data(split: str = "dev") -> List[Dict]:
    if split == "test":
        version = "v2"
    elif split == "dev":
        version = "v1"
    else:
        raise ValueError("Invalid split. Use 'dev' or 'test'")
    data_files = f"tasks/sqa/rubrics_{version}_recomputed.json"
    p = hf_hub_download(
        ASTA_BENCH_DATASET_REPO,
        data_files,
        token=os.getenv("HF_TOKEN"),
        repo_type="dataset",
        revision=ASTA_BENCH_DATASET_REVISION,
    )
    dataset = datasets.load_dataset(
        "json",
        data_files=p,
    )
    return list(map(dict, dataset["train"]))


@task
def sqa(
    scorer_model: str | Model = "google/gemini-2.5-flash",
    limit: int | None = None,
    is_retrieverless: bool = False,
    simplified_eval: bool = True,
    assess_jointly: bool = True,
    sentence_wise_cit_eval: bool = False,
    all_at_once: bool = True,
    temperature: float = 0.5,
    top_p: float = 0.95,
    with_search_tools: bool = True,
    with_report_editor_tool: bool = False,
    with_native_search_tools: bool = False,
    split: str = "dev",
    weights=None,
    aggregate: bool = True,
    excerpt_prompt: bool = True,
) -> Task:
    if weights is None:
        weights = [0.25, 0.25, 0.25, 0.25]
    rubrics_ds = load_rubrics_data(split)
    json_dataset = MemoryDataset(
        samples=[json_to_sample(item, excerpt_prompt) for item in rubrics_ds[:limit]],
    )
    tool_setups = []
    inserted_before = None
    # See https://github.com/allenai/nora-issues-research/issues/39#issuecomment-2881618717
    if split in {"dev", "test"}:
        inserted_before = "2025-05"
    else:
        logger.warning(
            f"Unknown split {split}. No inserted_before date will be set for search tools."
        )
    if with_search_tools:
        tool_setups.append(
            use_tools(make_asta_s2_mcp_tools(insertion_date=inserted_before))
        )
    if with_native_search_tools:
        native_search_tools = make_native_search_tools(inserted_before=inserted_before)
        if native_search_tools:
            tool_setups.append(use_tools(native_search_tools))
    if with_report_editor_tool:
        tool_setups.append(use_tools(report_editor()))
    scorer_list = [
        score_sqa(
            scorer_model,
            simplified_eval=simplified_eval,
            assess_jointly=assess_jointly,
            temperature=temperature,
            top_p=top_p,
        ),
        score_precision(scorer_model),
        score_citation(
            scorer_model,
            sentence_wise_cit_eval=sentence_wise_cit_eval,
            is_retrieverless=is_retrieverless,
            all_at_once=all_at_once,
            temperature=temperature,
            top_p=top_p,
        ),
    ]
    weights_map = {k: v for k, v in zip(METRIC_LABELS, weights)}

    return Task(
        dataset=json_dataset,
        setup=tool_setups or None,
        # Solver that just tells the user to specify a real solver via cli
        solver=not_implemented_solver(),
        scorer=aggregate_scorer(scorer_list, weights_map) if aggregate else scorer_list,
        fail_on_error=False,
    )


@task
def sqa_dev() -> Task:
    return sqa(split="dev")


@task
def sqa_test() -> Task:
    return sqa(split="test")
