from src.domain.entities.conversation import Conversation, ConversationMessage, MessageContent, MessageType
from src.domain.ports.repositories.conversation_repository import ConversationRepositoryPort
from src.domain.ports.repositories.charts_repository import ChartsRepositoryPort
from src.application.ports.llm_service_port import LLMServicePort
from src.application.dtos.LLM import LLMRequestDTO
from src.application.dtos.conversation import SendMessageRequestDTO, MessageResponseDTO, MessageContentDTO
import uuid
from datetime import datetime, timezone
from typing import List, Optional

class ChatConversationUseCase:
    def __init__(
        self, 
        conversation_repo: ConversationRepositoryPort,
        charts_repo: ChartsRepositoryPort,
        llm_service: LLMServicePort
    ):
        self._conversation_repo = conversation_repo
        self._charts_repo = charts_repo
        self._llm_service = llm_service

    async def execute(
        self, 
        conversation_id: str, 
        user_id: str, 
        request: SendMessageRequestDTO
    ) -> MessageResponseDTO:
        """Handle interactive chat conversation with support for multiple charts"""
        
        # Get conversation
        conversation = await self._conversation_repo.get_conversation_by_id(conversation_id)
        
        if not conversation:
            raise ValueError("Conversation not found")
        
        if conversation.user_id != user_id:
            raise ValueError("Conversation does not belong to user")

        # If chart is included in message, verify it exists and belongs to user
        if request.chart_image_id:
            chart_image = await self._charts_repo.get_by_id(request.chart_image_id)
            if not chart_image:
                raise ValueError("Chart image not found")
            
            if chart_image.user_id != user_id:
                raise ValueError("Chart image does not belong to user")
            
            # Add to referenced charts if not already there
            if request.chart_image_id not in conversation.referenced_charts:
                conversation.referenced_charts.append(request.chart_image_id)

        # Create user message content
        user_content = MessageContent(
            text=request.message,
            chart_image_id=request.chart_image_id
        )

        # Create user message
        user_message = ConversationMessage(
            id=str(uuid.uuid4()),
            conversation_id=conversation_id,
            sender_id=user_id,
            message_type=MessageType.TEXT,
            content=user_content,
            timestamp=datetime.now(timezone.utc)
        )

        # Add user message to conversation
        conversation.messages.append(user_message)
        conversation.updated_at = datetime.now(timezone.utc)

        # Get AI response if there's text in the message
        assistant_message = None
        if request.message.strip():
            # Build context from conversation history
            context = self._build_conversation_context(conversation.messages[:-1])  # Exclude current message
            
            # Create enhanced question with context
            enhanced_question = self._create_enhanced_question(request.message, context)

            # Get the most recent chart for analysis (or the one in current message)
            chart_to_analyze = request.chart_image_id
            if not chart_to_analyze and conversation.referenced_charts:
                chart_to_analyze = conversation.referenced_charts[-1]  # Most recent chart

            if chart_to_analyze:
                chart_image = await self._charts_repo.get_by_id(chart_to_analyze)
                if chart_image:
                    # Get image bytes (handles both file path and direct data)
                    image_bytes = chart_image.get_image_bytes()
                    
                    # Get AI response using LLM service
                    llm_request = LLMRequestDTO(
                        image_bytes=image_bytes,
                        question=enhanced_question
                    )

                    llm_response = await self._llm_service.analyze(llm_request)

                    # Create assistant message
                    assistant_content = MessageContent.create_text(llm_response.answer)
                    assistant_message = ConversationMessage(
                        id=str(uuid.uuid4()),
                        conversation_id=conversation_id,
                        sender_id="system",  # AI assistant
                        message_type=MessageType.TEXT,
                        content=assistant_content,
                        timestamp=datetime.now(timezone.utc)
                    )

                    # Add assistant message to conversation
                    conversation.messages.append(assistant_message)
                    conversation.updated_at = datetime.now(timezone.utc)

        # Save updated conversation
        await self._conversation_repo.update_conversation(conversation)

        # Return the assistant message if it exists, otherwise return user message
        message_to_return = assistant_message if assistant_message else user_message
        
        content_dto = MessageContentDTO(
            text=message_to_return.content.text,
            chart_image_id=message_to_return.content.chart_image_id
        )
        
        return MessageResponseDTO(
            id=message_to_return.id,
            message_type=message_to_return.message_type,
            content=content_dto,
            timestamp=message_to_return.timestamp.isoformat(),
            sender_id=message_to_return.sender_id
        )
    
    def _build_conversation_context(self, messages: List[ConversationMessage]) -> str:
        """Build context string from conversation history"""
        if not messages:
            return ""
        
        # Take last 5 messages for context to avoid token limits
        recent_messages = messages[-5:]
        
        context_parts = []
        for i, msg in enumerate(recent_messages, 1):
            role = "User" if msg.sender_id != "system" else "Assistant"
            content_text = msg.content.text or ""
            if msg.content.chart_image_id:
                content_text += f" [Chart: {msg.content.chart_image_id}]"
            context_parts.append(f"Turn {i} - {role}: {content_text}")
        
        return " | ".join(context_parts)

    def _create_enhanced_question(self, current_question: str, context: str) -> str:
        """Create enhanced question with conversation context"""
        if not context:
            return current_question
        
        return f"""Previous conversation context: {context}

Current question: {current_question}

Please provide a response that takes into account the conversation history and directly addresses the current question about the chart.""" 