app/services/chat_service.py

from sqlalchemy.orm import Session
from typing import List, Dict, Optional
import json

from app.database.models import Message, Session as ChatSession
from app.services.session_service import SessionService
from app.utils.helpers import generate_id


class ChatService:
    """Service for chat operations and agent orchestration."""
    
    @staticmethod
    def process_message(
        db: Session,
        message: str,
        session_id: Optional[str] = None,
        user_id: Optional[str] = None,
        policy_ids: Optional[List[str]] = None
    ) -> Dict:
        """
        Process a user message and generate response using LangGraph workflow.
        
        Args:
            db: Database session
            message: User message content
            session_id: Optional session ID
            user_id: Optional user ID
            policy_ids: Optional list of policy IDs to search within
            
        Returns:
            Dictionary with response message and metadata
        """
        from app.llm.graph import multi_agent_graph
        
        # Create or get session
        if not session_id:
            chat_session = SessionService.create_session(db, user_id)
            session_id = chat_session.id
        else:
            chat_session = SessionService.get_session_by_id(db, session_id)
            if not chat_session:
                raise ValueError("Session not found")
        
        # Save user message
        user_message = Message(
            id=generate_id(),
            session_id=session_id,
            role="user",
            content=message
        )
        db.add(user_message)
        db.commit()
        
        # Get chat history for context
        chat_history = ChatService.get_chat_history(db, session_id, limit=10)
        
        # Format chat history for the graph
        history_messages = [
            {"role": msg["role"], "content": msg["content"]}
            for msg in chat_history
        ]
        
        # Process query through LangGraph workflow
        response_data = multi_agent_graph.process_query(
            query=message,
            user_id=user_id,
            chat_history=history_messages,
            policy_ids=policy_ids
        )
        
        # Prepare metadata
        meta = {
            "agent": response_data.get("agent", "unknown"),
            "routing_reasoning": response_data.get("routing_reasoning", ""),
            "sources": response_data.get("sources", []),
            "metadata": response_data.get("metadata", {}),
            "policy_names": response_data.get("policy_names", [])  # Policy document names
        }
        
        print(f"[Chat Service] Saving meta with policy_names: {meta.get('policy_names', [])}")
        
        # Save assistant message
        assistant_message = Message(
            id=generate_id(),
            session_id=session_id,
            role="assistant",
            content=response_data.get("answer", "I apologize, but I couldn't generate a response."),
            meta=json.dumps(meta)
        )
        db.add(assistant_message)
        
        # Update session timestamp
        SessionService.update_session_timestamp(db, session_id)
        
        db.commit()
        db.refresh(assistant_message)
        
        # Auto-generate session title if this is the first exchange
        messages_count = db.query(Message).filter(Message.session_id == session_id).count()
        if messages_count == 2 and chat_session.title == "New Conversation":
            # Generate title from first user message
            title = ChatService._generate_session_title(message)
            SessionService.update_session_title(db, session_id, title)
        
        return {
            "message": assistant_message,
            "session_id": session_id,
            "agent": meta["agent"],
            "sources": meta.get("sources", [])
        }
    
    @staticmethod
    def get_chat_history(
        db: Session,
        session_id: str,
        limit: Optional[int] = None
    ) -> List[Dict]:
        """
        Get chat history for a session.
        
        Args:
            db: Database session
            session_id: Session ID
            limit: Optional limit on number of messages
            
        Returns:
            List of message dictionaries
        """
        query = db.query(Message).filter(
            Message.session_id == session_id
        ).order_by(Message.created_at.asc())
        
        if limit:
            # Get last N messages
            total = query.count()
            if total > limit:
                query = query.offset(total - limit)
        
        messages = query.all()
        
        return [
            {
                "id": msg.id,
                "role": msg.role,
                "content": msg.content,
                "meta": json.loads(msg.meta) if msg.meta else {},
                "created_at": msg.created_at.isoformat()
            }
            for msg in messages
        ]
    
    @staticmethod
    def _generate_session_title(first_message: str) -> str:
        """
        Generate a ChatGPT-style session title using LLM.
        
        Args:
            first_message: The first user message in the conversation
            
        Returns:
            A concise, descriptive title (max 50 characters)
        """
        from app.llm.client import llm_client
        
        try:
            prompt = f"""Generate a very short, concise title for a chat conversation that starts with this message:

"{first_message}"

Requirements:
- Maximum 50 characters
- Be specific and descriptive
- Capture the main topic/question
- Professional tone
- No quotes around the title
- Examples: "Building Code Requirements", "Fire Safety Regulations", "Basement Definition"

Return ONLY the title, nothing else:"""

            title = llm_client.get_completion(
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=20
            )
            
            # Clean up the title
            title = title.strip().strip('"').strip("'")
            
            # Ensure it's not too long
            if len(title) > 50:
                title = title[:47] + "..."
            
            # Fallback if empty or too short
            if len(title) < 3:
                title = first_message[:50].strip()
                if len(first_message) > 50:
                    title += "..."
            
            return title
            
        except Exception as e:
            print(f"[Chat Service] Error generating title: {e}")
            # Fallback to simple truncation
            title = first_message[:50].strip()
            if len(first_message) > 50:
                title += "..."
            return title


# Global chat service instance
chat_service = ChatService()