"""Agent learning memory system with BM25-based retrieval.

This module provides a memory system for trading agents to store and retrieve
experiences, decisions, and market observations using BM25 text indexing.
"""

import uuid
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Callable
from dataclasses import dataclass, field

from openclaw.utils.logging import get_logger
from openclaw.memory.bm25_index import BM25Index, MemoryDocument


class MemoryType(str, Enum):
    """Types of memories stored in the system."""

    TRADE = "trade_memory"
    MARKET = "market_memory"
    DECISION = "decision_memory"
    ERROR = "error_memory"


@dataclass
class TradeMemory:
    """Memory of a trade execution.

    Attributes:
        symbol: Trading symbol (e.g., "AAPL")
        action: Trade action ("buy" or "sell")
        quantity: Number of shares/contracts
        price: Execution price
        pnl: Profit/loss from trade
        timestamp: When trade occurred
        market_conditions: Market state at time of trade
        strategy: Strategy used for trade
        outcome: Trade outcome description
    """

    symbol: str
    action: str
    quantity: float
    price: float
    pnl: float
    timestamp: datetime = field(default_factory=datetime.now)
    market_conditions: Dict[str, Any] = field(default_factory=dict)
    strategy: str = ""
    outcome: str = ""

    def to_text(self) -> str:
        """Convert trade memory to searchable text."""
        return (
            f"Trade {self.action} {self.symbol} "
            f"quantity {self.quantity} price {self.price} "
            f"pnl {self.pnl:.2f} strategy {self.strategy} "
            f"outcome {self.outcome}"
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "symbol": self.symbol,
            "action": self.action,
            "quantity": self.quantity,
            "price": self.price,
            "pnl": self.pnl,
            "timestamp": self.timestamp.isoformat(),
            "market_conditions": self.market_conditions,
            "strategy": self.strategy,
            "outcome": self.outcome,
        }


@dataclass
class MarketMemory:
    """Memory of market state observation.

    Attributes:
        symbol: Trading symbol
        timestamp: Observation timestamp
        price_data: OHLCV data
        indicators: Technical indicator values
        market_regime: Market regime (trending, ranging, volatile)
        sentiment: Market sentiment
        events: Relevant market events
    """

    symbol: str
    timestamp: datetime = field(default_factory=datetime.now)
    price_data: Dict[str, float] = field(default_factory=dict)
    indicators: Dict[str, float] = field(default_factory=dict)
    market_regime: str = ""
    sentiment: str = "neutral"
    events: List[str] = field(default_factory=list)

    def to_text(self) -> str:
        """Convert market memory to searchable text."""
        events_text = " ".join(self.events) if self.events else "none"
        indicator_text = " ".join(f"{k}={v:.2f}" for k, v in self.indicators.items())
        return (
            f"Market {self.symbol} regime {self.market_regime} "
            f"sentiment {self.sentiment} indicators {indicator_text} "
            f"events {events_text}"
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "symbol": self.symbol,
            "timestamp": self.timestamp.isoformat(),
            "price_data": self.price_data,
            "indicators": self.indicators,
            "market_regime": self.market_regime,
            "sentiment": self.sentiment,
            "events": self.events,
        }


@dataclass
class DecisionMemory:
    """Memory of an agent decision.

    Attributes:
        decision_type: Type of decision (trade, hold, analyze, etc.)
        context: Decision context
        reasoning: Reasoning behind decision
        expected_outcome: Expected result
        actual_outcome: Actual result
        confidence: Decision confidence (0.0 to 1.0)
        timestamp: When decision was made
        factors: Factors considered
    """

    decision_type: str
    context: str = ""
    reasoning: str = ""
    expected_outcome: str = ""
    actual_outcome: str = ""
    confidence: float = 0.5
    timestamp: datetime = field(default_factory=datetime.now)
    factors: List[str] = field(default_factory=list)

    def to_text(self) -> str:
        """Convert decision memory to searchable text."""
        factors_text = " ".join(self.factors) if self.factors else "none"
        return (
            f"Decision {self.decision_type} context {self.context} "
            f"reasoning {self.reasoning} expected {self.expected_outcome} "
            f"actual {self.actual_outcome} confidence {self.confidence} "
            f"factors {factors_text}"
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "decision_type": self.decision_type,
            "context": self.context,
            "reasoning": self.reasoning,
            "expected_outcome": self.expected_outcome,
            "actual_outcome": self.actual_outcome,
            "confidence": self.confidence,
            "timestamp": self.timestamp.isoformat(),
            "factors": self.factors,
        }


@dataclass
class ErrorMemory:
    """Memory of errors or failures.

    Attributes:
        error_type: Type of error
        error_message: Error description
        context: Context where error occurred
        recovery_action: How error was resolved
        timestamp: When error occurred
        severity: Error severity (low, medium, high, critical)
        preventability: Whether error was preventable
    """

    error_type: str
    error_message: str
    context: str = ""
    recovery_action: str = ""
    timestamp: datetime = field(default_factory=datetime.now)
    severity: str = "medium"
    preventability: str = "unknown"

    def to_text(self) -> str:
        """Convert error memory to searchable text."""
        return (
            f"Error {self.error_type} message {self.error_message} "
            f"context {self.context} recovery {self.recovery_action} "
            f"severity {self.severity} preventable {self.preventability}"
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "error_type": self.error_type,
            "error_message": self.error_message,
            "context": self.context,
            "recovery_action": self.recovery_action,
            "timestamp": self.timestamp.isoformat(),
            "severity": self.severity,
            "preventability": self.preventability,
        }


class LearningMemory:
    """Agent learning memory system with BM25-based retrieval.

    Provides storage and retrieval of trading experiences, market observations,
    decisions, and errors using BM25 text indexing for fast similarity search.

    Attributes:
        agent_id: Unique identifier for the agent
        storage_dir: Directory for persistent storage
        max_memories: Maximum number of memories to retain
        decay_enabled: Whether to enable memory decay
        decay_days: Days after which memories start decaying
    """

    def __init__(
        self,
        agent_id: str,
        storage_dir: Optional[Path] = None,
        max_memories: int = 10000,
        decay_enabled: bool = True,
        decay_days: int = 30,
        bm25_k1: float = 1.5,
        bm25_b: float = 0.75,
    ):
        """Initialize learning memory system.

        Args:
            agent_id: Unique agent identifier
            storage_dir: Directory for persistent storage
            max_memories: Maximum memories to retain
            decay_enabled: Enable memory decay
            decay_days: Days before memory starts decaying
            bm25_k1: BM25 k1 parameter
            bm25_b: BM25 b parameter
        """
        self.agent_id = agent_id
        self.storage_dir = storage_dir or Path(f".memory/{agent_id}")
        self.max_memories = max_memories
        self.decay_enabled = decay_enabled
        self.decay_days = decay_days

        self.logger = get_logger(f"memory.{agent_id}")

        # Initialize BM25 index
        self.index = BM25Index(k1=bm25_k1, b=bm25_b)

        # Memory importance calculator
        self._importance_calculators: Dict[MemoryType, Callable[[Any], float]] = {
            MemoryType.TRADE: self._calculate_trade_importance,
            MemoryType.MARKET: self._calculate_market_importance,
            MemoryType.DECISION: self._calculate_decision_importance,
            MemoryType.ERROR: self._calculate_error_importance,
        }

        # Load existing memories
        self._load()

    def _calculate_trade_importance(self, memory: TradeMemory) -> float:
        """Calculate importance score for trade memory."""
        importance = 0.5

        # High P&L trades are important
        abs_pnl = abs(memory.pnl)
        if abs_pnl > 1000:
            importance += 0.3
        elif abs_pnl > 500:
            importance += 0.2
        elif abs_pnl > 100:
            importance += 0.1

        # Unusual outcomes are important
        if "unexpected" in memory.outcome.lower() or "anomaly" in memory.outcome.lower():
            importance += 0.2

        return min(1.0, importance)

    def _calculate_market_importance(self, memory: MarketMemory) -> float:
        """Calculate importance score for market memory."""
        importance = 0.3

        # Volatile regimes are more important
        if memory.market_regime == "volatile":
            importance += 0.3
        elif memory.market_regime == "crisis":
            importance += 0.4

        # Extreme sentiment is important
        if memory.sentiment in ["extreme_fear", "extreme_greed"]:
            importance += 0.2

        # Significant events are important
        if memory.events:
            importance += min(0.3, len(memory.events) * 0.1)

        return min(1.0, importance)

    def _calculate_decision_importance(self, memory: DecisionMemory) -> float:
        """Calculate importance score for decision memory."""
        importance = 0.4

        # Confidence affects importance
        importance += memory.confidence * 0.2

        # Outcome mismatch is important (learning opportunity)
        if memory.expected_outcome and memory.actual_outcome:
            if memory.expected_outcome != memory.actual_outcome:
                importance += 0.3

        return min(1.0, importance)

    def _calculate_error_importance(self, memory: ErrorMemory) -> float:
        """Calculate importance score for error memory."""
        importance = 0.5

        # Severity affects importance
        severity_map = {
            "critical": 0.4,
            "high": 0.3,
            "medium": 0.2,
            "low": 0.1,
        }
        importance += severity_map.get(memory.severity, 0.1)

        # Preventable errors are important lessons
        if memory.preventability == "yes":
            importance += 0.1

        return min(1.0, importance)

    def _generate_doc_id(self) -> str:
        """Generate unique document ID."""
        return f"{self.agent_id}_{uuid.uuid4().hex[:16]}_{int(datetime.now().timestamp())}"

    def _apply_decay(self) -> None:
        """Apply decay to old memories."""
        if not self.decay_enabled:
            return

        if self.decay_days <= 0:
            return

        now = datetime.now()

        docs_to_remove: List[str] = []
        for doc_id, doc in list(self.index.documents.items()):
            age_days = (now - doc.timestamp).days

            if age_days > self.decay_days:
                # Reduce importance based on age
                decay_factor = 1.0 - (age_days - self.decay_days) / (self.decay_days * 2)
                decay_factor = max(0.1, decay_factor)
                doc.importance *= decay_factor

                # Remove very old, low-importance memories
                if age_days > self.decay_days * 3 and doc.importance < 0.1:
                    docs_to_remove.append(doc_id)

        # Remove expired memories
        for doc_id in docs_to_remove:
            self.index.remove_document(doc_id)

        if docs_to_remove:
            self.logger.info(f"Removed {len(docs_to_remove)} decayed memories")

    def _enforce_memory_limit(self) -> None:
        """Enforce maximum memory limit by removing least important."""
        # Check if we need to make room for a new document
        # (current count should be less than max to allow adding one)
        if len(self.index.documents) < self.max_memories:
            return

        # Sort by importance (ascending) and remove oldest low-importance memories
        memories = list(self.index.documents.items())
        memories.sort(key=lambda x: (x[1].importance, x[1].timestamp))

        # Remove enough to make room for new document
        target_count = self.max_memories - 1
        to_remove = len(memories) - target_count
        for i in range(to_remove):
            self.index.remove_document(memories[i][0])

        if to_remove > 0:
            self.logger.info(f"Enforced memory limit: removed {to_remove} memories")

    def add_trade_memory(
        self,
        symbol: str,
        action: str,
        quantity: float,
        price: float,
        pnl: float,
        strategy: str = "",
        outcome: str = "",
        market_conditions: Optional[Dict[str, Any]] = None,
    ) -> str:
        """Add a trade memory.

        Args:
            symbol: Trading symbol
            action: Trade action (buy/sell)
            quantity: Trade quantity
            price: Execution price
            pnl: Profit/loss
            strategy: Strategy used
            outcome: Outcome description
            market_conditions: Market state

        Returns:
            Memory document ID
        """
        memory = TradeMemory(
            symbol=symbol,
            action=action,
            quantity=quantity,
            price=price,
            pnl=pnl,
            strategy=strategy,
            outcome=outcome,
            market_conditions=market_conditions or {},
        )

        importance = self._calculate_trade_importance(memory)
        doc_id = self._add_memory(memory.to_text(), MemoryType.TRADE, importance, memory.to_dict())
        self.logger.debug(f"Added trade memory: {symbol} {action} P&L={pnl}")
        return doc_id

    def add_market_memory(
        self,
        symbol: str,
        price_data: Optional[Dict[str, float]] = None,
        indicators: Optional[Dict[str, float]] = None,
        market_regime: str = "",
        sentiment: str = "neutral",
        events: Optional[List[str]] = None,
    ) -> str:
        """Add a market observation memory.

        Args:
            symbol: Trading symbol
            price_data: OHLCV data
            indicators: Technical indicators
            market_regime: Market regime
            sentiment: Market sentiment
            events: Market events

        Returns:
            Memory document ID
        """
        memory = MarketMemory(
            symbol=symbol,
            price_data=price_data or {},
            indicators=indicators or {},
            market_regime=market_regime,
            sentiment=sentiment,
            events=events or [],
        )

        importance = self._calculate_market_importance(memory)
        doc_id = self._add_memory(memory.to_text(), MemoryType.MARKET, importance, memory.to_dict())
        self.logger.debug(f"Added market memory: {symbol} regime={market_regime}")
        return doc_id

    def add_decision_memory(
        self,
        decision_type: str,
        context: str = "",
        reasoning: str = "",
        expected_outcome: str = "",
        actual_outcome: str = "",
        confidence: float = 0.5,
        factors: Optional[List[str]] = None,
    ) -> str:
        """Add a decision memory.

        Args:
            decision_type: Type of decision
            context: Decision context
            reasoning: Reasoning
            expected_outcome: Expected result
            actual_outcome: Actual result
            confidence: Confidence level
            factors: Decision factors

        Returns:
            Memory document ID
        """
        memory = DecisionMemory(
            decision_type=decision_type,
            context=context,
            reasoning=reasoning,
            expected_outcome=expected_outcome,
            actual_outcome=actual_outcome,
            confidence=confidence,
            factors=factors or [],
        )

        importance = self._calculate_decision_importance(memory)
        doc_id = self._add_memory(
            memory.to_text(), MemoryType.DECISION, importance, memory.to_dict()
        )
        self.logger.debug(f"Added decision memory: {decision_type}")
        return doc_id

    def add_error_memory(
        self,
        error_type: str,
        error_message: str,
        context: str = "",
        recovery_action: str = "",
        severity: str = "medium",
        preventability: str = "unknown",
    ) -> str:
        """Add an error memory.

        Args:
            error_type: Error type
            error_message: Error description
            context: Error context
            recovery_action: Recovery action
            severity: Error severity
            preventability: Whether preventable

        Returns:
            Memory document ID
        """
        memory = ErrorMemory(
            error_type=error_type,
            error_message=error_message,
            context=context,
            recovery_action=recovery_action,
            severity=severity,
            preventability=preventability,
        )

        importance = self._calculate_error_importance(memory)
        doc_id = self._add_memory(memory.to_text(), MemoryType.ERROR, importance, memory.to_dict())
        self.logger.debug(f"Added error memory: {error_type} severity={severity}")
        return doc_id

    def _add_memory(
        self, content: str, memory_type: MemoryType, importance: float, metadata: Dict[str, Any]
    ) -> str:
        """Internal method to add a memory document.

        Args:
            content: Text content for indexing
            memory_type: Type of memory
            importance: Importance score
            metadata: Structured metadata

        Returns:
            Document ID
        """
        # Apply decay and enforce limits
        self._apply_decay()
        self._enforce_memory_limit()

        # Create document
        doc_id = self._generate_doc_id()
        doc = MemoryDocument(
            doc_id=doc_id,
            content=content,
            memory_type=memory_type.value,
            importance=importance,
            metadata=metadata,
        )

        # Add to index
        self.index.add_document(doc)

        return doc_id

    def search_similar_trades(
        self,
        symbol: str = "",
        strategy: str = "",
        min_pnl: Optional[float] = None,
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        """Search for similar trades.

        Args:
            symbol: Symbol to match
            strategy: Strategy to match
            min_pnl: Minimum P&L filter
            top_k: Number of results

        Returns:
            List of matching trade memories with scores
        """
        query = f"trade {symbol} {strategy}".strip()
        results = self.index.search(query, top_k=top_k * 2, memory_type=MemoryType.TRADE.value)

        filtered_results = []
        for doc, score in results:
            metadata = doc.metadata
            if min_pnl is not None:
                if metadata.get("pnl", 0) < min_pnl:
                    continue
            filtered_results.append({
                "doc_id": doc.doc_id,
                "score": score,
                "memory_type": doc.memory_type,
                "timestamp": doc.timestamp.isoformat(),
                "importance": doc.importance,
                "data": metadata,
            })

        return filtered_results[:top_k]

    def search_similar_market_states(
        self,
        symbol: str = "",
        regime: str = "",
        indicators: Optional[Dict[str, float]] = None,
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        """Search for similar market states.

        Args:
            symbol: Symbol to match
            regime: Market regime to match
            indicators: Indicators to match
            top_k: Number of results

        Returns:
            List of matching market memories with scores
        """
        indicator_text = ""
        if indicators:
            indicator_text = " ".join(f"{k}={v}" for k, v in indicators.items())

        query = f"market {symbol} {regime} {indicator_text}".strip()
        results = self.index.search(query, top_k=top_k, memory_type=MemoryType.MARKET.value)

        return [
            {
                "doc_id": doc.doc_id,
                "score": score,
                "memory_type": doc.memory_type,
                "timestamp": doc.timestamp.isoformat(),
                "importance": doc.importance,
                "data": doc.metadata,
            }
            for doc, score in results
        ]

    def get_decision_suggestions(
        self,
        context: str,
        decision_type: str = "",
        top_k: int = 3,
    ) -> List[Dict[str, Any]]:
        """Get decision suggestions based on similar past decisions.

        Args:
            context: Current decision context
            decision_type: Type of decision
            top_k: Number of suggestions

        Returns:
            List of relevant decision memories with outcomes
        """
        query = f"{decision_type} {context}".strip()
        results = self.index.search(query, top_k=top_k, memory_type=MemoryType.DECISION.value)

        suggestions = []
        for doc, score in results:
            metadata = doc.metadata
            suggestions.append({
                "doc_id": doc.doc_id,
                "score": score,
                "decision_type": metadata.get("decision_type"),
                "context": metadata.get("context"),
                "reasoning": metadata.get("reasoning"),
                "expected_outcome": metadata.get("expected_outcome"),
                "actual_outcome": metadata.get("actual_outcome"),
                "confidence": metadata.get("confidence"),
                "timestamp": doc.timestamp.isoformat(),
            })

        return suggestions

    def get_error_lessons(
        self,
        error_type: str = "",
        context: str = "",
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        """Get lessons from past errors.

        Args:
            error_type: Type of error to look for
            context: Current context
            top_k: Number of lessons

        Returns:
            List of relevant error memories with recovery actions
        """
        # Replace underscores with spaces for better token matching
        # (e.g., "api_error" becomes "api error" to match tokenized content)
        error_type_normalized = error_type.replace("_", " ") if error_type else ""
        query = f"{error_type_normalized} {context}".strip()
        results = self.index.search(query, top_k=top_k, memory_type=MemoryType.ERROR.value)

        return [
            {
                "doc_id": doc.doc_id,
                "score": score,
                "error_type": doc.metadata.get("error_type"),
                "error_message": doc.metadata.get("error_message"),
                "context": doc.metadata.get("context"),
                "recovery_action": doc.metadata.get("recovery_action"),
                "severity": doc.metadata.get("severity"),
                "preventability": doc.metadata.get("preventability"),
                "timestamp": doc.timestamp.isoformat(),
            }
            for doc, score in results
        ]

    def get_memory_stats(self) -> Dict[str, Any]:
        """Get memory system statistics.

        Returns:
            Dictionary with memory statistics
        """
        index_stats = self.index.get_stats()

        # Calculate additional stats
        total_accesses = sum(doc.access_count for doc in self.index.documents.values())
        avg_importance = (
            sum(doc.importance for doc in self.index.documents.values()) / len(self.index.documents)
            if self.index.documents
            else 0.0
        )

        return {
            "agent_id": self.agent_id,
            "total_memories": len(self.index.documents),
            "memory_types": index_stats.get("memory_types", {}),
            "avg_importance": avg_importance,
            "total_accesses": total_accesses,
            "decay_enabled": self.decay_enabled,
            "decay_days": self.decay_days,
            **index_stats,
        }

    def update_memory_importance(self, doc_id: str, importance: float) -> bool:
        """Update the importance of a memory.

        Args:
            doc_id: Document ID
            importance: New importance value (0.0 to 1.0)

        Returns:
            True if updated, False if not found
        """
        importance = max(0.0, min(1.0, importance))
        return self.index.update_document(doc_id, importance=importance)

    def mark_important(self, doc_id: str) -> bool:
        """Mark a memory as important (max importance).

        Args:
            doc_id: Document ID

        Returns:
            True if marked, False if not found
        """
        return self.update_memory_importance(doc_id, 1.0)

    def delete_memory(self, doc_id: str) -> bool:
        """Delete a specific memory.

        Args:
            doc_id: Document ID

        Returns:
            True if deleted, False if not found
        """
        return self.index.remove_document(doc_id)

    def clear_all_memories(self) -> None:
        """Clear all memories."""
        for doc_id in list(self.index.documents.keys()):
            self.index.remove_document(doc_id)
        self.logger.info("All memories cleared")

    def save(self) -> None:
        """Save memories to disk."""
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        index_path = self.storage_dir / "index.pkl"
        self.index.save(index_path)
        self.logger.debug(f"Saved {len(self.index.documents)} memories to {index_path}")

    def _load(self) -> bool:
        """Load memories from disk.

        Returns:
            True if loaded, False otherwise
        """
        index_path = self.storage_dir / "index.pkl"
        if self.index.load(index_path):
            self.logger.info(f"Loaded {len(self.index.documents)} memories from {index_path}")
            return True
        return False
