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. Grade the generated article based on the following criteria, and give it a integer score between 1 and 5 inclusive, where 1 is the worst and 5 is the best.

## Criteria

{{CRITERION}}

## User Request

{{USER_REQUEST}}

## Generated Article

{{MODEL_OUTPUT}}

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:"


_CRITERIA = {
    "faithfulness": "Did the model say anything that isn’t supported by the facts?",
    "completeness": "Did the model include the important facts?",
    "style": "Does the model output follow the requested style, while keeping the facts intact?",
    "tone_and_intent": "Does the output follow the requested tone and intent i.e. what the text is trying to do emotionally/strategically?",
    "formality_and_register": "Does the output follow the requested formality and register i.e. how “corporate” vs “casual” the language is?",
    "point_of_view_and_voice": "Does the output follow the requested point of view and voice i.e. who is “speaking” and how direct it is?",
    "structure_and_formatting_habits": "Does the output follow the requested structure and formatting habits i.e. how the text is organized?",
}


class ArabicContentGenerationAnnotator(Annotator):
    name = "arabic_content_generation"

    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

        annotations: Dict[str, Dict[str, Union[int, str]]] = {}
        for criterion_id, criterion in _CRITERIA.items():
            annotator_prompt = (
                PROMPT_TEMPLATE.strip()
                .replace("{{CRITERION}}", criterion)
                .replace("{{USER_REQUEST}}", user_request_text)
                .replace("{{MODEL_OUTPUT}}", model_output_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

            annotations[criterion_id] = {
                "prompt": annotator_prompt,
                **self._parse_annotator_response(annotator_response_text),
            }
        return annotations

    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
