from typing import Dict, Optional, Set

from helm.benchmark.annotation.model_as_judge import AnnotatorModelInfo, LLMAsJuryAnnotator
from helm.clients.auto_client import AutoClient


PROMPT_TEMPLATE = """You are a mental health expert tasked with evaluating the quality
of a counselor response in a therapy conversation.
Your goal is to assess how well the generated response demonstrates therapeutic skills,
and how it compares to the gold response in terms of accuracy, completeness, and clarity.

The conversation history and question will be provided in these tags:
<conversation>
{QUESTION}
</conversation>

The response will be provided in these tags:
<response>
{RESPONSE}
</response>

The reference response will be provided in these tags:
<gold_response>
{GOLD_RESPONSE}
</gold_response>

Carefully analyze the <response> compared to the <gold_response> and the original <conversation>.
For each of the following criteria, rate the response on a scale of 1 to 5 (1 = very poor, 5 = excellent),
and provide a short justification for your score.

Evaluation Criteria:
Accuracy (1-5)
- Does the response use appropriate therapeutic techniques and show understanding of the client's concerns?

Completeness (1-5)
- Does the response address the client's key concerns and include important context from the conversation?

Clarity (1-5)
- Is the response easy to understand and communicated in an empathetic, therapeutic way?

Output Format:
Output the evaluation as a single valid JSON object matching the following structure:
{
    "accuracy": {
        "score": 0,
        "explanation": "Explain why this score was given."
    },
    "completeness": {
        "score": 0,
        "explanation": "Explain why this score was given."
    },
    "clarity": {
        "score": 0,
        "explanation": "Explain why this score was given."
    }
}

Ensure the output is valid JSON:
- Use **double quotes** (") for all keys and string values.
- When quoting text or sections inside the explanations, use escaped double quotes (\") to
  maintain valid JSON formatting.
- Do not include any additional information in the output.
"""

ANNOTATION_CRITERIA: Dict[str, Set[str]] = {
    "accuracy": {"score", "explanation"},
    "completeness": {"score", "explanation"},
    "clarity": {"score", "explanation"},
}


class MentalHealthAnnotator(LLMAsJuryAnnotator):
    """The MentalHealth autograder."""

    def __init__(
        self,
        auto_client: AutoClient,
        annotator_models: Dict[str, AnnotatorModelInfo],
        template_name: Optional[str] = None,
    ):
        super().__init__(
            name="mental_health",
            auto_client=auto_client,
            prompt_template=PROMPT_TEMPLATE,
            annotation_criteria=ANNOTATION_CRITERIA,
            annotator_models=annotator_models,
        )
