"""
Base agent class for MAI-DxO specialized agents.

Provides common functionality and interface for all MAI-DxO agents
based on the Microsoft Research architecture.
"""

from abc import abstractmethod
from datetime import datetime
from typing import Any, Dict, List, Optional

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.models import ChatCompletionClient

from ...information.gatekeeper import InformationGatekeeper
from ...models.mai_dxo import (
    AgentMessage,
    AgentOutput,
    AgentRole,
    DecisionContext,
    Evidence,
    Hypothesis,
    InformationType,
    MessageType,
)


class MAIDxOBaseAgent(AssistantAgent):
    """Base class for all MAI-DxO specialized agents.

    Provides common functionality for hypothesis management, evidence tracking,
    and communication with the Information Gatekeeper.

    Attributes:
        agent_role: The specialized role this agent performs
        domain_context: The domain context (e.g., 'business', 'technical')
        information_gatekeeper: Gateway for accessing external information
        current_hypotheses: List of hypotheses currently maintained by the agent
        gathered_evidence: List of evidence collected by the agent
        agent_outputs: List of all outputs generated by the agent
        decision_context: Current decision context being analyzed
        information_requests_made: Count of information requests made
    """

    def __init__(
        self,
        name: str,
        model_client: ChatCompletionClient,
        agent_role: AgentRole,
        domain_context: str = "business",
        information_gatekeeper: Optional[InformationGatekeeper] = None,
        **kwargs,
    ):
        """Initialize the MAI-DxO base agent.

        Args:
            name: Display name for the agent
            model_client: Chat completion client for LLM interactions
            agent_role: The specialized role this agent performs
            domain_context: Domain context for specialized prompts
            information_gatekeeper: Gateway for accessing external information
            **kwargs: Additional arguments passed to parent class
        """
        # Get the specialized system prompt for this agent role
        system_message = self._get_system_prompt(agent_role, domain_context)

        super().__init__(
            name=name,
            model_client=model_client,
            system_message=system_message,
            **kwargs,
        )

        # Agent configuration
        self.agent_role = agent_role
        self.domain_context = domain_context
        self.information_gatekeeper = information_gatekeeper

        # Agent state tracking
        self.current_hypotheses: List[Hypothesis] = []
        self.gathered_evidence: List[Evidence] = []
        self.agent_outputs: List[AgentOutput] = []
        self.decision_context: Optional[DecisionContext] = None

        # Performance tracking
        self.information_requests_made = 0

    @abstractmethod
    def _get_system_prompt(self, agent_role: AgentRole, domain_context: str) -> str:
        """Get the specialized system prompt for this agent role and domain.

        Args:
            agent_role: The role this agent performs
            domain_context: The domain context for specialization

        Returns:
            String containing the system prompt for this agent
        """
        pass

    async def process_decision_context(self, context: DecisionContext) -> AgentOutput:
        """Process a decision context and generate initial analysis.

        Args:
            context: The decision context to analyze

        Returns:
            AgentOutput containing the agent's analysis and findings
        """
        self.decision_context = context

        # Generate initial analysis based on agent specialization
        initial_analysis = await self._generate_initial_analysis(context)

        # Create agent output
        output = AgentOutput(
            agent_role=self.agent_role,
            message_type=MessageType.ANALYSIS,
            content=initial_analysis,
            confidence_level=self._assess_confidence(initial_analysis),
            supporting_evidence=self.gathered_evidence.copy(),
            timestamp=datetime.now().isoformat(),
            reasoning=f"Initial {self.agent_role.value} analysis of the decision context",
        )

        self.agent_outputs.append(output)
        return output

    async def request_information(
        self,
        query: str,
        information_type: InformationType,
        priority: float = 0.5,
    ) -> Optional[str]:
        """Request information through the Information Gatekeeper.

        Args:
            query: The information query to request
            information_type: Type of information being requested
            priority: Priority level of the request (0.0-1.0)

        Returns:
            Information content if successful, None if failed
        """
        if not self.information_gatekeeper:
            return None

        try:
            context = {
                "agent_role": self.agent_role.value,
                "domain": self.domain_context,
                "current_hypotheses": len(self.current_hypotheses),
                "decision_context_id": self.decision_context.id
                if self.decision_context
                else None,
            }

            response = await self.information_gatekeeper.process_information_request(
                requesting_agent=self.agent_role,
                query=query,
                context=context,
                information_type=information_type,
                priority=priority,
            )

            # Track the request
            self.information_requests_made += 1

            # Convert response to evidence
            evidence = Evidence(
                source=response.source_type,
                content=response.information,
                reliability_score=response.reliability_score,
                relevance_score=priority,  # Use priority as relevance proxy
                timestamp=response.timestamp,
            )
            self.gathered_evidence.append(evidence)

            return response.information

        except Exception as e:
            # Handle information request failures gracefully
            print(f"Information request failed for {self.agent_role}: {e}")
            return None

    def update_hypotheses(self, new_hypotheses: List[Hypothesis]) -> None:
        """Update the current hypotheses based on new evidence.

        Args:
            new_hypotheses: List of updated hypotheses to replace current ones
        """
        self.current_hypotheses = new_hypotheses

    def get_current_hypotheses(self) -> List[Hypothesis]:
        """Get the current hypotheses maintained by this agent.

        Returns:
            List of current hypotheses (copy to prevent external modification)
        """
        return self.current_hypotheses.copy()

    def get_performance_metrics(self) -> Dict[str, Any]:
        """Get performance metrics for this agent.

        Returns:
            Dictionary containing various performance metrics and statistics
        """
        return {
            "agent_role": self.agent_role.value,
            "information_requests_made": self.information_requests_made,
            "evidence_gathered": len(self.gathered_evidence),
            "outputs_generated": len(self.agent_outputs),
            "current_hypotheses_count": len(self.current_hypotheses),
            "average_confidence": (
                sum(output.confidence_level for output in self.agent_outputs)
                / len(self.agent_outputs)
                if self.agent_outputs
                else 0.0
            ),
        }

    @abstractmethod
    async def _generate_initial_analysis(
        self, context: DecisionContext
    ) -> Dict[str, Any]:
        """Generate initial analysis specific to this agent's role.

        Args:
            context: The decision context to analyze

        Returns:
            Dictionary containing the agent's specialized analysis
        """
        pass

    def _assess_confidence(self, analysis: Dict[str, Any]) -> float:
        """Assess confidence level for the analysis.

        Args:
            analysis: The analysis content to assess

        Returns:
            Confidence score between 0.0 and 1.0
        """
        # Base confidence assessment - can be overridden by specialized agents
        if not analysis:
            return 0.0

        # Simple heuristic based on content richness
        content_keys = len(analysis.keys())
        evidence_count = len(self.gathered_evidence)

        base_confidence = min(0.9, 0.3 + (content_keys * 0.1) + (evidence_count * 0.05))
        return round(base_confidence, 2)

    def _create_agent_message(
        self,
        recipient: str,
        message_type: MessageType,
        content: Dict[str, Any],
        confidence: float,
    ) -> AgentMessage:
        """Create a standardized agent message.

        Args:
            recipient: The intended recipient of the message
            message_type: Type of message being sent
            content: Message content
            confidence: Confidence level in the message

        Returns:
            AgentMessage object ready for transmission
        """
        return AgentMessage(
            sender=self.agent_role,
            recipient=recipient,
            message_type=message_type,
            content=content,
            confidence_level=confidence,
            supporting_evidence=self.gathered_evidence.copy(),
            timestamp=datetime.now().isoformat(),
        )
