"""
Chat Assistant Module
Implements a conversational AI assistant using Ollama for policy and compliance questions.
"""

import ollama
from typing import List, Dict, Optional, Any
from pydantic import BaseModel
from fastapi import HTTPException
from src.utils import initialize_a_logger


class ChatMessage(BaseModel):
    """Represents a single chat message."""
    role: str  # "user" or "assistant"
    content: str


class ChatRequest(BaseModel):
    """Request model for chat messages."""
    message: str
    conversation_history: Optional[List[ChatMessage]] = None
    context: Optional[str] = None  # Additional context about policies/documents


class ChatResponse(BaseModel):
    """Response model for chat messages."""
    message: str
    model: str
    tokens_used: Optional[int] = None


class ChatAssistant:
    """
    Chat Assistant for answering policy and compliance questions.
    Uses Ollama to provide intelligent responses about IT governance, compliance, and policies.
    """

    def __init__(self, model_name: str = "qwen2.5:3b", config: Optional[Dict] = None):
        """
        Initialize the Chat Assistant.
        
        Args:
            model_name: Name of the Ollama model to use
            config: Optional configuration dictionary
        """
        self.model_name = model_name
        self.config = config or {}
        self.logger = initialize_a_logger('./logs/ChatAssistant.log')
        self.logger.info(f"ChatAssistant initialized with model: {model_name}")
        
        # Check Ollama connectivity on initialization
        self._check_ollama_connection()

        # System prompt for the assistant
        self.system_prompt = """You are a helpful compliance and policy assistant. Your role is to:
1. Answer questions about IT governance, compliance policies, and regulatory requirements
2. Provide guidance on ISO 27001, NIST, GDPR, and other compliance frameworks
3. Help users understand workflow processes and requirements
4. Verify content against compliance standards
5. Be accurate, helpful, and concise in your responses

Always provide clear, actionable advice. If you're unsure about something, say so rather than guessing.
When discussing compliance requirements, cite specific standards or frameworks when possible."""

    def _check_ollama_connection(self):
        """Check if Ollama is accessible and model is available."""
        try:
            models = ollama.list()
            model_names = [m.get("name", "") for m in models.get("models", [])]
            self.logger.info(f"Ollama connected. Available models: {model_names}")
            
            # Check if our model is available
            if self.model_name not in model_names:
                self.logger.warning(
                    f"Model '{self.model_name}' not found in available models. "
                    f"Available models: {model_names}. "
                    f"Trying to use it anyway - it may need to be pulled."
                )
        except Exception as e:
            self.logger.error(f"Failed to connect to Ollama: {e}")
            self.logger.warning(
                "Ollama connection check failed. "
                "The service may not be available. "
                "Chat functionality may not work until Ollama is running."
            )

    def _build_messages(
        self, 
        user_message: str, 
        conversation_history: Optional[List[ChatMessage]] = None,
        context: Optional[str] = None
    ) -> List[Dict[str, str]]:
        """
        Build the message list for Ollama API.
        
        Args:
            user_message: The current user message
            conversation_history: Previous messages in the conversation
            context: Additional context to include
            
        Returns:
            List of message dictionaries for Ollama
        """
        messages = []
        
        # Add system prompt
        messages.append({
            "role": "system",
            "content": self.system_prompt
        })
        
        # Add context if provided
        if context:
            messages.append({
                "role": "system",
                "content": f"Additional context: {context}"
            })
        
        # Add conversation history
        if conversation_history:
            for msg in conversation_history[-10:]:  # Keep last 10 messages for context
                messages.append({
                    "role": msg.role,
                    "content": msg.content
                })
        
        # Add current user message
        messages.append({
            "role": "user",
            "content": user_message
        })
        
        return messages

    async def chat(self, request: ChatRequest) -> ChatResponse:
        """
        Process a chat message and return a response.
        
        Args:
            request: Chat request with message and optional history
            
        Returns:
            Chat response with assistant's message
            
        Raises:
            HTTPException: If the chat request fails
        """
        try:
            self.logger.info(f"Processing chat message: {request.message[:100] if request.message else 'Empty message'}...")
            
            # Validate request
            if not request.message or not request.message.strip():
                raise ValueError("Message cannot be empty")
            
            # Build messages for Ollama
            messages = self._build_messages(
                user_message=request.message,
                conversation_history=request.conversation_history,
                context=request.context
            )
            
            self.logger.debug(f"Sending {len(messages)} messages to Ollama model: {self.model_name}")
            
            # Call Ollama API
            try:
                response = ollama.chat(
                    model=self.model_name,
                    messages=messages,
                    options={
                        "temperature": 0.7,  # Slightly creative for conversational responses
                        "top_p": 0.9,
                    }
                )
            except ConnectionError as e:
                self.logger.error(f"Ollama connection error: {e}", exc_info=True)
                raise HTTPException(
                    status_code=503,
                    detail="Cannot connect to Ollama service. Please ensure Ollama is running and accessible."
                )
            except Exception as ollama_error:
                error_str = str(ollama_error).lower()
                self.logger.error(f"Ollama API error: {ollama_error}", exc_info=True)
                
                # Check if it's a connection error
                if "connection" in error_str or "refused" in error_str or "connect" in error_str:
                    raise HTTPException(
                        status_code=503,
                        detail="Ollama service is not available. Please ensure Ollama is running on localhost:11434."
                    )
                # Check if model is not found
                if "not found" in error_str or ("model" in error_str and "not" in error_str):
                    raise HTTPException(
                        status_code=404,
                        detail=f"Model '{self.model_name}' not found. Please ensure the model is available. Try: ollama pull {self.model_name}"
                    )
                # Generic Ollama error
                raise HTTPException(
                    status_code=500,
                    detail=f"Ollama error: {str(ollama_error)}"
                )
            
            # Extract response content
            assistant_message = response.get("message", {}).get("content", "")
            
            if not assistant_message:
                self.logger.warning("Empty response from Ollama")
                assistant_message = "I apologize, but I couldn't generate a response. Please try again."
            
            self.logger.info(f"Generated response: {assistant_message[:100]}...")
            
            return ChatResponse(
                message=assistant_message,
                model=self.model_name,
                tokens_used=response.get("eval_count", 0)
            )
            
        except HTTPException:
            # Re-raise HTTP exceptions as-is
            raise
        except Exception as e:
            error_message = str(e)
            self.logger.error(f"Error in chat: {error_message}", exc_info=True)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to generate chat response: {error_message}"
            )

    async def verify_content(
        self, 
        field_name: str, 
        content: str,
        verification_criteria: Optional[str] = None
    ) -> ChatResponse:
        """
        Verify content against compliance requirements.
        
        Args:
            field_name: Name of the field being verified
            content: Content to verify
            verification_criteria: Optional specific criteria to check against
            
        Returns:
            Chat response with verification feedback
        """
        try:
            self.logger.info(f"Verifying content for field: {field_name}")
            
            # Build verification prompt
            verification_prompt = f"""Please review and verify the following content from the field "{field_name}".

Content to verify:
"{content}"

Please provide:
1. Whether the content meets compliance requirements
2. Any suggestions for improvement
3. Any potential issues or concerns
4. Specific recommendations

"""
            
            if verification_criteria:
                verification_prompt += f"\nSpecific criteria to check:\n{verification_criteria}\n"
            
            verification_prompt += "\nProvide a detailed, helpful verification response."
            
            # Create chat request
            chat_request = ChatRequest(
                message=verification_prompt,
                context=f"Content verification for field: {field_name}"
            )
            
            return await self.chat(chat_request)
            
        except HTTPException:
            raise
        except Exception as e:
            error_message = str(e)
            self.logger.error(f"Error verifying content: {error_message}", exc_info=True)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to verify content: {error_message}"
            )

    async def answer_policy_question(
        self, 
        question: str,
        policy_context: Optional[str] = None
    ) -> ChatResponse:
        """
        Answer a question about policies or compliance.
        
        Args:
            question: The user's question
            policy_context: Optional context about specific policies
            
        Returns:
            Chat response with the answer
        """
        try:
            self.logger.info(f"Answering policy question: {question[:100]}...")
            
            # Enhance question with policy context
            enhanced_question = question
            if policy_context:
                enhanced_question = f"Context: {policy_context}\n\nQuestion: {question}"
            
            chat_request = ChatRequest(
                message=enhanced_question,
                context="Policy and compliance question answering"
            )
            
            return await self.chat(chat_request)
            
        except HTTPException:
            raise
        except Exception as e:
            error_message = str(e)
            self.logger.error(f"Error answering policy question: {error_message}", exc_info=True)
            raise HTTPException(
                status_code=500,
                detail=f"Failed to answer question: {error_message}"
            )


# Global instance
_chat_assistant_instance: Optional[ChatAssistant] = None


def get_chat_assistant(model_name: str = "qwen2.5:3b") -> ChatAssistant:
    """
    Get or create the global ChatAssistant instance.
    
    Args:
        model_name: Name of the Ollama model to use
        
    Returns:
        ChatAssistant instance
    """
    global _chat_assistant_instance
    if _chat_assistant_instance is None:
        _chat_assistant_instance = ChatAssistant(model_name=model_name)
    return _chat_assistant_instance
