import logging
import re
from io import BytesIO
from pathlib import Path
from typing import Any

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.model import ChatMessage, ChatMessageUser, ContentImage, ContentText
from inspect_ai.scorer import (
    INCORRECT,
    AnswerPattern,
    Score,
    Scorer,
    Target,
    accuracy,
    scorer,
    stderr,
)
from inspect_ai.solver import (
    Generate,
    Solver,
    TaskState,
    solver,
)
from PIL import Image
from platformdirs import user_cache_dir

logger = logging.getLogger(__name__)

FREEFORM_TEMPLATE = r"""
Answer the following question. The entire content of your response should be of the following format: 'ANSWER: $ANSWER' (without quotes) where $ANSWER is your answer.

{question}
"""

IMAGE_BASE_DIR = Path(user_cache_dir("inspect_evals")) / "docvqa_images"


def _levenshtein_distance(str1: str, str2: str) -> int:
    """Computes a Levenshtein distance, same as Levenshtein.distance in the python-Levenshtein package."""
    # Create a matrix of size (len(str1) + 1) x (len(str2) + 1)
    matrix = [[0 for j in range(len(str2) + 1)] for i in range(len(str1) + 1)]

    # Initialize the first row and column
    for i in range(len(str1) + 1):
        matrix[i][0] = i
    for j in range(len(str2) + 1):
        matrix[0][j] = j

    # Fill in the rest of the matrix
    for i in range(1, len(str1) + 1):
        for j in range(1, len(str2) + 1):
            matrix[i][j] = min(
                matrix[i - 1][j] + 1,  # deletion
                matrix[i][j - 1] + 1,  # insertion
                matrix[i - 1][j - 1] + int(str1[i - 1] != str2[j - 1]),  # substitution
            )

    return matrix[len(str1)][len(str2)]


def _best_normalized_levenshtein_similiarity(
    completion: str, ground_truths: list[str], threshold: float
) -> float:
    """
    Compute a best normalized Levenshtein similiarity, an input into the Average Normalized Levenshtein Similiarity (ANLS)

    The Average Normalized Levenshtein Similarity (ANLS) is defined in equation (1) of
    https://arxiv.org/pdf/1907.00490.pdf

    Note that the "average" is computed by the accuracy metric -- not here. This function computes
    the term inside the summation of equation (1).
    """
    best_score = 0.0
    for ground_truth in ground_truths:
        if len(ground_truth) == 0 and len(completion) == 0:
            best_score = 1
            break
        levenshtein_distance = _levenshtein_distance(
            completion.lower(), ground_truth.lower()
        )
        normed_levenshtein_distance = levenshtein_distance / max(
            len(completion), len(ground_truth)
        )
        if normed_levenshtein_distance < threshold:
            score = 1.0 - normed_levenshtein_distance
        else:
            score = 0.0
        if score > best_score:
            best_score = score
    return best_score


@task
def docvqa() -> Task:
    dataset = hf_dataset(
        path="lmms-lab/DocVQA",
        name="DocVQA",
        split="validation",  # "answers" in the "test" split are held back by the authors
        sample_fields=record_to_sample,
        trust=True,
        shuffle=True,
    )

    return Task(
        dataset=dataset,
        solver=[docvqa_solver()],
        scorer=docvqa_scorer(),
    )


@scorer(metrics=[accuracy(), stderr(cluster="document_id")])
def docvqa_scorer() -> Scorer:
    async def normalized_levenshtein_similiarity_score(
        state: TaskState, target: Target
    ) -> Score:
        threshold = 0.5
        ground_truths = target.target
        match = re.search(
            AnswerPattern.LINE,
            state.output.completion,
            re.IGNORECASE,
        )
        if match:
            completion = match.groups()[0]
            return Score(
                value=_best_normalized_levenshtein_similiarity(
                    completion, ground_truths, threshold
                ),
                answer=completion,
            )

        else:
            # didn't find the scoring pattern
            return Score(
                value=INCORRECT,
                explanation="Scoring pattern not matched in output: "
                + f"{state.output.completion}",
            )

    return normalized_levenshtein_similiarity_score


@solver
def docvqa_solver() -> Solver:
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        state.user_prompt.text = FREEFORM_TEMPLATE.format(
            question=state.user_prompt.text
        )
        return await generate(state)

    return solve


def record_to_sample(record: dict[str, Any]) -> Sample:
    # extract image
    image_path = Path(IMAGE_BASE_DIR / record["image"]["path"])

    image_bytes = record["image"]["bytes"]
    assert is_image_png(image_bytes)

    if not image_path.exists():
        logger.debug(f"Extracting {image_path.name}")
        # ensure parent
        image_path.parent.mkdir(exist_ok=True, parents=True)
        # reduce the image size
        img = Image.open(BytesIO(image_bytes))
        img.thumbnail((1024, 1024))
        # save preserving format
        img.save(image_path, format=img.format)

    message: list[ChatMessage] = [
        ChatMessageUser(
            content=[
                ContentText(text=record["question"]),
                ContentImage(image=image_path.as_posix()),
            ]
        )
    ]

    return Sample(
        input=message,
        target=record["answers"],
        id=record["questionId"],
        metadata={"document_id": record["docId"]},
    )


def is_image_png(image_bytes: bytes) -> bool:
    return image_bytes[:8] == b"\x89\x50\x4e\x47\x0d\x0a\x1a\x0a"
