from typing import Dict, Optional, Union

import afnio
import afnio.cognitive as cog
from afnio._variable import Variable
from afnio.models import ChatCompletionModel

SENTIMENT_RESPONSE_FORMAT = {
    "type": "json_schema",
    "json_schema": {
        "strict": True,
        "name": "sentiment_response_schema",
        "schema": {
            "type": "object",
            "properties": {
                "sentiment": {
                    "type": "string",
                    "enum": ["positive", "neutral", "negative"],
                },
            },
            "additionalProperties": False,
            "required": ["sentiment"],
        },
    },
}


class SentimentAnalyzer(cog.Module):
    """
    An agent that uses the [`ChatCompletion`][afnio.cognitive.ChatCompletion] module
    to analyze the sentiment of input messages. The sentiment is classified into
    three categories: positive, neutral, or negative.

    Examples:
        >>> import os
        >>>
        >>> import afnio
        >>> import afnio.utils.agents as agents
        >>> from afnio.models.openai import AsyncOpenAI
        >>>
        >>> os.environ["OPENAI_API_KEY"] = "sk-..."  # Replace with your actual key
        >>>
        >>> fwd_model = AsyncOpenAI()
        >>> agent = agents.SentimentAnalyzer()
        >>>
        >>> response = agent(
        ...     fwd_model,
        ...     inputs={"message": "I've been a satisfied client of ProCare for a year"},
        ...     model="gpt-4.1-nano",
        ...     temperature=0.0,
        ... )
        >>> print(response.data)
        {"sentiment":"positive"}
    """  # noqa: E501

    def __init__(self):
        super().__init__()

        self.sentiment_task = cog.Parameter(
            data="Read the provided message and determine the sentiment.",
            role="system prompt for sentiment classification",
            requires_grad=True,
        )
        self.sentiment_user = afnio.Variable(
            data="**Message:**\n\n{message}\n\n",
            role="input template to sentiment classifier",
        )
        self.sentiment_classifier = cog.ChatCompletion()

    def forward(
        self,
        fwd_model: Optional[ChatCompletionModel],
        inputs: Optional[Dict[str, Union[str, Variable]]] = None,
        **completion_args,
    ):
        sentiment_messages = [
            {"role": "system", "content": [self.sentiment_task]},
            {"role": "user", "content": [self.sentiment_user]},
        ]
        return self.sentiment_classifier(
            fwd_model,
            sentiment_messages,
            inputs=inputs,
            response_format=SENTIMENT_RESPONSE_FORMAT,
            **completion_args,
        )
