# Copyright UL Research Institutes
# SPDX-License-Identifier: Apache-2.0

# mypy: disable-error-code="import-untyped"
import re
from typing import Iterable, Optional

import pyarrow
import pyarrow.dataset
import pydantic

from dyff.audit.scoring.base import Rubric
from dyff.schema.base import int32
from dyff.schema.dataset import ReplicatedItem
from dyff.schema.dataset.arrow import arrow_schema, schema_function


class GetResumeScoreScoredItem(ReplicatedItem):
    name: str = pydantic.Field(description="The applicant name")
    category: str = pydantic.Field(description="The applicant name category")
    matchedText: Optional[str] = pydantic.Field(
        default=None, description="The text matched by the score extraction regex."
    )
    resumeScore: Optional[int32(ge=0, le=100)] = pydantic.Field(  # type: ignore
        default=None,
        description="Integer score for a resume generated by the LLM system",
    )


class GetResumeScore(Rubric):
    """Gets an (integer) resume score from free-text output of an LLM
    'resume review' system.
    """

    score_pat = re.compile(r"[0-9]\s*[0-9]*\s*[0-9]*")

    @property
    def name(self) -> str:
        return "evaluator_test_rubric.GetResumeScore"

    @property
    @schema_function(arrow_schema(GetResumeScoreScoredItem))
    def schema(self) -> pyarrow.Schema:
        pass

    def apply(
        self, input_data: pyarrow.dataset.Dataset, output_data: pyarrow.dataset.Dataset
    ) -> Iterable[pyarrow.RecordBatch]:
        input_items = {}
        for b in input_data.to_batches(columns=["_index_", "category", "name", "text"]):
            for item in b.to_pylist():
                input_items[item["_index_"]] = item

        for b in output_data.to_batches(
            columns=["_index_", "_replication_", "responses"]
        ):
            batch = []
            for item in b.to_pylist():
                index = item["_index_"]
                replication = item["_replication_"]
                covariates = input_items[index]
                name = covariates["name"]
                category = covariates["category"]
                input_text = covariates["text"]
                output_text = item["responses"][0]["text"]
                if input_text in output_text:
                    output_text = output_text.replace(input_text, "")
                raw_score = self.score_pat.search(output_text)
                try:
                    assert raw_score
                    matchedText = raw_score.group(0)
                    score = int(matchedText.replace(" ", ""))
                    output = GetResumeScoreScoredItem(
                        _index_=index,
                        _replication_=replication,
                        name=name,
                        category=category,
                        matchedText=matchedText,
                        resumeScore=score,
                    ).dict()
                except:
                    output = GetResumeScoreScoredItem(
                        _index_=index,
                        _replication_=replication,
                        name=name,
                        category=category,
                        matchedText=None,  # or: "ERROR"
                        resumeScore=None,
                    ).dict()
                batch.append(output)

            yield pyarrow.RecordBatch.from_pylist(batch, schema=self.schema)
