import re

from inspect_ai.scorer import Score, Scorer, Target, mean, scorer, stderr
from inspect_ai.solver import TaskState

SENTENCE_TRANSFORMER_REPO = (
    "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)
ROUGE_L_THRESHOLD = 0.2


def split_sentences(text: str, is_zh: bool) -> list[str]:
    """Split text into sentences (for ROUGE-L scoring)."""
    if is_zh:
        sentences = re.findall(r"[^。！？]+[。！？]?", text)
    else:
        sentences = re.split(r"(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|!)\s", text)
    sentences = [s.strip() for s in sentences if s.strip()]
    return sentences


def tokenize_sentence(sentence: str, is_zh: bool) -> str:
    """Tokenize sentences (for ROUGE-L scoring)."""
    # Import jieba here to avoid loading it in the global scope
    import jieba  # type: ignore

    return " ".join(jieba.cut(sentence)) if is_zh else sentence


@scorer(metrics=[mean(), stderr()])
def rouge_l_scorer(is_zh: bool) -> Scorer:
    """ROUGE-L scorer for the SEVENLLM benchmark.

    This scorer calculates the ROUGE-L score to evaluate the similarity between
    the model-generated output and the reference text on a per-sentence basis.
    The ROUGE-L metric focuses on the longest common subsequence (LCS) between
    the sentences, which captures the structure and fluency of the text. The scorer
    is designed to handle both English and Chinese texts, with sentence tokenization
    and splitting adjusted based on the `is_zh` flag.

    The routine:
    - Splits the target reference and model-generated text into sentences.
    - Tokenizes each sentence (with support for Chinese tokenization).
    - Computes the ROUGE-L similarity score for each pair of sentences.
    - Aggregates the scores, with sentences scoring above a defined threshold
      (default: 0.2) contributing to the final score proportionally.

    The total score reflects the percentage of sentences in the reference text
    that have sufficiently similar counterparts in the generated output,
    scaled to 100%.

    Args:
        is_zh (bool): A flag indicating whether the input text is in Chinese (`True`)
        or English (`False`).

    Returns:
        Scorer: A scoring function that computes the aggregated ROUGE-L score
        and provides the result alongside the model-generated output and a
        detailed explanation.
    """
    # Import rouge here to avoid loading it in the global scope
    from rouge import Rouge  # type: ignore

    rouge = Rouge()

    async def score(state: TaskState, target: Target) -> Score:
        test_sentences = split_sentences(target.text, is_zh)
        generated_sentences = split_sentences(state.output.completion, is_zh)
        num_sentences = len(test_sentences)
        per_sentence_score = 100.0 / num_sentences  # Adjust scaling as needed
        total_score = 0.0

        for test_sentence in test_sentences:
            max_similarity = 0.0
            for gen_sentence in generated_sentences:
                test_sentence_tok = tokenize_sentence(test_sentence, is_zh)
                gen_sentence_tok = tokenize_sentence(gen_sentence, is_zh)

                scores = rouge.get_scores(test_sentence_tok, gen_sentence_tok)
                similarity = scores[0]["rouge-l"]["f"]
                if similarity > max_similarity:
                    max_similarity = similarity
            if max_similarity >= ROUGE_L_THRESHOLD:
                total_score += per_sentence_score

        return Score(
            value=total_score,
            answer=state.output.completion,
            explanation=f"Accumulated ROUGE-L score: {total_score:.4f}",
        )

    return score


@scorer(metrics=[mean(), stderr()])
def semantic_similarity_scorer() -> Scorer:
    """Semantic similarity scorer for the SEVENLLM benchmark.

    This scorer evaluates the semantic similarity between the model-generated output
    and the reference target text using a pre-trained SentenceTransformer model.
    It computes cosine similarity between the embeddings of the two texts, providing
    a quantitative measure of their alignment in meaning.

    The routine:
    - Encodes both the model's output and the reference target text into embeddings
      using a SentenceTransformer.
    - Computes the cosine similarity between the two embeddings as the similarity score.
    - Returns the score along with the model output and an explanation string.

    The score is aggregated using mean and standard error (stderr) metrics to evaluate
    model performance across the benchmark.

    Returns:
        Scorer: A scoring function that computes semantic similarity and provides
        both the similarity score and a detailed explanation of the result.
    """
    # Import SentenceTransformer here to avoid loading it in the global scope
    from sentence_transformers import SentenceTransformer, util  # type: ignore

    model = SentenceTransformer(SENTENCE_TRANSFORMER_REPO)

    async def score(state: TaskState, target: Target) -> Score:
        model_output = state.output.completion
        target_text = target.text
        embeddings = model.encode([model_output, target_text], convert_to_tensor=True)
        cosine_similarity = util.pytorch_cos_sim(embeddings[0], embeddings[1])
        similarity_score = cosine_similarity.item()

        return Score(
            value=similarity_score,
            answer=model_output,
            explanation=f"Semantic similarity score: {similarity_score:.4f}",
        )

    return score
