app/agents/study_orchestrator.py

"""
Study Orchestrator Agent

The central coordination agent that orchestrates all learning agents:
1. Coordinates DoubtPredictorAgent for proactive doubt capture
2. Manages BehavioralAgent for signal analysis
3. Triggers KnowledgeGraphAgent for graph updates
4. Schedules RecallAgent for spaced repetition
5. Integrates with Notion for permanent storage
6. Syncs with Supabase for cross-device access
"""

import asyncio
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum

from .doubt_predictor import DoubtPredictorAgent, DoubtPrediction
from .behavioral_agent import BehavioralAgent, BehavioralSignal
from .knowledge_graph_agent import KnowledgeGraphAgent
from .recall_agent import RecallAgent, RecallCard
from .peer_learning_agent import PeerLearningAgent


class SessionPhase(Enum):
    PRE_LEARNING = "pre_learning"
    ACTIVE_LEARNING = "active_learning"
    REVIEW = "review"
    BREAK = "break"
    POST_LEARNING = "post_learning"


@dataclass
class LearningSession:
    session_id: str
    user_id: str
    topic: str
    phase: SessionPhase
    started_at: datetime
    ended_at: Optional[datetime] = None
    predictions: List[DoubtPrediction] = field(default_factory=list)
    captured_doubts: List[Dict] = field(default_factory=list)
    behavioral_signals: List[BehavioralSignal] = field(default_factory=list)
    recommendations: List[str] = field(default_factory=list)
    xp_earned: int = 0
    notes: str = ""


@dataclass
class OrchestratorState:
    current_session: Optional[LearningSession] = None
    active_predictions: List[DoubtPrediction] = field(default_factory=list)
    pending_recalls: List[RecallCard] = field(default_factory=list)
    peer_insights: List[Dict] = field(default_factory=list)
    gamification_state: Dict = field(default_factory=dict)


class StudyOrchestrator:
    """
    Central orchestration agent that coordinates all learning agents.
    
    Workflow:
    1. PRE_LEARNING: Load predictions, check recall queue, get peer insights
    2. ACTIVE_LEARNING: Monitor behavioral signals, update predictions, capture doubts
    3. REVIEW: Trigger spaced repetition, update knowledge graph
    4. POST_LEARNING: Sync to Notion, update gamification, generate session summary
    """
    
    def __init__(self, user_id: str, config: Optional[Dict] = None):
        self.user_id = user_id
        self.config = config or {}
        
        self.doubt_predictor = DoubtPredictorAgent(user_id, config)
        self.behavioral_agent = BehavioralAgent(user_id, config)
        self.knowledge_graph = KnowledgeGraphAgent(user_id, config)
        self.recall_agent = RecallAgent(user_id, config)
        self.peer_agent = PeerLearningAgent(user_id, config)
        
        self.state = OrchestratorState()
        
        self.session_history = []
        
    async def start_session(self, topic: str, subtopic: str = "") -> LearningSession:
        """Start a new learning session"""
        session = LearningSession(
            session_id=f"session_{datetime.now().timestamp()}",
            user_id=self.user_id,
            topic=topic,
            phase=SessionPhase.PRE_LEARNING,
            started_at=datetime.now()
        )
        
        self.state.current_session = session
        
        learning_context = await self._build_learning_context(topic, subtopic)
        
        predictions = self.doubt_predictor.predict_doubts(learning_context, top_k=5)
        session.predictions = predictions
        self.state.active_predictions = predictions
        
        recalls = await self.recall_agent.get_due_recalls(topic)
        self.state.pending_recalls = recalls
        
        peer_insights = await self.peer_agent.get_peer_insights(topic)
        self.state.peer_insights = peer_insights
        
        return session
    
    async def update_session(
        self, 
        behavioral_data: Dict,
        captured_doubt: Optional[Dict] = None
    ):
        """Update session with new behavioral data and captured doubts"""
        if not self.state.current_session:
            return
        
        gesture_signal = behavioral_data.get('gesture_signal')
        if gesture_signal:
            self.behavioral_agent.add_gesture_signal(gesture_signal)
            
            learning_context = await self._build_learning_context(
                self.state.current_session.topic,
                ''
            )
            learning_context['gesture_signal'] = gesture_signal
            
            if gesture_signal.get('signal_type') in ['confusion', 'cognitive_load', 'doubt_intent']:
                new_predictions = self.doubt_predictor.predict_doubts(
                    learning_context, 
                    top_k=3,
                    gesture_influence=gesture_signal.get('confidence', 0.5)
                )
                for pred in new_predictions:
                    if pred.confidence > 0.5:
                        self.state.active_predictions.append(pred)
        
        signals = self.behavioral_agent.process_signals(behavioral_data)
        self.state.current_session.behavioral_signals.extend(signals)
        
        if captured_doubt:
            self.state.current_session.captured_doubts.append(captured_doubt)
            
            self.doubt_predictor.update_policy(
                state=self.doubt_predictor.get_current_state(behavioral_data),
                predicted_doubt=captured_doubt.get('predicted_from', ''),
                actual_doubt=captured_doubt.get('doubt_text', ''),
                reward=captured_doubt.get('reward', 1.0)
            )
            
            await self.knowledge_graph.add_doubt_to_graph(captured_doubt)
        
        confusion = self.behavioral_agent.calculate_confusion_score(signals)
        
        if confusion > 0.7 and len(self.state.current_session.captured_doubts) < 3:
            learning_context = await self._build_learning_context(
                self.state.current_session.topic,
                ''
            )
            learning_context['confusion_score'] = confusion
            
            new_predictions = self.doubt_predictor.predict_doubts(learning_context, top_k=3)
            
            for pred in new_predictions:
                if pred.confidence > 0.5:
                    self.state.active_predictions.append(pred)
    
    async def trigger_review(self) -> List[RecallCard]:
        """Trigger spaced repetition review"""
        recalls = await self.recall_agent.get_due_recalls(
            self.state.current_session.topic if self.state.current_session else None
        )
        return recalls
    
    async def complete_review(
        self, 
        recall_id: str, 
        quality: int
    ) -> Dict:
        """Complete a recall card review"""
        result = await self.recall_agent.complete_review(recall_id, quality)
        
        if self.state.current_session:
            xp = self._calculate_xp_for_review(quality)
            self.state.current_session.xp_earned += xp
            
            self.state.gamification_state = await self._update_gamification(xp)
        
        return result
    
    async def end_session(self) -> Dict:
        """End the current session and generate summary"""
        if not self.state.current_session:
            return {}
        
        session = self.state.current_session
        session.ended_at = datetime.now()
        
        session_summary = {
            'session_id': session.session_id,
            'duration': (session.ended_at - session.started_at).total_seconds(),
            'topic': session.topic,
            'doubts_captured': len(session.captured_doubts),
            'predictions_made': len(session.predictions),
            'xp_earned': session.xp_earned,
            'predictions_accuracy': self._calculate_prediction_accuracy(session),
            'confusion_peaks': self._find_confusion_peaks(session.behavioral_signals),
            'topics_covered': list(set([
                d.get('topic', '') for d in session.captured_doubts
            ]))
        }
        
        self.session_history.append(session_summary)
        
        await self.knowledge_graph.sync_to_graph()
        
        await self._sync_to_notion(session)
        
        await self._sync_to_supabase(session_summary)
        
        self.state.current_session = None
        
        return session_summary
    
    async def _build_learning_context(
        self, 
        topic: str, 
        subtopic: str
    ) -> Dict:
        """Build comprehensive learning context"""
        context = {
            'topic': topic,
            'subtopic': subtopic,
            'progress': 0.0,
            'time_spent': 0,
            'confusion_score': 0.0,
            'eye_confidence': 0.0,
            'scroll_reversals': 0,
            'selections': 0,
            'prev_doubts': 0,
            'mastery': 0.0,
            'difficulty': 0.5,
            'streak': 0
        }
        
        if self.state.current_session:
            context['time_spent'] = (
                datetime.now() - self.state.current_session.started_at
            ).total_seconds()
            context['prev_doubts'] = len(self.state.current_session.captured_doubts)
            
            if self.state.current_session.behavioral_signals:
                signals = self.state.current_session.behavioral_signals[-10:]
                context['confusion_score'] = self.behavioral_agent.calculate_confusion_score(signals)
        
        return context
    
    def _calculate_xp_for_review(self, quality: int) -> int:
        """Calculate XP earned for review"""
        base_xp = {1: 5, 2: 8, 3: 10, 4: 15, 5: 25}
        return base_xp.get(quality, 5)
    
    async def _update_gamification(self, xp: int) -> Dict:
        """Update gamification state"""
        if 'total_xp' not in self.state.gamification_state:
            self.state.gamification_state = {
                'total_xp': 0,
                'level': 1,
                'streak': 0,
                'fish_xp': 0,
                'achievements': []
            }
        
        self.state.gamification_state['total_xp'] += xp
        self.state.gamification_state['fish_xp'] += xp // 2
        
        self.state.gamification_state['level'] = self._calculate_level(
            self.state.gamification_state['total_xp']
        )
        
        return self.state.gamification_state
    
    def _calculate_level(self, xp: int) -> int:
        """Calculate level from XP"""
        level_thresholds = [0, 100, 300, 600, 1000, 1500, 2200, 3000, 4000, 5500]
        for i, threshold in enumerate(level_thresholds):
            if xp < threshold:
                return max(1, i)
        return len(level_thresholds)
    
    def _calculate_prediction_accuracy(self, session: LearningSession) -> float:
        """Calculate accuracy of doubt predictions"""
        if not session.predictions:
            return 0.0
        
        correct = 0
        for captured in session.captured_doubts:
            predicted = captured.get('predicted_from', '')
            actual = captured.get('doubt_text', '')
            
            for pred in session.predictions:
                if pred.predicted_doubt.lower() in actual.lower():
                    correct += 1
                    break
        
        return correct / max(len(session.captured_doubts), 1)
    
    def _find_confusion_peaks(self, signals: List[BehavioralSignal]) -> List[Dict]:
        """Find moments of peak confusion"""
        peaks = []
        
        confusion_values = [
            self.behavioral_agent.calculate_confusion_score([s]) 
            for s in signals
        ]
        
        threshold = 0.7
        in_peak = False
        peak_start = 0
        
        for i, val in enumerate(confusion_values):
            if val > threshold and not in_peak:
                in_peak = True
                peak_start = i
            elif val < threshold and in_peak:
                in_peak = False
                peaks.append({
                    'start_index': peak_start,
                    'end_index': i,
                    'max_value': max(confusion_values[peak_start:i])
                })
        
        return peaks
    
    async def _sync_to_notion(self, session: LearningSession):
        """Sync session data to Notion"""
        pass
    
    async def _sync_to_supabase(self, session_summary: Dict):
        """Sync session data to Supabase"""
        pass
    
    def get_active_insights(self) -> Dict:
        """Get current active insights for display"""
        return {
            'predictions': [
                {
                    'doubt': p.predicted_doubt,
                    'confidence': p.confidence,
                    'explanation': p.suggested_explanation
                }
                for p in self.state.active_predictions[:3]
            ],
            'pending_reviews': len(self.state.pending_recalls),
            'peer_insights_count': len(self.state.peer_insights),
            'gamification': self.state.gamification_state,
            'session_active': self.state.current_session is not None
        }