from typing import Dict, List, Union
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


# Define input schema
# Define the updated input schema for RecommendationAgent.
class RecommendationInputSchema(BaseIOSchema):
    """
    Schema for aggregated customer feedback for recommendations.
   
    Attributes:
      overall_sentiment_score: An integer score representing the overall sentiment.
      topic_sentiment: A dictionary mapping business topics to sentiment (e.g., {"Waiting Time": "negative"}).
      emotion_count: A dictionary counting emotions (e.g., {"joy": 2, "anger": 12}).
      words_sentiment: A dictionary mapping key words to sentiment (e.g., {"delicious": "positive"}).
      sentiment_breakdown: A dictionary representing the percentage breakdown of sentiments (e.g., {"positive": 70, "negative": 20, "mixed": 2, "neutral": 8}).
    """
    overall_sentiment_score: int
    topic_sentiment: Dict[str, Dict[str, int]]  # <- Changed
    emotion_count: Dict[str, int]              # <- Changed
    words_sentiment: Dict[str, Dict[str, int]] # <- Changed
    sentiment_breakdown: Dict[str, int]
    total_reviews: int

 
class RecommendationRowScema(BaseIOSchema):
    """_summary_
 
    Args:
        BaseIOSchema (_type_): _description_
 
    Returns:
        _type_: _description_
    """
    recommendation: str = Field(..., description="Recommendation for a specific topic.")
    topic: str = Field(..., description="Topic for which the recommendation is made.")
    severity: str = Field(..., description="Severity of the recommendation.")
    rank: int = Field(..., description="Rank of the recommendation.")
   
 
# Define the output schema for recommendations.
class RecommendationOutputSchema(BaseIOSchema):
    """
    Schema for business recommendations based on feedback.
   
    Attributes:
      recommendations: A dictionary mapping business topics to recommended actions.
    """
    recommendations: List[RecommendationRowScema] = Field(
        ..., description="Dictionary mapping business topics to actionable recommendations."
    )

# Define the AI-based recommendation agent
class RecommendationAgent:
    def __init__(self, client, model: str = "gpt-4o-mini", temperature: float = 0.2):
        """
        Initialize the RecommendationAgent.
 
        Args:
            client: The LLM client (e.g., from instructor.from_openai).
            model: The identifier for the model to use.
            temperature: Sampling temperature for the LLM.
        """
        # from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
 
        # Define a system prompt that instructs the agent on generating recommendations.
        system_prompt_generator = SystemPromptGenerator(
            background=["You are an expert business consultant specializing in customer feedback analysis."],
            steps=[
                "Review the aggregated customer feedback provided in the input.",
                "Consider the overall sentiment score, topic sentiment, emotion counts, words sentiment, and the sentiment breakdown.",
                "Identify the areas needing improvement (e.g., if 'Waiting Time' is negative or if emotion counts indicate high anger).",
                "Generate clear and actionable recommendations for each relevant topic or words to help the business improve."
            ],
            output_instructions=[
            "Return the output as a JSON object with a single key 'recommendations' which is a list of recommendation entries.",
            "Each entry in the list should be a dictionary containing: 'recommendation', 'topic', 'severity', and 'rank'.",
            "Include recommendations for each topic or word that shows significant negative or mixed sentiment, or important business opportunity.",
            "Rank recommendations based on urgency and severity, where 1 is most urgent.",
            "Do not group multiple topics into one recommendation — each should be separate."
        ]

        )
 
        # Configure the agent with the custom input and output schemas.
        config = AgentConfig(
            client=client,
            model=model,
            temperature=temperature,
            system_prompt_generator=system_prompt_generator
        )
        self.agent = AtomicAgent[RecommendationInputSchema,RecommendationOutputSchema](config=config)
 
    def generate_recommendations(self, aggregated_feedback: dict) -> RecommendationOutputSchema:
        """
        Generate business recommendations based on aggregated feedback.
 
        Args:
            aggregated_feedback: A dictionary containing overall_sentiment_score, topic_sentiment, emotion_count, words_sentiment, and sentiment_breakdown.
 
        Returns:
            An instance of RecommendationOutputSchema with the recommendations.
        """
        input_data = RecommendationInputSchema(**aggregated_feedback)
        result = self.agent.run(input_data)
        return result