"""
Agent Decision Tracing Utilities for Daita Agents - Fixed Complete Version

Provides simple tools for agents to trace their decision-making process,
reasoning chains, and confidence scores. Integrates seamlessly with the
TraceManager for automatic decision observability.

FIXED ISSUES:
- Completed the trace_decision decorator implementation
- Fixed async context manager issues
- Added missing helper functions
- Improved error handling
- Fixed circular import issues
"""

import asyncio
import logging
import time
import functools
from typing import Dict, Any, Optional, List, Union, Callable, Tuple
from contextlib import asynccontextmanager
from dataclasses import dataclass
from enum import Enum

logger = logging.getLogger(__name__)


class DecisionEventType(str, Enum):
    """Types of decision events that can be streamed in real-time."""

    DECISION_STARTED = "decision_started"
    CONFIDENCE_UPDATED = "confidence_updated"
    REASONING_ADDED = "reasoning_added"
    ALTERNATIVE_ADDED = "alternative_added"
    FACTOR_SET = "factor_set"
    DECISION_COMPLETED = "decision_completed"


@dataclass
class DecisionEvent:
    """Real-time decision event for streaming."""

    event_type: DecisionEventType
    decision_point: str
    span_id: Optional[str]
    timestamp: float
    data: Dict[str, Any]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "event_type": self.event_type.value,
            "decision_point": self.decision_point,
            "span_id": self.span_id,
            "timestamp": self.timestamp,
            "data": self.data,
        }


class DecisionType(str, Enum):
    """Common types of agent decisions."""

    CLASSIFICATION = "classification"
    ANALYSIS = "analysis"
    RECOMMENDATION = "recommendation"
    ROUTING = "routing"
    VALIDATION = "validation"
    SELECTION = "selection"
    PRIORITIZATION = "prioritization"
    CUSTOM = "custom"


@dataclass
class DecisionContext:
    """Simple container for decision context and metadata."""

    decision_point: str
    decision_type: DecisionType
    confidence_score: float = 0.0
    reasoning_chain: List[str] = None
    alternatives: List[str] = None
    factors: Dict[str, Any] = None

    def __post_init__(self):
        if self.reasoning_chain is None:
            self.reasoning_chain = []
        if self.alternatives is None:
            self.alternatives = []
        if self.factors is None:
            self.factors = {}


class DecisionRecorder:
    """
    Context manager for recording agent decisions.

    Provides a simple interface for agents to record their decision-making
    process without needing to understand the underlying tracing system.
    """

    def __init__(
        self,
        decision_point: str,
        decision_type: Union[str, DecisionType] = DecisionType.CUSTOM,
        agent_id: Optional[str] = None,
        stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
    ):
        self.decision_point = decision_point
        self.decision_type = (
            DecisionType(decision_type)
            if isinstance(decision_type, str)
            else decision_type
        )
        self.agent_id = agent_id
        self.span_id = None
        self.context = DecisionContext(decision_point, self.decision_type)
        self.stream_callback = stream_callback

        # Import here to avoid circular imports
        from .tracing import get_trace_manager

        self.trace_manager = get_trace_manager()

    def _emit_event(self, event_type: DecisionEventType, data: Dict[str, Any]):
        """Emit a real-time decision event if callback is provided."""
        event = DecisionEvent(
            event_type=event_type,
            decision_point=self.decision_point,
            span_id=self.span_id,
            timestamp=time.time(),
            data=data,
        )

        # Emit to local callback if provided
        if self.stream_callback:
            try:
                self.stream_callback(event)
            except Exception as e:
                logger.debug(f"Error calling local decision stream callback: {e}")

        # Also emit to TraceManager for centralized streaming
        if self.agent_id:
            try:
                self.trace_manager.emit_decision_event(self.agent_id, event)
            except Exception as e:
                logger.debug(f"Error emitting decision event to TraceManager: {e}")

    async def __aenter__(self):
        """Start decision tracing."""
        try:
            # Start the decision span
            self.span_id = self.trace_manager.start_span(
                operation_name=f"decision_{self.decision_point}",
                trace_type="decision_trace",  # Use string to avoid enum import issues
                agent_id=self.agent_id,
                decision_point=self.decision_point,
                decision_type=self.decision_type.value,
            )

            logger.debug(f"Started decision trace: {self.decision_point}")

            # Emit decision started event
            self._emit_event(
                DecisionEventType.DECISION_STARTED,
                {"decision_type": self.decision_type.value, "agent_id": self.agent_id},
            )

            return self

        except Exception as e:
            logger.error(f"Failed to start decision trace: {e}")
            # Return self anyway so calling code doesn't break
            return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """End decision tracing with recorded context."""
        try:
            if self.span_id:
                # Record the decision details
                self.trace_manager.record_decision(
                    span_id=self.span_id,
                    confidence=self.context.confidence_score,
                    reasoning=self.context.reasoning_chain,
                    alternatives=self.context.alternatives,
                    **self.context.factors,
                )

                # End the span
                from .tracing import TraceStatus

                status = TraceStatus.ERROR if exc_type else TraceStatus.SUCCESS
                error_message = str(exc_val) if exc_val else None

                self.trace_manager.end_span(
                    span_id=self.span_id,
                    status=status,
                    error_message=error_message,
                    output_data=self.get_summary(),
                )

                logger.debug(
                    f"Ended decision trace: {self.decision_point} (confidence: {self.context.confidence_score:.2f})"
                )

                # Emit decision completed event
                self._emit_event(
                    DecisionEventType.DECISION_COMPLETED,
                    {
                        "final_confidence": self.context.confidence_score,
                        "total_reasoning_steps": len(self.context.reasoning_chain),
                        "total_alternatives": len(self.context.alternatives),
                        "status": (
                            status.value if hasattr(status, "value") else str(status)
                        ),
                        "success": exc_type is None,
                    },
                )

        except Exception as e:
            logger.error(f"Failed to end decision trace: {e}")

    def set_confidence(self, confidence: float):
        """Set confidence score (0.0 to 1.0)."""
        old_confidence = self.context.confidence_score
        self.context.confidence_score = max(0.0, min(1.0, confidence))

        # Emit confidence updated event
        self._emit_event(
            DecisionEventType.CONFIDENCE_UPDATED,
            {
                "old_confidence": old_confidence,
                "new_confidence": self.context.confidence_score,
                "confidence_change": self.context.confidence_score - old_confidence,
            },
        )

    def add_reasoning(self, reasoning: str):
        """Add a reasoning step to the decision chain."""
        self.context.reasoning_chain.append(reasoning)

        # Emit reasoning added event
        self._emit_event(
            DecisionEventType.REASONING_ADDED,
            {
                "reasoning": reasoning,
                "step_number": len(self.context.reasoning_chain),
                "total_steps": len(self.context.reasoning_chain),
            },
        )

    def add_alternative(self, alternative: str):
        """Add an alternative option that was considered."""
        self.context.alternatives.append(alternative)

        # Emit alternative added event
        self._emit_event(
            DecisionEventType.ALTERNATIVE_ADDED,
            {
                "alternative": alternative,
                "total_alternatives": len(self.context.alternatives),
            },
        )

    def set_factor(self, key: str, value: Any):
        """Set a decision factor (e.g., data_quality: 0.9)."""
        self.context.factors[key] = value

        # Emit factor set event
        self._emit_event(
            DecisionEventType.FACTOR_SET,
            {
                "factor_key": key,
                "factor_value": value,
                "total_factors": len(self.context.factors),
            },
        )

    def get_summary(self) -> Dict[str, Any]:
        """Get a summary of the decision for logging or return."""
        return {
            "decision_point": self.decision_point,
            "decision_type": self.decision_type.value,
            "confidence": self.context.confidence_score,
            "reasoning_steps": len(self.context.reasoning_chain),
            "alternatives_considered": len(self.context.alternatives),
            "factors": list(self.context.factors.keys()),
        }


# Context manager factory
@asynccontextmanager
async def record_decision_point(
    decision_point: str,
    decision_type: Union[str, DecisionType] = DecisionType.CUSTOM,
    agent_id: Optional[str] = None,
    stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
):
    """
    Context manager for recording a decision point.

    Usage:
        async with record_decision_point("data_classification") as decision:
            result = classify(data)
            decision.set_confidence(0.85)
            decision.add_reasoning("Pattern match found")
            return result
    """
    recorder = DecisionRecorder(
        decision_point, decision_type, agent_id, stream_callback
    )
    async with recorder as decision:
        yield decision


# Complete decorator for automatic decision tracing
def trace_decision(
    decision_point: str,
    decision_type: Union[str, DecisionType] = DecisionType.CUSTOM,
    extract_confidence: bool = True,
    extract_reasoning: bool = True,
):
    """
    Decorator for automatic decision tracing.

    The decorated function can return:
    1. Just the result
    2. (result, confidence)
    3. (result, confidence, reasoning_list)
    4. (result, {"confidence": X, "reasoning": [...], "alternatives": [...]})

    Usage:
        @trace_decision("classification", DecisionType.CLASSIFICATION)
        async def classify_data(self, data):
            # Your logic here
            result = {"class": "positive"}
            confidence = 0.85
            reasoning = ["Pattern A detected", "Threshold met"]
            return result, confidence, reasoning
    """

    def decorator(func: Callable):
        if asyncio.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                # Extract agent_id from self if available
                agent_id = None
                if args and hasattr(args[0], "agent_id"):
                    agent_id = args[0].agent_id

                async with record_decision_point(
                    decision_point, decision_type, agent_id
                ) as decision:
                    try:
                        result = await func(*args, **kwargs)

                        # Extract decision metadata from result
                        confidence, reasoning, alternatives = (
                            _extract_decision_metadata(
                                result, extract_confidence, extract_reasoning
                            )
                        )

                        if confidence is not None:
                            decision.set_confidence(confidence)

                        if reasoning:
                            for reason in reasoning:
                                decision.add_reasoning(reason)

                        if alternatives:
                            for alt in alternatives:
                                decision.add_alternative(alt)

                        # Return just the main result (strip metadata)
                        if isinstance(result, tuple) and len(result) > 1:
                            return result[0]
                        return result

                    except Exception as e:
                        logger.error(f"Decision function {func.__name__} failed: {e}")
                        raise

            return async_wrapper
        else:

            @functools.wraps(func)
            def sync_wrapper(*args, **kwargs):
                # For sync functions, use asyncio.run if needed
                async def async_exec():
                    agent_id = None
                    if args and hasattr(args[0], "agent_id"):
                        agent_id = args[0].agent_id

                    async with record_decision_point(
                        decision_point, decision_type, agent_id
                    ) as decision:
                        try:
                            result = func(*args, **kwargs)

                            confidence, reasoning, alternatives = (
                                _extract_decision_metadata(
                                    result, extract_confidence, extract_reasoning
                                )
                            )

                            if confidence is not None:
                                decision.set_confidence(confidence)

                            if reasoning:
                                for reason in reasoning:
                                    decision.add_reasoning(reason)

                            if alternatives:
                                for alt in alternatives:
                                    decision.add_alternative(alt)

                            if isinstance(result, tuple) and len(result) > 1:
                                return result[0]
                            return result

                        except Exception as e:
                            logger.error(
                                f"Decision function {func.__name__} failed: {e}"
                            )
                            raise

                try:
                    # Try to get current event loop
                    loop = asyncio.get_event_loop()
                    if loop.is_running():
                        # We're already in an async context
                        task = asyncio.create_task(async_exec())
                        task.add_done_callback(
                            lambda t: t.exception() if not t.cancelled() else None
                        )
                        return task
                    else:
                        return loop.run_until_complete(async_exec())
                except RuntimeError:
                    # No event loop, create one
                    return asyncio.run(async_exec())

            return sync_wrapper

    return decorator


def _extract_decision_metadata(
    result, extract_confidence: bool, extract_reasoning: bool
) -> Tuple[Optional[float], List[str], List[str]]:
    """Extract confidence, reasoning, and alternatives from function result."""
    confidence = None
    reasoning = []
    alternatives = []

    try:
        if isinstance(result, tuple):
            if len(result) >= 2 and extract_confidence:
                # (result, confidence) or (result, confidence, reasoning)
                conf_value = result[1]
                if isinstance(conf_value, (int, float)):
                    confidence = float(conf_value)
                elif isinstance(conf_value, dict):
                    # (result, metadata_dict)
                    confidence = conf_value.get("confidence")
                    reasoning = conf_value.get("reasoning", [])
                    alternatives = conf_value.get("alternatives", [])

            if len(result) >= 3 and extract_reasoning:
                # (result, confidence, reasoning)
                reasoning_value = result[2]
                if isinstance(reasoning_value, list):
                    reasoning = reasoning_value
                elif isinstance(reasoning_value, str):
                    reasoning = [reasoning_value]

            if len(result) >= 4:
                # (result, confidence, reasoning, alternatives)
                alt_value = result[3]
                if isinstance(alt_value, list):
                    alternatives = alt_value

        elif isinstance(result, dict) and "confidence" in result:
            # Result is a dict with metadata
            confidence = result.get("confidence")
            reasoning = result.get("reasoning", [])
            alternatives = result.get("alternatives", [])

    except Exception as e:
        logger.debug(f"Error extracting decision metadata: {e}")

    return confidence, reasoning, alternatives


# Helper functions for common decision patterns


async def record_decision(
    decision_point: str,
    decision_type: DecisionType,
    result: Any,
    confidence: float,
    reasoning_items: List[str],
    factors: Optional[Dict[str, Any]] = None,
    factor_prefix: str = "",
    alternatives: Optional[List[str]] = None,
    result_key: str = "result",
    agent_id: Optional[str] = None,
    stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
) -> Dict[str, Any]:
    """
    Record a decision point with confidence, reasoning, and optional factors.

    This is the single entry point for all decision recording. The three legacy
    helpers (record_classification_decision, record_analysis_decision,
    record_recommendation_decision) are thin aliases over this function.

    Args:
        decision_point: Identifier for the decision (e.g. "sentiment_analysis")
        decision_type: DecisionType enum value (CLASSIFICATION, ANALYSIS, RECOMMENDATION)
        result: The decision output (str, dict, or any serializable value)
        confidence: Confidence score between 0.0 and 1.0
        reasoning_items: List of reasoning strings explaining the decision
        factors: Optional dict of named factors to record (e.g. feature weights, risk levels)
        factor_prefix: Optional prefix applied to each factor key (e.g. "weight_", "risk_")
        alternatives: Optional list of alternatives considered
        result_key: Key name used for the result in the returned dict (default: "result")
        agent_id: Optional agent ID for tracing context
        stream_callback: Optional callback for streaming decision events

    Usage:
        # Classification
        result = await record_decision(
            "sentiment_analysis", DecisionType.CLASSIFICATION,
            result="positive", confidence=0.87,
            reasoning_items=["Strong positive word sentiment", "Context supports positive tone"],
            factors={"word_sentiment": 0.6, "context_analysis": 0.4}, factor_prefix="weight_",
            alternatives=["neutral", "negative"],
        )

        # Recommendation
        result = await record_decision(
            "investment_advice", DecisionType.RECOMMENDATION,
            result="buy", confidence=0.78,
            reasoning_items=["Strong fundamentals", "Market conditions favorable"],
            factors={"market_volatility": "medium", "liquidity": "high"}, factor_prefix="risk_",
            alternatives=["hold", "sell"],
        )
    """
    async with record_decision_point(
        decision_point, decision_type, agent_id, stream_callback
    ) as decision:
        decision.set_confidence(confidence)
        for item in reasoning_items:
            decision.add_reasoning(item)
        if alternatives:
            for alt in alternatives:
                decision.add_alternative(alt)
        if factors:
            for key, value in factors.items():
                decision.set_factor(f"{factor_prefix}{key}", value)
        return {
            result_key: result,
            "confidence": confidence,
            "decision_summary": decision.get_summary(),
        }


# Backward-compatible aliases — prefer record_decision() for new code.
async def record_classification_decision(
    decision_point: str,
    classification_result: str,
    confidence: float,
    features_used: List[str],
    alternatives_considered: Optional[List[str]] = None,
    feature_weights: Optional[Dict[str, float]] = None,
    agent_id: Optional[str] = None,
    stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
) -> Dict[str, Any]:
    """Alias for record_decision() with DecisionType.CLASSIFICATION."""
    return await record_decision(
        decision_point,
        DecisionType.CLASSIFICATION,
        classification_result,
        confidence,
        reasoning_items=[f"Feature used: {f}" for f in features_used],
        factors=feature_weights,
        factor_prefix="weight_",
        alternatives=alternatives_considered,
        result_key="classification",
        agent_id=agent_id,
        stream_callback=stream_callback,
    )


async def record_analysis_decision(
    decision_point: str,
    analysis_result: Dict[str, Any],
    confidence: float,
    key_insights: List[str],
    data_quality_factors: Optional[Dict[str, float]] = None,
    agent_id: Optional[str] = None,
    stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
) -> Dict[str, Any]:
    """Alias for record_decision() with DecisionType.ANALYSIS."""
    return await record_decision(
        decision_point,
        DecisionType.ANALYSIS,
        analysis_result,
        confidence,
        reasoning_items=key_insights,
        factors=data_quality_factors,
        result_key="analysis",
        agent_id=agent_id,
        stream_callback=stream_callback,
    )


async def record_recommendation_decision(
    decision_point: str,
    recommendation: str,
    confidence: float,
    rationale: List[str],
    alternatives_considered: Optional[List[str]] = None,
    risk_factors: Optional[Dict[str, str]] = None,
    agent_id: Optional[str] = None,
    stream_callback: Optional[Callable[[DecisionEvent], None]] = None,
) -> Dict[str, Any]:
    """Alias for record_decision() with DecisionType.RECOMMENDATION."""
    return await record_decision(
        decision_point,
        DecisionType.RECOMMENDATION,
        recommendation,
        confidence,
        reasoning_items=rationale,
        factors=risk_factors,
        factor_prefix="risk_",
        alternatives=alternatives_considered,
        result_key="recommendation",
        agent_id=agent_id,
        stream_callback=stream_callback,
    )


# Utility functions for decision analysis


def get_recent_decisions(
    agent_id: Optional[str] = None, decision_type: Optional[str] = None, limit: int = 20
) -> List[Dict[str, Any]]:
    """Get recent decision traces."""
    try:
        from .tracing import get_trace_manager

        trace_manager = get_trace_manager()
        operations = trace_manager.get_recent_operations(
            agent_id=agent_id, limit=limit * 2
        )

        # Filter for decision traces
        decisions = [op for op in operations if op.get("type") == "decision_trace"]

        # Filter by decision type if specified
        if decision_type:
            decisions = [
                op
                for op in decisions
                if op.get("metadata", {}).get("decision_type") == decision_type
            ]

        return decisions[:limit]
    except Exception as e:
        logger.error(f"Error getting recent decisions: {e}")
        return []


def register_agent_decision_stream(
    agent_id: str, callback: Callable[[DecisionEvent], None]
) -> None:
    """
    Register a callback to receive streaming decision events for a specific agent.

    This provides centralized decision event streaming through the TraceManager.

    Usage:
        def my_callback(event: DecisionEvent):
            print(f"Decision event: {event.event_type} - {event.data}")

        register_agent_decision_stream("my-agent-id", my_callback)
    """
    try:
        from .tracing import get_trace_manager

        trace_manager = get_trace_manager()
        trace_manager.register_decision_stream_callback(agent_id, callback)
        logger.info(f"Registered decision stream for agent {agent_id}")
    except Exception as e:
        logger.error(f"Failed to register agent decision stream: {e}")


def unregister_agent_decision_stream(
    agent_id: str, callback: Callable[[DecisionEvent], None]
) -> None:
    """
    Unregister a decision stream callback for a specific agent.

    Usage:
        unregister_agent_decision_stream("my-agent-id", my_callback)
    """
    try:
        from .tracing import get_trace_manager

        trace_manager = get_trace_manager()
        trace_manager.unregister_decision_stream_callback(agent_id, callback)
        logger.info(f"Unregistered decision stream for agent {agent_id}")
    except Exception as e:
        logger.error(f"Failed to unregister agent decision stream: {e}")


def get_streaming_agents() -> List[str]:
    """Get list of agents that have decision streaming enabled."""
    try:
        from .tracing import get_trace_manager

        trace_manager = get_trace_manager()
        return trace_manager.get_streaming_agents()
    except Exception as e:
        logger.error(f"Failed to get streaming agents: {e}")
        return []


def get_decision_stats(
    agent_id: Optional[str] = None, decision_type: Optional[str] = None
) -> Dict[str, Any]:
    """Get decision statistics for analysis."""
    try:
        decisions = get_recent_decisions(agent_id, decision_type, limit=50)

        if not decisions:
            return {"total_decisions": 0, "average_confidence": 0}

        # Calculate statistics
        total_decisions = len(decisions)
        successful_decisions = len(
            [d for d in decisions if d.get("status") == "success"]
        )

        # Extract confidence scores
        confidences = []
        for decision in decisions:
            confidence = decision.get("metadata", {}).get("confidence_score")
            if confidence is not None:
                confidences.append(confidence)

        avg_confidence = sum(confidences) / len(confidences) if confidences else 0

        # Decision type distribution
        type_counts = {}
        for decision in decisions:
            dec_type = decision.get("metadata", {}).get("decision_type", "unknown")
            type_counts[dec_type] = type_counts.get(dec_type, 0) + 1

        return {
            "total_decisions": total_decisions,
            "successful_decisions": successful_decisions,
            "success_rate": (
                successful_decisions / total_decisions if total_decisions > 0 else 0
            ),
            "average_confidence": avg_confidence,
            "confidence_count": len(confidences),
            "decision_types": type_counts,
            "agent_id": agent_id,
            "filter_type": decision_type,
        }
    except Exception as e:
        logger.error(f"Error getting decision stats: {e}")
        return {"total_decisions": 0, "average_confidence": 0}


# Export everything
__all__ = [
    # Enums
    "DecisionType",
    "DecisionEventType",
    # Event classes
    "DecisionEvent",
    # Main interfaces
    "record_decision_point",
    "trace_decision",
    # Helper functions
    "record_classification_decision",
    "record_analysis_decision",
    "record_recommendation_decision",
    # Analysis functions
    "get_recent_decisions",
    "get_decision_stats",
    # Streaming functions
    "register_agent_decision_stream",
    "unregister_agent_decision_stream",
    "get_streaming_agents",
    # Classes (for advanced usage)
    "DecisionRecorder",
    "DecisionContext",
]
