"""Learning Manager for Bat Genie's self-improvement capabilities."""
import json
import os
from datetime import datetime
from typing import Dict, List, Optional, Any
from pathlib import Path
import logging
from dataclasses import dataclass, asdict
import traceback

logger = logging.getLogger(__name__)

@dataclass
class LearningEvent:
    """Represents a learning event from success, failure, or correction."""
    timestamp: str
    event_type: str  # 'error', 'correction', 'success'
    context: Dict[str, Any]
    description: str
    solution: Optional[str] = None
    source_file: Optional[str] = None
    line_number: Optional[int] = None
    stack_trace: Optional[str] = None

class LearningManager:
    """Manages Bat Genie's learning and self-improvement capabilities."""
    
    def __init__(self, workspace_root: str):
        """Initialize the Learning Manager."""
        self.workspace_root = Path(workspace_root)
        self.learning_db_path = self.workspace_root / '.cursor' / 'learning_db.json'
        self.learning_db_path.parent.mkdir(parents=True, exist_ok=True)
        self.load_learning_db()

    def load_learning_db(self):
        """Load the learning database from disk."""
        if self.learning_db_path.exists():
            try:
                with open(self.learning_db_path, 'r') as f:
                    self.learning_db = json.load(f)
            except json.JSONDecodeError:
                logger.error("Error loading learning database, creating new one")
                self.learning_db = {'events': [], 'patterns': {}}
        else:
            self.learning_db = {'events': [], 'patterns': {}}

    def save_learning_db(self):
        """Save the learning database to disk."""
        with open(self.learning_db_path, 'w') as f:
            json.dump(self.learning_db, f, indent=2)

    def record_error(self, error: Exception, context: Dict[str, Any]):
        """Record an error and its context for learning."""
        tb = traceback.extract_tb(error.__traceback__)
        if tb:
            filename, line_number, _, _ = tb[-1]
        else:
            filename, line_number = None, None

        event = LearningEvent(
            timestamp=datetime.now().isoformat(),
            event_type='error',
            context=context,
            description=str(error),
            source_file=filename,
            line_number=line_number,
            stack_trace=traceback.format_exc()
        )
        
        self.learning_db['events'].append(asdict(event))
        self._analyze_error_pattern(event)
        self.save_learning_db()

    def record_correction(self, original: str, correction: str, context: Dict[str, Any]):
        """Record a correction made by the user."""
        event = LearningEvent(
            timestamp=datetime.now().isoformat(),
            event_type='correction',
            context=context,
            description=original,
            solution=correction
        )
        
        self.learning_db['events'].append(asdict(event))
        self._analyze_correction_pattern(event)
        self.save_learning_db()

    def record_success(self, action: str, context: Dict[str, Any]):
        """Record a successful action for reinforcement."""
        event = LearningEvent(
            timestamp=datetime.now().isoformat(),
            event_type='success',
            context=context,
            description=action
        )
        
        self.learning_db['events'].append(asdict(event))
        self._analyze_success_pattern(event)
        self.save_learning_db()

    def _analyze_error_pattern(self, event: LearningEvent):
        """Analyze error patterns to identify common issues."""
        error_type = type(event.description).__name__
        if error_type not in self.learning_db['patterns']:
            self.learning_db['patterns'][error_type] = {
                'count': 0,
                'contexts': [],
                'solutions': []
            }
        
        pattern = self.learning_db['patterns'][error_type]
        pattern['count'] += 1
        pattern['contexts'].append(event.context)
        
        # Look for similar errors that were solved
        similar_events = [
            e for e in self.learning_db['events']
            if e['event_type'] == 'correction' and
            self._calculate_similarity(e['description'], event.description) > 0.7
        ]
        
        if similar_events:
            pattern['solutions'].extend([e['solution'] for e in similar_events])

    def _analyze_correction_pattern(self, event: LearningEvent):
        """Analyze correction patterns to improve future suggestions."""
        correction_type = self._categorize_correction(event.description, event.solution)
        if correction_type not in self.learning_db['patterns']:
            self.learning_db['patterns'][correction_type] = {
                'count': 0,
                'examples': []
            }
        
        pattern = self.learning_db['patterns'][correction_type]
        pattern['count'] += 1
        pattern['examples'].append({
            'original': event.description,
            'correction': event.solution,
            'context': event.context
        })

    def _analyze_success_pattern(self, event: LearningEvent):
        """Analyze successful patterns to reinforce good practices."""
        action_type = self._categorize_action(event.description)
        if action_type not in self.learning_db['patterns']:
            self.learning_db['patterns'][action_type] = {
                'count': 0,
                'successful_contexts': []
            }
        
        pattern = self.learning_db['patterns'][action_type]
        pattern['count'] += 1
        pattern['successful_contexts'].append(event.context)

    def get_similar_cases(self, query: str, context: Dict[str, Any], n: int = 5) -> List[Dict]:
        """Find similar cases from past experiences."""
        similar_cases = []
        for event in self.learning_db['events']:
            similarity = self._calculate_similarity(event['description'], query)
            if similarity > 0.5:  # Threshold for similarity
                similar_cases.append({
                    'event': event,
                    'similarity': similarity
                })
        
        # Sort by similarity and return top N
        similar_cases.sort(key=lambda x: x['similarity'], reverse=True)
        return similar_cases[:n]

    def get_learned_patterns(self) -> Dict[str, Any]:
        """Get all learned patterns and their frequencies."""
        return self.learning_db['patterns']

    def _calculate_similarity(self, text1: str, text2: str) -> float:
        """Calculate similarity between two texts."""
        # Implement text similarity calculation
        # For now, using a simple implementation
        words1 = set(text1.lower().split())
        words2 = set(text2.lower().split())
        intersection = words1.intersection(words2)
        union = words1.union(words2)
        return len(intersection) / len(union) if union else 0

    def _categorize_correction(self, original: str, correction: str) -> str:
        """Categorize the type of correction made."""
        # Implement correction categorization
        # For now, using a simple implementation
        if 'security' in correction.lower():
            return 'security_fix'
        elif 'performance' in correction.lower():
            return 'performance_optimization'
        elif 'bug' in correction.lower():
            return 'bug_fix'
        else:
            return 'general_improvement'

    def _categorize_action(self, action: str) -> str:
        """Categorize the type of successful action."""
        # Implement action categorization
        # For now, using a simple implementation
        if 'test' in action.lower():
            return 'testing'
        elif 'refactor' in action.lower():
            return 'refactoring'
        elif 'fix' in action.lower():
            return 'bug_fixing'
        else:
            return 'development' 