from typing import Any, Dict, Union

from helm.benchmark.adaptation.request_state import RequestState
from helm.benchmark.annotation.annotator import Annotator
from helm.clients.auto_client import AutoClient
from helm.common.request import Request


PROMPT_TEMPLATE = """The following is an article generated by a LLM based on a user request, followed by a reference article. Grade the generated article with integer score between 1 and 5 inclusive, where:

- 1 - reference article is much better
- 2 - reference article is slightly better
- 3 - both articles are about as good
- 4 - generated article is slightly better
- 5 - generated article is much better

## User Request

{{USER_REQUEST}}

## Generated Article

{{MODEL_OUTPUT}}

## Reference Article

{{REFERENCE_ARTICLE}}

Respond with a short paragraph starting with "REASONING: " explaining your reasoning in English. Then respond with a paragraph starting with "SCORE: " followed by a single number from 1 to 5. Do not respond with anything else.
"""  # noqa: E501


_REASONING_TAG = "REASONING:"
_SCORE_TAG = "SCORE:"


class ArabicContentGenerationRelativeAnnotator(Annotator):
    name = "arabic_content_generation_reference_comparison"

    def __init__(self, auto_client: AutoClient):
        self._auto_client = auto_client

    def annotate(self, request_state: RequestState) -> Any:
        assert request_state.result

        user_request_text = request_state.request.prompt

        assert len(request_state.result.completions) == 1
        model_output_text = request_state.result.completions[0].text

        assert len(request_state.instance.references) == 1
        reference_text = request_state.instance.references[0].output.text

        annotator_prompt = (
            PROMPT_TEMPLATE.strip()
            .replace("{{USER_REQUEST}}", user_request_text)
            .replace("{{MODEL_OUTPUT}}", model_output_text)
            .replace("{{REFERENCE_ARTICLE}}", reference_text)
        )

        annotator_request = Request(
            model="openai/gpt-5.1-2025-11-13",
            model_deployment="openai/gpt-5.1-2025-11-13",
            prompt=annotator_prompt,
            temperature=0.0,
            max_tokens=2000,
        )
        annotator_response = self._auto_client.make_request(annotator_request)
        assert len(annotator_response.completions) == 1
        annotator_response_text = annotator_response.completions[0].text

        return {"prompt": annotator_prompt, **self._parse_annotator_response(annotator_response_text)}

    def _parse_annotator_response(self, annotator_response_text: str) -> Dict[str, Union[int, str]]:
        result: Dict[str, Union[int, str]] = {}
        for line in annotator_response_text.split("\n"):
            if line.startswith(_REASONING_TAG):
                result["reasoning"] = line.removeprefix(_REASONING_TAG).strip()
            elif line.startswith(_SCORE_TAG):
                result["score"] = int(line.removeprefix(_SCORE_TAG).strip())
        return result
