from typing import Dict, List, Any, Optional
from llamarch.common.llm import LLM
from llamarch.common.llm_embedding import LLMEmbedding
import numpy as np
from dataclasses import dataclass
import logging


@dataclass
class AgentResponse:
    """
    Data class to store responses generated by an agent.

    Attributes
    ----------
    agent_id : str
        Unique identifier for the agent.
    response : Any
        The generated response from the language model.
    confidence : float
        Confidence score of the response, typically between 0 and 1.
    metadata : Dict[str, Any]
        Additional metadata about the response, such as model information.
    """
    agent_id: str
    response: Any
    confidence: float
    metadata: Dict[str, Any]


class GenerativeAIAgent:
    """
    A generative AI agent that can generate responses and track performance.

    Parameters
    ----------
    agent_id : str
        Unique identifier for the agent.
    llm : LLM
        An instance of the LLM class, initialized with the desired model.
    embedding : LLMEmbedding, optional
        An instance of the LLMEmbedding class, initialized with the desired embedding model (default is None).
    """

    def __init__(self, agent_id: str, llm: LLM, embedding: Optional[LLMEmbedding] = None):
        self.agent_id = agent_id
        self.llm = llm
        self.embedding = embedding
        self.performance_history: List[float] = []
        self.logger = logging.getLogger(f"Agent-{agent_id}")

    async def generate_response(self, query: str) -> AgentResponse:
        """
        Generate a response to the given query using the language model.

        Parameters
        ----------
        query : str
            The input query to respond to.

        Returns
        -------
        AgentResponse
            An object containing the response text, confidence score, and metadata.

        Notes
        -----
        Currently, the confidence score is set to a default value of 1.0 and can be adjusted based on future requirements.
        """
        response_text = self.llm.generate(query)
        confidence = 1.0  # Placeholder; this could be dynamically calculated if desired
        metadata = {
            "model_name": self.llm.model_name
        }
        return AgentResponse(agent_id=self.agent_id, response=response_text, confidence=confidence, metadata=metadata)

    def update_performance(self, score: float):
        """
        Update the agent's performance history with a new score.

        Parameters
        ----------
        score : float
            The performance score to add to the agent's history.

        Notes
        -----
        This method also logs the updated performance score.
        """
        self.performance_history.append(score)
        self.logger.info(f"Agent {self.agent_id} performance updated: {score}")

    @property
    def average_performance(self) -> float:
        """
        Calculate the agent's average performance score.

        Returns
        -------
        float
            The average score from the agent's performance history. Returns 0.0 if no history is available.

        Notes
        -----
        This property uses NumPy's `mean` function to compute the average.
        """
        return np.mean(self.performance_history) if self.performance_history else 0.0
