"""
Learning and adaptation capabilities for the AI agent
Handles experience tracking, pattern recognition, and strategy optimization
"""
from typing import Dict, List, Any, Optional
import json
from datetime import datetime
from pathlib import Path
import numpy as np
from dataclasses import dataclass
import pickle

@dataclass
class Experience:
    action_type: str
    context: Dict[str, Any]
    result: Dict[str, Any]
    timestamp: datetime
    feedback: Optional[str] = None

class LearningManager:
    def __init__(self, workspace_path: str):
        self.workspace_path = Path(workspace_path)
        self.experience_path = self.workspace_path / '.agent' / 'experience.pkl'
        self.experiences: List[Experience] = []
        self.patterns: Dict[str, Any] = {}
        self.load_experiences()
        
    def record_experience(self, 
                       action_type: str,
                       context: Dict[str, Any],
                       result: Dict[str, Any],
                       feedback: Optional[str] = None):
        """Record a new experience"""
        experience = Experience(
            action_type=action_type,
            context=context,
            result=result,
            timestamp=datetime.now(),
            feedback=feedback
        )
        self.experiences.append(experience)
        self._update_patterns()
        self.save_experiences()
        
    def _update_patterns(self):
        """Update recognized patterns from experiences"""
        patterns = {
            "success_patterns": self._analyze_success_patterns(),
            "error_patterns": self._analyze_error_patterns(),
            "context_patterns": self._analyze_context_patterns()
        }
        self.patterns.update(patterns)
        
    def _analyze_success_patterns(self) -> Dict[str, Any]:
        """Analyze patterns in successful actions"""
        success_patterns = {}
        
        for exp in self.experiences:
            if exp.result.get("success", False):
                # Record successful action patterns
                action_type = exp.action_type
                if action_type not in success_patterns:
                    success_patterns[action_type] = {
                        "count": 0,
                        "contexts": [],
                        "common_factors": set()
                    }
                
                success_patterns[action_type]["count"] += 1
                success_patterns[action_type]["contexts"].append(exp.context)
                
                # Find common factors in successful contexts
                context_factors = self._extract_context_factors(exp.context)
                if success_patterns[action_type]["common_factors"]:
                    success_patterns[action_type]["common_factors"] &= context_factors
                else:
                    success_patterns[action_type]["common_factors"] = context_factors
                    
        return success_patterns
        
    def _analyze_error_patterns(self) -> Dict[str, Any]:
        """Analyze patterns in failed actions"""
        error_patterns = {}
        
        for exp in self.experiences:
            if not exp.result.get("success", False):
                error = exp.result.get("error", "unknown_error")
                if error not in error_patterns:
                    error_patterns[error] = {
                        "count": 0,
                        "contexts": [],
                        "solutions": set()
                    }
                
                error_patterns[error]["count"] += 1
                error_patterns[error]["contexts"].append(exp.context)
                
                # Record successful solutions if this error was later resolved
                if exp.feedback and "resolved" in exp.feedback.lower():
                    error_patterns[error]["solutions"].add(exp.feedback)
                    
        return error_patterns
        
    def _analyze_context_patterns(self) -> Dict[str, Any]:
        """Analyze patterns in different contexts"""
        context_patterns = {}
        
        for exp in self.experiences:
            for key, value in exp.context.items():
                if key not in context_patterns:
                    context_patterns[key] = {
                        "values": set(),
                        "success_rate": {},
                        "common_actions": set()
                    }
                
                context_patterns[key]["values"].add(str(value))
                
                # Track success rate for this context value
                success = exp.result.get("success", False)
                if str(value) not in context_patterns[key]["success_rate"]:
                    context_patterns[key]["success_rate"][str(value)] = {
                        "success": 0,
                        "total": 0
                    }
                
                context_patterns[key]["success_rate"][str(value)]["total"] += 1
                if success:
                    context_patterns[key]["success_rate"][str(value)]["success"] += 1
                    context_patterns[key]["common_actions"].add(exp.action_type)
                    
        return context_patterns
        
    def _extract_context_factors(self, context: Dict) -> set:
        """Extract key factors from a context"""
        factors = set()
        
        def extract_value(value: Any):
            if isinstance(value, (str, int, float, bool)):
                return str(value)
            elif isinstance(value, (list, tuple)):
                return tuple(extract_value(v) for v in value)
            elif isinstance(value, dict):
                return tuple((k, extract_value(v)) for k, v in sorted(value.items()))
            return str(type(value))
            
        for key, value in context.items():
            factors.add((key, extract_value(value)))
            
        return factors
        
    def get_strategy_recommendations(self, 
                                 context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Get recommended strategies based on past experiences"""
        recommendations = []
        
        # Find similar contexts
        similar_experiences = self._find_similar_experiences(context)
        
        # Analyze successful strategies
        success_strategies = {}
        for exp in similar_experiences:
            if exp.result.get("success", False):
                strategy = {
                    "action_type": exp.action_type,
                    "context_match": self._calculate_context_similarity(context, exp.context),
                    "timestamp": exp.timestamp
                }
                
                if exp.action_type not in success_strategies:
                    success_strategies[exp.action_type] = []
                success_strategies[exp.action_type].append(strategy)
                
        # Sort and format recommendations
        for action_type, strategies in success_strategies.items():
            # Calculate success rate
            total_similar = len([e for e in similar_experiences 
                               if e.action_type == action_type])
            success_count = len(strategies)
            success_rate = success_count / total_similar if total_similar > 0 else 0
            
            # Get most recent successful strategy
            most_recent = max(strategies, key=lambda x: x["timestamp"])
            
            recommendations.append({
                "action_type": action_type,
                "success_rate": success_rate,
                "context_similarity": most_recent["context_match"],
                "last_success": most_recent["timestamp"],
                "total_attempts": total_similar
            })
            
        # Sort by success rate and context similarity
        recommendations.sort(
            key=lambda x: (x["success_rate"], x["context_similarity"]),
            reverse=True
        )
        
        return recommendations
        
    def _find_similar_experiences(self, 
                              context: Dict[str, Any],
                              threshold: float = 0.7) -> List[Experience]:
        """Find experiences with similar contexts"""
        similar_experiences = []
        
        for exp in self.experiences:
            similarity = self._calculate_context_similarity(context, exp.context)
            if similarity >= threshold:
                similar_experiences.append(exp)
                
        return similar_experiences
        
    def _calculate_context_similarity(self, 
                                  context1: Dict[str, Any],
                                  context2: Dict[str, Any]) -> float:
        """Calculate similarity between two contexts"""
        # Get all keys
        all_keys = set(context1.keys()) | set(context2.keys())
        
        if not all_keys:
            return 0.0
            
        similarity_scores = []
        
        for key in all_keys:
            val1 = context1.get(key)
            val2 = context2.get(key)
            
            # Calculate similarity for this key
            if val1 is None or val2 is None:
                similarity_scores.append(0.0)
            elif isinstance(val1, (int, float)) and isinstance(val2, (int, float)):
                # Numeric similarity
                max_val = max(abs(val1), abs(val2))
                if max_val == 0:
                    similarity_scores.append(1.0)
                else:
                    similarity_scores.append(1 - abs(val1 - val2) / max_val)
            elif isinstance(val1, str) and isinstance(val2, str):
                # String similarity
                from difflib import SequenceMatcher
                similarity_scores.append(
                    SequenceMatcher(None, val1, val2).ratio()
                )
            else:
                # Basic equality check for other types
                similarity_scores.append(1.0 if val1 == val2 else 0.0)
                
        return sum(similarity_scores) / len(similarity_scores)
        
    def save_experiences(self):
        """Save experiences to file"""
        self.experience_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.experience_path, 'wb') as f:
            pickle.dump(self.experiences, f)
            
    def load_experiences(self):
        """Load experiences from file"""
        if self.experience_path.exists():
            with open(self.experience_path, 'rb') as f:
                self.experiences = pickle.load(f)
            self._update_patterns()
