import asyncio
import json
import logging
from enum import StrEnum
from typing import Any, Literal, Sequence, TypedDict

from inspect_ai.model import get_model
from pydantic import BaseModel, ConfigDict

from astabench.evals.paper_finder.datamodel import CorpusId
from astabench.evals.utils import extract_json_from_response

logger = logging.getLogger(__name__)


"""
###################################################### types ###########################################################
"""

DEFAULT_CONTENT_RELEVANCE_CRITERION_NAME = "Relevance Criterion"
grader_model = get_model("openai/gpt-4o-2024-11-20")


class Relevance(StrEnum):
    PERFECT = "perfectly_relevant_papers"
    HIGHLY = "highly_relevant_papers"
    SOMEWHAT = "somewhat_relevant_papers"
    NOT_RELEVANT = "not_relevant_papers"


class Document(BaseModel):
    corpus_id: str
    markdown: str


class RelevanceJudgement(BaseModel):
    relevance: int
    corpus_id: str


class RelevanceCriterion(BaseModel):
    name: str
    description: str
    weight: float

    model_config = ConfigDict(frozen=True)


pf_rj_2_bench_rj = {
    0: Relevance.NOT_RELEVANT.value,
    1: Relevance.SOMEWHAT.value,
    2: Relevance.HIGHLY.value,
    3: Relevance.PERFECT.value,
}
bench_rj_2_pf_rj = {v: k for k, v in pf_rj_2_bench_rj.items()}

"""
##################################################### prompts ##########################################################
"""

relevance_criteria_judgement_prompt_with_relevant_snippets_after = """
Judge how relevant the following paper is to each of the provided criteria. For each criterion, consider its entire description when making your judgement.

For each criterion, provide the following outputs:
- `relevance` (str): one of "Perfectly Relevant", "Somewhat Relevant", "Not Relevant".
- `relevant_snippet` (str | null): a snippet from the document that best show the relevance of the paper to the criterion. To be clear, copy EXACT text ONLY. Choose one short text span that best shows the relevance in a concrete and specific way, up to 20 words. ONLY IF NECESSARY, you can add another few-words-long span (e.g. for coreference, disambiguation, necessary context), separated by ` ... `. If relevance is "Not Relevant" output null. The snippet may contain citations, but make sure to only take snippets that directly show the relevance of this paper. 

Also provide another field `relevance_summary` (str): a short summary explanation of how the paper is relevant to the criteria in general. null if it is not relevant to any of the criteria.
- This should be short but convey the most useful and specific information for a user skimming a list of papers, up to 30 words.
- No need to mention which are perfectly relevant, somewhat relevant, or not relevant. Just provide new information that was not mentioned in the criteria.
- Start with perfectly relevant ones, include a specific and interesting detail about what matches them. Then go on to somewhat relevant ones, saying why it is close but not a perfect match. No need to add extra info for not relevant ones.
- Start with actionable info. Instead of saying "The paper uses X to...", just say "Uses X to...".

Output a JSON:
- top-level key `criteria`. Under it, for every criterion name (exactly as given in the provided criteria), there should be an object containing two fields: `relevance` and `relevant_snippet`.
- top-level key `relevance_summary` with string value or null.

Criteria:
```
{criteria}
```"""  # noqa: E501


class RelevanceJudgementInput(TypedDict):
    doc_id: str
    criteria: str
    document: str


class RelevanceCriterionJudgementValue(BaseModel):
    relevant_snippet: str | None
    relevance: Literal["Perfectly Relevant", "Somewhat Relevant", "Not Relevant"]


class RelevanceJudgementOutput(BaseModel):
    criteria: dict[str, RelevanceCriterionJudgementValue]
    relevance_summary: str | None


ModelName = str


rj_4l_codes = {
    "Perfectly Relevant": 3,
    "Highly Relevant": 2,
    "Somewhat Relevant": 1,
    "Not Relevant": 0,
}

"""
#################################################### functions #########################################################
"""


async def load_relevance_judgement(
    entities: Sequence[Document],
    relevance_criteria: list[RelevanceCriterion] | None = None,
) -> dict[CorpusId, str]:
    relevance_judgement_results = await _llm_judge_relevance_with_criteria(
        entities, relevance_criteria
    )

    id2relevance = _parse_relevance_judgement_results(relevance_judgement_results)

    relevance_judgements = _build_relevance_judgement_results(
        entities, id2relevance, relevance_criteria
    )

    translate_to_eval_scores = {
        cid: pf_rj_2_bench_rj[pf_rj] for cid, pf_rj in relevance_judgements.items()
    }

    return translate_to_eval_scores


def _format_documents_for_relevance_judgement(
    documents: Sequence[Document], criteria_to_use: list[RelevanceCriterion]
) -> list[RelevanceJudgementInput]:
    formatted_criteria = json.dumps(
        [{"name": r.name, "description": r.description} for r in criteria_to_use],
        indent=2,
    )
    doc_texts: list[RelevanceJudgementInput] = [
        RelevanceJudgementInput(
            document=doc.markdown,
            criteria=formatted_criteria,
            doc_id=doc.corpus_id,
        )
        for doc in documents
    ]
    return doc_texts


async def _llm_judge_relevance_with_criteria(
    documents: Sequence[Document],
    relevance_criteria: list[RelevanceCriterion] | None,
) -> dict[CorpusId, RelevanceJudgementOutput]:
    doc_texts = _format_documents_for_relevance_judgement(documents, relevance_criteria)
    if len(doc_texts) == 0:
        logger.info("No documents to run relevance judgement on.")
        return dict()
    logger.info(f"Running relevance judgement for {len(doc_texts)} documents")

    judging_futures = []
    for doc_text in doc_texts:
        score_prompt = (
            relevance_criteria_judgement_prompt_with_relevant_snippets_after.format(
                criteria=doc_text
            )
        )
        judging_futures.append(grader_model.generate(score_prompt))

    relevance_judgement_results = dict()
    for raw, doc_id in zip(
        await asyncio.gather(*judging_futures), [d.corpus_id for d in documents]
    ):
        try:
            relevance_judgement_results[doc_id] = RelevanceJudgementOutput(
                **extract_json_from_response(raw.completion)
            )
        except (TypeError, ValueError):
            # lets skip a single doc's failure..
            continue

    return relevance_judgement_results


def _parse_relevance_judgement_results(
    relevance_judgement_results: dict[CorpusId, RelevanceJudgementOutput],
) -> dict[str, list[dict[str, Any]]]:
    id2relevance: dict[str, list[dict[str, Any]]] = dict()
    for doc_id, relevance_judgement in relevance_judgement_results.items():
        doc_relevance_criterion_judgements = [
            {
                # NOTE: removing underscore as it's a common llm error for naming fields
                "name": criterion_name.replace("_", " "),
                "relevance": rj_4l_codes[judgement.relevance],
                "relevant_snippet": judgement.relevant_snippet,
            }
            for criterion_name, judgement in relevance_judgement.criteria.items()
        ]
        id2relevance[doc_id] = doc_relevance_criterion_judgements

    return id2relevance


def _build_relevance_judgement_results(
    entities: Sequence[Document],
    id2relevance: dict[str, list[dict[str, Any]]],
    relevance_criteria: list[RelevanceCriterion] | None = None,
) -> dict[CorpusId, int]:
    results = dict()

    for doc in entities:
        rcj_result = id2relevance.get(doc.corpus_id)

        if relevance_criteria is None or rcj_result is None:
            logger.warning(
                f"Relevance criteria and/or judgements not found for document {doc.corpus_id}. Skipping."
            )
            continue

        # check that all required criteria names are present in rcj
        required_criteria_names = [criterion.name for criterion in relevance_criteria]
        judgement_criteria_names = [criterion["name"] for criterion in rcj_result]
        if any(
            name not in judgement_criteria_names for name in required_criteria_names
        ):
            logger.warning(
                f"Required relevance criteria not found for document {doc.corpus_id}. Skipping. Required criteria: {required_criteria_names}. Judged criteria: {judgement_criteria_names}"
            )
            continue

        try:
            relevance_criteria_score = calculate_relevance_criteria_score(
                relevance_criteria, rcj_result
            )
        except Exception:
            logger.exception(
                f"Failed to calculate relevance criteria score for document {doc.corpus_id}. Skipping."
            )
            continue
        relevance_0_to_3 = calculate_0_to_3_relevance(relevance_criteria_score)
        results[doc.corpus_id] = relevance_0_to_3
    return results


def calculate_0_to_3_relevance(relevance_criteria_score: float) -> int:
    if relevance_criteria_score <= 0.25:
        return 0
    if relevance_criteria_score <= 0.67:
        return 1
    if relevance_criteria_score <= 0.99:
        return 2
    return 3


def calculate_relevance_criteria_score(
    relevance_criteria: list[RelevanceCriterion], judgements: list[dict[str, Any]]
) -> float:
    criterion_name_to_weight = {}
    for criteria in relevance_criteria:
        criterion_name_to_weight[criteria.name] = criteria.weight

    score = 0
    for judgement in judgements:
        # relevance is 0-3; divide by 3 to normalize to a 0-1 score
        score += (
            criterion_name_to_weight[judgement["name"]] * judgement["relevance"] / 3
        )

    score = min(1.0, score)

    return score
