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 services.ollama_service import OllamaService
from utils.json_handler import JSONHandler

class RoleplayEngine:
    def __init__(self, groq_service=None, ollama_service=None):
        # Use provided services or create new ones (for backward compatibility)
        self.groq_service = groq_service or GroqService()
        self.ollama_service = ollama_service or OllamaService()
        self.json_handler = JSONHandler()

    def _get_ai_service(self, use_groq: bool):
        """Get the appropriate AI service based on groqRoleplay flag"""
        return self.groq_service if use_groq else self.ollama_service

    def start_session(self, scenario: RoleplayScenario, is_admin: int, force_restart: bool = False) -> str:
        """Start a new roleplay session"""
        session_id = scenario.session_id  # Use the session_id from scenario

        # Always create new conversation if force_restart=True, otherwise only if no conversation exists
        existing_conversation = self.json_handler.load_conversation(session_id, is_admin)

        if force_restart:
            print(f"DEBUG: Force restarting conversation - replacing existing conversation")
        elif existing_conversation:
            print(f"DEBUG: Conversation already exists, not creating new one")
            return session_id
        else:
            print(f"DEBUG: No existing conversation found, creating new one")

        # Initialize NEW conversation with scenario introduction and first AI message
        conversation_data = {
            'session_id': session_id,
            'scenario_id': scenario.session_id,
            'started_at': datetime.now().isoformat(),
            'status': 'active',
            'turns': [
                {
                    'turn_number': 1,
                    'speaker': 'system',
                    'message': scenario.scenario_intro,
                    'timestamp': datetime.now().isoformat()
                },
                {
                    'turn_number': 2,
                    'speaker': 'ai',
                    'message': scenario.conversation_starter,
                    'timestamp': datetime.now().isoformat()
                }
            ]
        }

        # Save initial conversation state
        self.json_handler.save_conversation(session_id, conversation_data, is_admin)

        return session_id

    def add_learner_response(self, session_id: str, scenario: RoleplayScenario, learner_message: str, is_admin: int, use_groq: bool = True) -> Optional[str]:
        """Add learner response and get AI reply"""

        # Load current conversation
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        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

        # Always generate dynamic AI response based on user input
        # 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'
        ]

        # Determine if this should be a conclusion response
        # After exactly 4 learner turns, provide a professional conclusion (only once)
        is_conclusion = len(learner_turns) == 4

        # Get AI response using selected service
        ai_service = self._get_ai_service(use_groq)
        ai_response = ai_service.play_character(
            scenario.to_dict(),
            conversation_history,
            learner_message,
            is_conclusion
        )

        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, is_admin)

        return ai_response

    def end_session(self, session_id: str, is_admin: int) -> bool:
        """End a roleplay session"""
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        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, is_admin)

    def get_conversation_history(self, session_id: str, is_admin: int) -> Optional[List[Dict[str, Any]]]:
        """Get full conversation history"""
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        if not conversation_data:
            return None

        return conversation_data['turns']

    def get_conversation_turns_for_assessment(self, session_id: str, is_admin: int) -> Optional[List[ConversationTurn]]:
        """Get conversation turns formatted for assessment"""
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        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, is_admin: int) -> bool:
        """Check if a session is still active"""
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        if not conversation_data:
            return False

        return conversation_data.get('status', 'active') == 'active'

    def get_session_info(self, session_id: str, is_admin: int) -> Optional[Dict[str, Any]]:
        """Get basic session information"""
        conversation_data = self.json_handler.load_conversation(session_id, is_admin)
        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')
        }