# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
# learningEngine.py
# Learning engine for adaptive Dynamic mode

import logging
from typing import Dict, Any, List
from datetime import datetime, timezone

logger = logging.getLogger(__name__)

class LearningEngine:
    """Learns from feedback and adapts future behavior"""
    
    def __init__(self):
        self.strategies = {}
        self.feedbackHistory = []
    
    def learnFromFeedback(self, feedback: Dict[str, Any], context: Any, taskIntent: Dict[str, Any]):
        """Learns from feedback and updates strategies - works on TASK level, not workflow level"""
        try:
            # Store feedback
            self.feedbackHistory.append({
                "feedback": feedback,
                "context": self._serializeContext(context),
                "taskIntent": taskIntent,  # Changed from intent to taskIntent
                "timestamp": datetime.now(timezone.utc).timestamp()
            })
            
            # Update strategies based on feedback (using taskIntent)
            self._updateStrategies(feedback, taskIntent)
            
            # Normalize scores for safe logging
            _qs = feedback.get('qualityScore', 0.0)
            _im = feedback.get('intentMatchScore', 0.0)
            try:
                _qs = float(0.0 if _qs is None else _qs)
            except Exception:
                _qs = 0.0
            try:
                _im = float(0.0 if _im is None else _im)
            except Exception:
                _im = 0.0
            logger.info(
                f"Learning from feedback: {feedback.get('actionAttempted', 'unknown')} - "
                f"Quality: {_qs:.2f}, Intent Match: {_im:.2f}"
            )
            
        except Exception as e:
            logger.error(f"Error learning from feedback: {str(e)}")
    
    def getImprovedStrategy(self, context: Any, taskIntent: Dict[str, Any]) -> Dict[str, Any]:
        """Returns improved strategy based on learning - works on TASK level"""
        try:
            # Get strategy key based on taskIntent
            strategyKey = self._getStrategyKey(taskIntent)
            
            # Get existing strategy or create default
            if strategyKey in self.strategies:
                strategy = self.strategies[strategyKey]
                logger.info(f"Using learned strategy for {strategyKey}: {strategy}")
                return strategy
            else:
                # Create default strategy
                defaultStrategy = self._createDefaultStrategy(taskIntent)
                self.strategies[strategyKey] = defaultStrategy
                logger.info(f"Created default strategy for {strategyKey}")
                return defaultStrategy
                
        except Exception as e:
            logger.error(f"Error getting improved strategy: {str(e)}")
            return self._createDefaultStrategy(taskIntent)
    
    def _updateStrategies(self, feedback: Dict[str, Any], taskIntent: Dict[str, Any]):
        """Updates strategies based on feedback - works on TASK level"""
        strategyKey = self._getStrategyKey(taskIntent)
        actionAttempted = feedback.get('actionAttempted', 'unknown')
        # Coerce possibly None or non-numeric to floats
        qs_raw = feedback.get('qualityScore', 0.0)
        im_raw = feedback.get('intentMatchScore', 0.0)
        try:
            qualityScore = float(0.0 if qs_raw is None else qs_raw)
        except Exception:
            qualityScore = 0.0
        try:
            intentMatchScore = float(0.0 if im_raw is None else im_raw)
        except Exception:
            intentMatchScore = 0.0
        
        # Get or create strategy
        if strategyKey not in self.strategies:
            self.strategies[strategyKey] = self._createDefaultStrategy(taskIntent)
        
        strategy = self.strategies[strategyKey]
        
        # Update based on success/failure
        if qualityScore > 0.7 and intentMatchScore > 0.7:
            # Successful action - reinforce it
            if 'successfulActions' not in strategy:
                strategy['successfulActions'] = []
            if actionAttempted not in strategy['successfulActions']:
                strategy['successfulActions'].append(actionAttempted)
            strategy['successRate'] = min(strategy.get('successRate', 0.5) + 0.1, 1.0)
            logger.info(f"Reinforced successful action: {actionAttempted}")
            
        elif qualityScore < 0.3 or intentMatchScore < 0.3:
            # Failed action - avoid it
            if 'failedActions' not in strategy:
                strategy['failedActions'] = []
            if actionAttempted not in strategy['failedActions']:
                strategy['failedActions'].append(actionAttempted)
            strategy['successRate'] = max(strategy.get('successRate', 0.5) - 0.1, 0.0)
            logger.info(f"Marked failed action to avoid: {actionAttempted}")
        
        # Update last modified
        strategy['lastModified'] = datetime.now(timezone.utc).timestamp()
    
    def _getStrategyKey(self, taskIntent: Dict[str, Any]) -> str:
        """Gets strategy key based on taskIntent"""
        dataType = taskIntent.get('dataType', 'unknown')
        expectedFormats = taskIntent.get('expectedFormats', [])
        formatKey = '_'.join(expectedFormats) if expectedFormats else 'unknown'
        return f"{dataType}_{formatKey}"
    
    def _createDefaultStrategy(self, taskIntent: Dict[str, Any]) -> Dict[str, Any]:
        """Creates a default strategy for the taskIntent"""
        dataType = taskIntent.get('dataType', 'unknown')
        expectedFormats = taskIntent.get('expectedFormats', [])
        formatStr = ', '.join(expectedFormats) if expectedFormats else 'any'
        formatKey = '_'.join(expectedFormats) if expectedFormats else 'unknown'
        
        # Create strategy based on intent type
        if dataType == 'numbers':
            return {
                'strategyId': f"numbers_{formatKey}",
                'successfulActions': [],
                'failedActions': [],
                'successRate': 0.5,
                'lastModified': datetime.now(timezone.utc).timestamp(),
                'recommendedPrompt': f"Deliver {dataType} data in {formatStr} format. Provide actual numbers, not code to generate them.",
                'avoidPrompt': "Do not ask AI to write code when user wants data. Deliver the data directly."
            }
        elif dataType == 'text':
            return {
                'strategyId': f"text_{formatKey}",
                'successfulActions': [],
                'failedActions': [],
                'successRate': 0.5,
                'lastModified': datetime.now(timezone.utc).timestamp(),
                'recommendedPrompt': f"Generate {dataType} content in {formatStr} format.",
                'avoidPrompt': "Ensure content is readable and well-structured."
            }
        elif dataType == 'documents':
            return {
                'strategyId': f"documents_{formatKey}",
                'successfulActions': [],
                'failedActions': [],
                'successRate': 0.5,
                'lastModified': datetime.now(timezone.utc).timestamp(),
                'recommendedPrompt': f"Create {dataType} in {formatStr} format with proper structure.",
                'avoidPrompt': "Ensure document is properly formatted and organized."
            }
        else:
            return {
                'strategyId': f"unknown_{formatKey}",
                'successfulActions': [],
                'failedActions': [],
                'successRate': 0.5,
                'lastModified': datetime.now(timezone.utc).timestamp(),
                'recommendedPrompt': f"Deliver {dataType} content in {formatStr} format.",
                'avoidPrompt': "Ensure content matches user requirements."
            }
    
    def _serializeContext(self, context: Any) -> Dict[str, Any]:
        """Serializes context for storage - task-level context"""
        try:
            taskObjective = ""
            if hasattr(context, 'taskStep') and context.taskStep:
                if hasattr(context.taskStep, 'objective'):
                    taskObjective = context.taskStep.objective
                elif isinstance(context.taskStep, dict):
                    taskObjective = context.taskStep.get('objective', '')
            
            return {
                "taskObjective": taskObjective,
                "workflowId": getattr(context, 'workflowId', ''),
                "availableDocuments": getattr(context, 'availableDocuments', [])
            }
        except Exception:
            return {}
    
    def getLearningSummary(self) -> Dict[str, Any]:
        """Gets a summary of what has been learned"""
        return {
            "totalStrategies": len(self.strategies),
            "totalFeedback": len(self.feedbackHistory),
            "strategies": list(self.strategies.keys()),
            "averageSuccessRate": sum(s.get('successRate', 0) for s in self.strategies.values()) / max(len(self.strategies), 1)
        }
