from typing import Dict, List, Any
from datetime import datetime

from inspect_ai.model import Model, GenerateConfig

from astabench.evals.sqa.pydantic_models import answer_precision_json_schema
from astabench.evals.sqa.rubric import extract_json_from_response
from astabench.evals.sqa.retry_utils import generate_with_retry

import re
import difflib


def _create_prompt_logs(
    prompt: str,
    query: str,
    answer: str,
    model_output: str,
    temperature: float,
    metadata: Dict[str, Any],
) -> Dict[str, Any]:
    prompt_id = str(hash(prompt))
    data = {
        "step_name": "precision_eval",
        "query": query,
        "answer": answer,
        "prompt": prompt,
        "temperature": temperature,
        "top_p": 0.95,
        "timestamp": str(datetime.now()),
        **metadata,
    }
    return {k: [{"prompt_id": prompt_id, "data": {k: v}}] for k, v in data.items()}


PRECICION_EVAL_PROMPT = """You are given a query and a corresponding long answer.

Goal: find irrelevant paragraphs in the answer. These are paragraphs that don't directly answer the query and shouldn't be in the answer.

For instance, if the query is about datasets for scientific question answering, a paragraph about multilingual question answering datasets that don't contain scientific text would be considered irrelevant.

Explicitly consider whether something may be indirectly relevant. For example, if the question is about the conditions of horses in South Africa, a paragraph about general animal welfare in South Africa is potentially relevant while not being precisely about horses. On the other hand, a paragraph about pig welfare in South Africa is irrelevant.

Note that subtle differences can make the text irrelevant to the query. For instance, text about scientific survey paper generation is not relevant to a query about automatic paper review generation. Even though they seem related, they are about very different tasks.

Also, useful background in general is relevant. If the question is about an approach to creating liver-related proteins, some information about liver-related proteins could contextualize other parts of the answer. If a paragraph contextualizes another part of the answer, then it is relevant.

Go through the answer and output a list of irrelevant paragraphs. Every single paragraph needs to be considered, one by one. Our goal is to catch all the irrelevant paragraphs, so please be thorough.
Return your result as a JSON object with a single key `irrelevant_paragraphs` whose value is a list of objects, each having keys `reason`, and `answer_text` as follows:
{{"irrelevant_paragraphs":[
{{
"reason": "discuss why something is irrelevant (not indirectly relevant)",
"answer_text": "exact ENTIRE paragraph (not just a part of it) from the answer that is irrelevant"
}},
...
]
}}
Make sure all the irrelevant paragraphs are included.

Question: {query}
Answer: {answer}
"""


class PrecisionEval:
    def __init__(self, model: Model, temperature: float = 0.5):
        self.model = model
        self.temperature = temperature

    @staticmethod
    def map_output_to_precision_score(
        parsed_json: Dict[str, Any], answer: str
    ) -> float:
        metadata = {}
        irrelevant_texts = [
            json_list["answer_text"]
            for json_list in parsed_json["irrelevant_paragraphs"]
        ]
        metadata["irrelevant_texts"] = irrelevant_texts

        # Split answer into paragraphs
        paragraphs = re.split(r"\n\s*\n", answer.strip())
        metadata["answer_paragraphs"] = paragraphs

        # Track which paragraphs contain which sentences
        matching_paragraph_indices = set()
        metadata["matching_sentence_paragraph_pairs"] = []

        for sentence in irrelevant_texts:
            for idx, para in enumerate(paragraphs):
                # Check for approximate match
                # Using approximate string matching instead of exact matching
                # to allow for some variations from LLM extracted text
                if (
                    difflib.SequenceMatcher(None, sentence.strip(), para).ratio()
                    >= 0.85
                ):
                    matching_paragraph_indices.add(idx)
                    metadata["matching_sentence_paragraph_pairs"].append(
                        {"sentence": sentence, "paragraph": para}
                    )
                    break  # Found a match for this sentence

        return 1 - (len(matching_paragraph_indices) / len(paragraphs)), metadata

    async def score_output(self, query: str, answer: str) -> Dict[str, Any]:
        prompt = PRECICION_EVAL_PROMPT.format(query=query, answer=answer)
        config = GenerateConfig(
            response_schema=answer_precision_json_schema,
            temperature=self.temperature,
        )
        result, parsed, num_retries = await generate_with_retry(
            self.model, prompt, config
        )
        score, metadata = self.map_output_to_precision_score(parsed, answer)
        metadata["num_retries"] = num_retries
        prompt_logs = _create_prompt_logs(
            prompt, query, answer, result.completion, self.temperature, metadata
        )

        metadata = {"prompt_logs": prompt_logs}
        return {"answer_precision": score}, metadata
