from typing import List, Dict
from pydantic import Field
# from atomic_agents.agents.base_agent import BaseAgent, BaseAgentConfig, BaseIOSchema
# NEW (v2.0)
from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema, BaseTool, BaseToolConfig
from atomic_agents.context import ChatHistory, SystemPromptGenerator, BaseDynamicContextProvider

import time

# Define the input schema for the review
class ReviewSentimentInputSchema(BaseIOSchema):
    """
    Schema for a batch of user reviews to be analyzed for sentiment.
    """
    reviews: List[str] = Field(..., description="List of user reviews to analyze.")


# Define the output schema for the sentiment analysis
class ReviewSentimentOutputSchema(BaseIOSchema):
    """
    Schema for analyzed sentiment data from the review.
    """
    overall_sentiment: str = Field(..., description="Overall sentiment of the review.")
    relevant_emotions: List[str] = Field(..., description="List of relevant emotions.")
    arousal: str = Field(..., description="Overall arousal level: active or passive.")
    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 ReviewSentimentAnalyzer:
    def __init__(self, client, model: str = "gpt-4", temperature: float = 0.2):
        """
        Initialize the ReviewSentimentAnalyzer 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
        """
        # from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator

        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 sentiment for 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 with a single key 'results' that contains a list of dicts.
                Each dict includes:
                - overall_sentiment
                - relevant_emotions
                - arousal
                - relevant_topics_sentiment
                - important_words_sentiment
                """
            ]
        )

        config = AgentConfig(
            client=client,
            model=model,
            temperature=temperature,
            system_prompt_generator=system_prompt_generator
        )

        self.agent = AtomicAgent[ReviewSentimentInputSchema,ReviewSentimentOutputSchema,](config=config)

    def analyze_reviews(self, reviews: List[str]) -> List[Dict]:
        """
        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 = ReviewSentimentInputSchema(reviews=reviews)
        result = self.agent.run(input_data)
        return result.results


    def analyze_review(self, review: str) -> Dict:
        """
        Analyzes a single review by wrapping it in a list and extracting the first result.
        """
        results = self.analyze_reviews([review])
        return results[0] if results else {}