from typing import Dict, List, Any, Optional
from datetime import datetime
import uuid
from models.scenario import RoleplayScenario
from models.assessment import ConversationTurn
from services.groq_service import GroqService
from utils.json_handler import JSONHandler

class RoleplayEngine:
    def __init__(self):
        self.groq_service = GroqService()
        self.json_handler = JSONHandler()
    
    def start_session(self, scenario: RoleplayScenario) -> str:
        """Start a new roleplay session"""
        session_id = str(uuid.uuid4())
        
        # Initialize conversation with scenario introduction 
        conversation_data = {
            'session_id': session_id,
            'scenario_id': scenario.id,
            'started_at': datetime.now().isoformat(),
            'status': 'active',
            'turns': [
                {
                    'turn_number': 1,
                    'speaker': 'system',
                    'message': scenario.scenario_intro,
                    'timestamp': datetime.now().isoformat()
                }
            ]
        }
        
        # Save initial conversation state
        self.json_handler.save_conversation(session_id, conversation_data)
        
        return session_id
    
    def add_learner_response(self, session_id: str, scenario: RoleplayScenario, learner_message: str) -> Optional[str]:
        """Add learner response and get AI reply"""
        
        # Load current conversation
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return None
        
        # Add learner's message
        turn_number = len(conversation_data['turns']) + 1
        learner_turn = {
            'turn_number': turn_number,
            'speaker': 'learner',
            'message': learner_message,
            'timestamp': datetime.now().isoformat()
        }
        conversation_data['turns'].append(learner_turn)
        
        # Check if this is the first learner response (after system intro)
        learner_turns = [turn for turn in conversation_data['turns'] if turn['speaker'] == 'learner']
        is_first_learner_response = len(learner_turns) == 1
        
        if is_first_learner_response:
            # Use the conversation starter for AI's first response
            ai_response = scenario.conversation_starter
        else:
            # Get conversation history for AI context (exclude system messages)
            conversation_history = [
                {
                    'speaker': turn['speaker'],
                    'message': turn['message']
                }
                for turn in conversation_data['turns'][:-1] if turn['speaker'] != 'system'
            ]
            
            # Get AI response
            ai_response = self.groq_service.play_character(
                scenario.to_dict(),
                conversation_history,
                learner_message
            )
        
        if not ai_response:
            return None
        
        # Add AI's response
        ai_turn = {
            'turn_number': turn_number + 1,
            'speaker': 'ai',
            'message': ai_response,
            'timestamp': datetime.now().isoformat()
        }
        conversation_data['turns'].append(ai_turn)
        
        # Save updated conversation
        self.json_handler.save_conversation(session_id, conversation_data)
        
        return ai_response
    
    def end_session(self, session_id: str) -> bool:
        """End a roleplay session"""
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return False
        
        conversation_data['status'] = 'completed'
        conversation_data['ended_at'] = datetime.now().isoformat()
        
        return self.json_handler.save_conversation(session_id, conversation_data)
    
    def get_conversation_history(self, session_id: str) -> Optional[List[Dict[str, Any]]]:
        """Get full conversation history"""
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return None
        
        return conversation_data['turns']
    
    def get_conversation_turns_for_assessment(self, session_id: str) -> Optional[List[ConversationTurn]]:
        """Get conversation turns formatted for assessment"""
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return None
        
        turns = []
        for turn_data in conversation_data['turns']:
            turn = ConversationTurn(
                turn_number=turn_data['turn_number'],
                speaker=turn_data['speaker'],
                message=turn_data['message'],
                timestamp=turn_data['timestamp']
            )
            turns.append(turn)
        
        return turns
    
    def is_session_active(self, session_id: str) -> bool:
        """Check if a session is still active"""
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return False
        
        return conversation_data.get('status', 'active') == 'active'
    
    def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Get basic session information"""
        conversation_data = self.json_handler.load_conversation(session_id)
        if not conversation_data:
            return None
        
        return {
            'session_id': conversation_data['session_id'],
            'scenario_id': conversation_data['scenario_id'],
            'started_at': conversation_data['started_at'],
            'status': conversation_data.get('status', 'active'),
            'turn_count': len(conversation_data['turns']),
            'ended_at': conversation_data.get('ended_at')
        }