from typing import List, Dict
from pydantic import Field
# from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig, BaseIOSchema
# from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator

# NEW (v2.0)
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema, BaseTool, BaseToolConfig
from atomic_agents.context import ChatHistory, SystemPromptGenerator, BaseDynamicContextProvider


class BatchReviewSentimentInputSchema(BaseIOSchema):
    """Schema for multiple user review contents."""
    reviews: List[str] = Field(
        ...,
        description="List of text contents of user reviews."
    )


class ReviewSentimentOutputSchema(BaseIOSchema):
    """
    Schema for analyzed sentiment data from the review.

    overall_sentiment: One of (positive, negative, mixed, neutral)
    relevant_emotions_with_arousal: List of dict mapping emotion to arousal (e.g. [{"anger": "passive"}])
    relevant_topics_sentiment: Dict mapping topics (e.g. Food Quality) to sentiment (e.g. "Positive")
    important_words_sentiment: Dict mapping important words to sentiment (e.g. {"delicious": "Positive"})
    """
    overall_sentiment: str = Field(..., description="Overall sentiment of the review.")
    relevant_emotions_with_arousal: List[Dict[str, str]] = Field(
        ..., description="List of relevant emotions with their arousal levels."
    )
    relevant_topics_sentiment: Dict[str, str] = Field(
        ..., description="Sentiment for each relevant topic."
    )
    important_words_sentiment: Dict[str, str] = Field(
        ..., description="Sentiment for important words extracted from the review."
    )
    
class BatchReviewSentimentOutputSchema(BaseIOSchema):
    """Schema for analyzed sentiment data from multiple reviews."""
    sentiments: List[ReviewSentimentOutputSchema] = Field(
        ...,
        description=(
            "List of analysis results for each review containing: "
            "overall_sentiment, relevant_emotions_with_arousal, "
            "relevant_topics_sentiment, and important_words_sentiment"
        )
    )

class BatchReviewSentimentAnalyzer:
    def __init__(self, client, model: str = "gpt-4", temperature: float = 0.2):
        """
        Initialize the BatchReviewSentimentAnalyzer for batch processing.

        Args:
            client: The LLM client (e.g., from instructor.from_openai)
            model: The model identifier to use
            temperature: Sampling temperature for the LLM
        """
        system_prompt_generator = SystemPromptGenerator(
            background=[
                "You are an expert sentiment analysis system that processes multiple reviews at once."
            ],
            steps=[
                "Analyze each review in the provided list independently.",
                "For each review, determine:",
                "1. Overall sentiment (positive, negative, mixed, neutral)",
                "2. Relevant emotions with arousal levels (e.g., [{'anger': 'passive'}])",
                "3. Sentiment for relevant topics (Service, Food Quality, etc.)",
                "   Evaluate each of the following topics: ",
                "Service Experience, Cleanliness & Hygiene, Ambiance & Atmosphere, ",
                "Pricing & Value, Online Ordering & Delivery, Complaint Resolution, ",
                "Customer Loyalty, Menu Preferences, Waiting Time, Portion Size, ",
                "Staff Friendliness, Payment Experience, Special Requests Handling, ",
                "Reservation Experience, Parking & Accessibility, Food Presentation, ",
                "Seasonal Trends, Offer & Promotions, Food Quality.",
                "4. Sentiment for important words in the review",
                "Maintain the exact same order of reviews in your output as received in input.",
                "Return all results in a single structured response."
            ],
            output_instructions=[
                "Return a JSON object whose value is "
                "an array of JSON objects. Each object must have the keys: "
                "overall_sentiment (str), relevant_emotions_with_arousal (list of dict), "
                "relevant_topics_sentiment (dict), and important_words_sentiment (dict).",
                "Do not wrap the JSON in markdown or add any extra text—only emit the raw JSON."
            ]
        )

        config = AgentConfig(
            client=client,
            model=model,
            temperature=temperature,
            system_prompt_generator=system_prompt_generator
        )
        self.agent = AtomicAgent[BatchReviewSentimentInputSchema, BatchReviewSentimentOutputSchema](config=config)

    def analyze_reviews(self, reviews: List[str]) -> List[ReviewSentimentOutputSchema]:
        """
        Analyze multiple reviews in a single batch request.

        Args:
            reviews: List of review texts to analyze

        Returns:
            List of analysis results in the same order as input reviews
        """
        input_data = BatchReviewSentimentInputSchema(reviews=reviews)
        result = self.agent.run(input_data)

        # print("result is ============", result)
        return result