"""Learning Manager - Manages workflow learning states and transitions"""
import logging
from typing import Dict, Optional, List
from dataclasses import dataclass, field
from datetime import datetime
from ..models.workflow_graph import LearningState, Workflow

logger = logging.getLogger(__name__)

@dataclass
class WorkflowStats:
    """Statistics for a workflow"""
    workflow_id: str
    learning_state: LearningState
    observation_count: int = 0
    execution_count: int = 0
    success_count: int = 0
    failure_count: int = 0
    last_execution: Optional[datetime] = None
    confidence_scores: List[float] = field(default_factory=list)
    created_at: datetime = field(default_factory=datetime.now)
    
    @property
    def success_rate(self) -> float:
        """Calculate success rate"""
        if self.execution_count == 0:
            return 0.0
        return self.success_count / self.execution_count
    
    @property
    def avg_confidence(self) -> float:
        """Calculate average confidence"""
        if not self.confidence_scores:
            return 0.0
        return sum(self.confidence_scores) / len(self.confidence_scores)

class LearningManager:
    """Manages workflow learning states and transitions"""
    
    def __init__(self):
        self.workflows: Dict[str, WorkflowStats] = {}
        logger.info("LearningManager initialized")
    
    def register_workflow(self, workflow: Workflow) -> None:
        """Register a new workflow for learning"""
        wf_id = workflow.workflow_id
        if wf_id not in self.workflows:
            self.workflows[wf_id] = WorkflowStats(
                workflow_id=wf_id,
                learning_state=workflow.learning_state
            )
            logger.info(f"Registered workflow: {wf_id} (state={workflow.learning_state})")
    
    def record_observation(self, workflow_id: str) -> None:
        """Record an observation of the workflow"""
        if workflow_id not in self.workflows:
            logger.warning(f"Unknown workflow: {workflow_id}")
            return
        
        stats = self.workflows[workflow_id]
        stats.observation_count += 1
        logger.debug(f"Observation recorded for {workflow_id} (count={stats.observation_count})")
        
        self._check_state_transition(workflow_id)
    
    def record_execution(self, workflow_id: str, success: bool, confidence: float) -> None:
        """Record an execution result"""
        if workflow_id not in self.workflows:
            logger.warning(f"Unknown workflow: {workflow_id}")
            return
        
        stats = self.workflows[workflow_id]
        stats.execution_count += 1
        stats.last_execution = datetime.now()
        stats.confidence_scores.append(confidence)
        
        if success:
            stats.success_count += 1
        else:
            stats.failure_count += 1
        
        logger.info(
            f"Execution recorded for {workflow_id}: "
            f"success={success}, confidence={confidence:.2f}, "
            f"success_rate={stats.success_rate:.2f}"
        )
        
        self._check_state_transition(workflow_id)
    
    def _check_state_transition(self, workflow_id: str) -> None:
        """Check if workflow should transition to next learning state"""
        stats = self.workflows[workflow_id]
        current_state = stats.learning_state
        new_state = None
        reason = ""
        
        if current_state == LearningState.OBSERVATION:
            if self._can_transition_to_coaching(stats):
                new_state = LearningState.COACHING
                reason = f"5+ observations ({stats.observation_count}), avg confidence > 0.90"
        
        elif current_state == LearningState.COACHING:
            if self._can_transition_to_auto_candidate(stats):
                new_state = LearningState.AUTO_CANDIDATE
                reason = f"10+ assists ({stats.execution_count}), success rate > 0.90"
        
        elif current_state == LearningState.AUTO_CANDIDATE:
            if self._can_transition_to_auto_confirmed(stats):
                new_state = LearningState.AUTO_CONFIRMED
                reason = f"20+ executions ({stats.execution_count}), success rate > 0.95"
        
        elif current_state == LearningState.AUTO_CONFIRMED:
            if self._should_rollback(stats):
                new_state = LearningState.COACHING
                reason = f"Confidence dropped below 0.90 (avg={stats.avg_confidence:.2f})"
        
        if new_state:
            self._transition_state(workflow_id, new_state, reason)
    
    def _can_transition_to_coaching(self, stats: WorkflowStats) -> bool:
        """Check if can transition from OBSERVATION to COACHING"""
        return (
            stats.observation_count >= 5 and
            stats.avg_confidence >= 0.90
        )
    
    def _can_transition_to_auto_candidate(self, stats: WorkflowStats) -> bool:
        """Check if can transition from COACHING to AUTO_CANDIDATE"""
        return (
            stats.execution_count >= 10 and
            stats.success_rate >= 0.90
        )
    
    def _can_transition_to_auto_confirmed(self, stats: WorkflowStats) -> bool:
        """Check if can transition from AUTO_CANDIDATE to AUTO_CONFIRMED"""
        return (
            stats.execution_count >= 20 and
            stats.success_rate >= 0.95
        )
    
    def _should_rollback(self, stats: WorkflowStats) -> bool:
        """Check if should rollback from AUTO_CONFIRMED to COACHING"""
        recent_scores = stats.confidence_scores[-10:] if len(stats.confidence_scores) >= 10 else stats.confidence_scores
        if not recent_scores:
            return False
        
        recent_avg = sum(recent_scores) / len(recent_scores)
        return recent_avg < 0.90
    
    def _transition_state(self, workflow_id: str, new_state: LearningState, reason: str) -> None:
        """Transition workflow to new learning state"""
        stats = self.workflows[workflow_id]
        old_state = stats.learning_state
        stats.learning_state = new_state
        
        logger.info(
            f"State transition for {workflow_id}: "
            f"{old_state.value} → {new_state.value} "
            f"(reason: {reason})"
        )
    
    def get_workflow_state(self, workflow_id: str) -> Optional[LearningState]:
        """Get current learning state of workflow"""
        if workflow_id in self.workflows:
            return self.workflows[workflow_id].learning_state
        return None
    
    def get_workflow_stats(self, workflow_id: str) -> Optional[WorkflowStats]:
        """Get statistics for workflow"""
        return self.workflows.get(workflow_id)
    
    def should_execute_automatically(self, workflow_id: str) -> bool:
        """Check if workflow should execute automatically"""
        state = self.get_workflow_state(workflow_id)
        return state in [LearningState.AUTO_CANDIDATE, LearningState.AUTO_CONFIRMED]
    
    def should_ask_confirmation(self, workflow_id: str) -> bool:
        """Check if should ask user confirmation before execution"""
        state = self.get_workflow_state(workflow_id)
        return state == LearningState.COACHING