src/agent/viettelpay_agent.py

"""
ViettelPay AI Agent using LangGraph
Multi-turn conversation support with short-term memory using InMemorySaver
"""

import os
from typing import Dict, Optional
from functools import partial
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage

from src.agent.nodes import (
    ViettelPayState,
    classify_intent_node,
    query_enhancement_node,
    knowledge_retrieval_node,
    script_response_node,
    generate_response_node,
    route_after_intent_classification,
    route_after_query_enhancement,
    route_after_knowledge_retrieval,
)

# Import configuration utility
from src.utils.config import get_knowledge_base_path, get_llm_provider


class ViettelPayAgent:
    """Main ViettelPay AI Agent using LangGraph workflow with multi-turn conversation support"""

    def __init__(
        self,
        knowledge_base_path: str = None,
        scripts_file: Optional[str] = None,
        llm_provider: str = None,
    ):
        knowledge_base_path = knowledge_base_path or get_knowledge_base_path()
        scripts_file = scripts_file or "./viettelpay_docs/processed/kich_ban.csv"
        llm_provider = llm_provider or get_llm_provider()

        self.knowledge_base_path = knowledge_base_path
        self.scripts_file = scripts_file
        self.llm_provider = llm_provider

        # Initialize LLM client once during agent creation
        print(f"🧠 Initializing LLM client ({self.llm_provider})...")
        from src.llm.llm_client import LLMClientFactory

        self.llm_client = LLMClientFactory.create_client(self.llm_provider)
        print(f"✅ LLM client initialized and ready")

        # Initialize knowledge retriever once during agent creation
        print(f"📚 Initializing knowledge retriever...")
        try:
            from src.knowledge_base.viettel_knowledge_base import ViettelKnowledgeBase

            self.knowledge_base = ViettelKnowledgeBase()
            ensemble_retriever = self.knowledge_base.load_knowledge_base(
                knowledge_base_path
            )
            if not ensemble_retriever:
                raise ValueError(
                    f"Knowledge base not found at {knowledge_base_path}. Run build_database_script.py first."
                )
            print(f"✅ Knowledge retriever initialized and ready")
        except Exception as e:
            print(f"⚠️ Knowledge retriever initialization failed: {e}")
            self.knowledge_base = None

        # Initialize checkpointer for short-term memory
        self.checkpointer = InMemorySaver()

        # Build workflow with pre-initialized components
        self.workflow = self._build_workflow()
        self.app = self.workflow.compile(checkpointer=self.checkpointer)

        print("✅ ViettelPay Agent initialized with multi-turn conversation support")

    def _build_workflow(self) -> StateGraph:
        """Build LangGraph workflow with pre-initialized components"""

        # Create workflow graph
        workflow = StateGraph(ViettelPayState)

        # Create node functions with pre-bound components using functools.partial
        # This eliminates the need to initialize components in each node call
        classify_intent_with_llm = partial(
            classify_intent_node, llm_client=self.llm_client
        )
        query_enhancement_with_llm = partial(
            query_enhancement_node, llm_client=self.llm_client
        )
        knowledge_retrieval_with_retriever = partial(
            knowledge_retrieval_node, knowledge_retriever=self.knowledge_base
        )
        generate_response_with_llm = partial(
            generate_response_node, llm_client=self.llm_client
        )

        # Add nodes (some with pre-bound components, some without)
        workflow.add_node("classify_intent", classify_intent_with_llm)
        workflow.add_node("query_enhancement", query_enhancement_with_llm)
        workflow.add_node("knowledge_retrieval", knowledge_retrieval_with_retriever)
        workflow.add_node(
            "script_response", script_response_node
        )  # No pre-bound components needed
        workflow.add_node("generate_response", generate_response_with_llm)

        # Set entry point
        workflow.set_entry_point("classify_intent")

        # Add conditional routing after intent classification
        workflow.add_conditional_edges(
            "classify_intent",
            route_after_intent_classification,
            {
                "script_response": "script_response",
                "query_enhancement": "query_enhancement",
            },
        )

        # Script responses go directly to end
        workflow.add_edge("script_response", END)

        # Query enhancement goes to knowledge retrieval
        workflow.add_edge("query_enhancement", "knowledge_retrieval")

        # Knowledge retrieval goes to response generation
        workflow.add_edge("knowledge_retrieval", "generate_response")
        workflow.add_edge("generate_response", END)

        print("🔄 LangGraph workflow built successfully with optimized component usage")
        return workflow

    def process_message(self, user_message: str, thread_id: str = "default") -> Dict:
        """Process a user message in a multi-turn conversation"""

        print(f"\n💬 Processing message: '{user_message}' (thread: {thread_id})")
        print("=" * 50)

        # Create configuration with thread_id for conversation memory
        config = {"configurable": {"thread_id": thread_id}}

        try:
            # Create human message
            human_message = HumanMessage(content=user_message)

            # Initialize state with the new message
            # Note: conversation_context is set to None so it gets recomputed with fresh message history
            initial_state = {
                "messages": [human_message],
                "intent": None,
                "confidence": None,
                "enhanced_query": None,
                "retrieved_docs": None,
                "conversation_context": None,  # Reset to ensure fresh context computation
                "response_type": None,
                "error": None,
                "processing_info": None,
            }

            # Run workflow with memory
            result = self.app.invoke(initial_state, config)

            # Extract response from the last AI message
            messages = result.get("messages", [])
            if messages:
                # Get the last AI message
                last_message = messages[-1]
                if hasattr(last_message, "content"):
                    response = last_message.content
                else:
                    response = str(last_message)
            else:
                response = "Xin lỗi, em không thể xử lý yêu cầu này."

            response_type = result.get("response_type", "unknown")
            intent = result.get("intent", "unknown")
            confidence = result.get("confidence", 0.0)
            enhanced_query = result.get("enhanced_query", "")
            error = result.get("error")

            # Build response info
            response_info = {
                "response": response,
                "intent": intent,
                "confidence": confidence,
                "response_type": response_type,
                "enhanced_query": enhanced_query,
                "success": error is None,
                "error": error,
                "thread_id": thread_id,
                "message_count": len(messages),
            }

            print(f"✅ Response generated successfully")
            print(f"   Intent: {intent} (confidence: {confidence})")
            print(f"   Type: {response_type}")
            if enhanced_query and enhanced_query != user_message:
                print(f"   Enhanced query: {enhanced_query}")
            print(f"   Thread: {thread_id}")

            return response_info

        except Exception as e:
            print(f"❌ Workflow error: {e}")

            return {
                "response": "Xin lỗi, em gặp lỗi kỹ thuật. Vui lòng thử lại sau.",
                "intent": "error",
                "confidence": 0.0,
                "response_type": "error",
                "enhanced_query": "",
                "success": False,
                "error": str(e),
                "thread_id": thread_id,
                "message_count": 0,
            }

    def chat(self, user_message: str, thread_id: str = "default") -> str:
        """Simple chat interface - returns just the response text"""
        result = self.process_message(user_message, thread_id)
        return result["response"]

    def get_conversation_history(self, thread_id: str = "default") -> list:
        """Get conversation history for a specific thread"""
        try:
            config = {"configurable": {"thread_id": thread_id}}

            # Get the current state to access message history
            current_state = self.app.get_state(config)

            if current_state and current_state.values.get("messages"):
                messages = current_state.values["messages"]
                history = []

                for msg in messages:
                    if hasattr(msg, "type") and hasattr(msg, "content"):
                        role = "user" if msg.type == "human" else "assistant"
                        history.append({"role": role, "content": msg.content})
                    elif hasattr(msg, "role") and hasattr(msg, "content"):
                        history.append({"role": msg.role, "content": msg.content})

                return history
            else:
                return []

        except Exception as e:
            print(f"❌ Error getting conversation history: {e}")
            return []

    def clear_conversation(self, thread_id: str = "default") -> bool:
        """Clear conversation history for a specific thread"""
        try:
            # Note: InMemorySaver doesn't have a direct clear method
            # The conversation will be cleared when the app is restarted
            # For persistent memory, you'd need to implement a clear method
            print(f"📝 Conversation clearing requested for thread: {thread_id}")
            print("   Note: InMemorySaver conversations clear on app restart")
            return True
        except Exception as e:
            print(f"❌ Error clearing conversation: {e}")
            return False

    def get_workflow_info(self) -> Dict:
        """Get information about the workflow structure"""
        return {
            "nodes": [
                "classify_intent",
                "query_enhancement",
                "knowledge_retrieval",
                "script_response",
                "generate_response",
            ],
            "entry_point": "classify_intent",
            "knowledge_base_path": self.knowledge_base_path,
            "scripts_file": self.scripts_file,
            "llm_provider": self.llm_provider,
            "memory_type": "InMemorySaver",
            "multi_turn": True,
            "query_enhancement": True,
            "optimizations": {
                "llm_client": "Single initialization with functools.partial",
                "knowledge_retriever": "Single initialization with functools.partial",
                "conversation_context": "Cached in state to avoid repeated computation",
            },
        }

    def health_check(self) -> Dict:
        """Check if all components are working"""

        health_status = {
            "agent": True,
            "workflow": True,
            "memory": True,
            "llm": False,
            "knowledge_base": False,
            "scripts": False,
            "overall": False,
        }

        try:
            # Test LLM client (already initialized)
            test_response = self.llm_client.generate("Hello", temperature=0.1)
            health_status["llm"] = bool(test_response)
            print("✅ LLM client working")

        except Exception as e:
            print(f"⚠️ LLM health check failed: {e}")
            health_status["llm"] = False

        try:
            # Test memory/checkpointer
            test_config = {"configurable": {"thread_id": "health_check"}}
            test_state = {"messages": [HumanMessage(content="test")]}

            # Try to invoke with memory
            self.app.invoke(test_state, test_config)
            health_status["memory"] = True
            print("✅ Memory/checkpointer working")

        except Exception as e:
            print(f"⚠️ Memory health check failed: {e}")
            health_status["memory"] = False

        try:
            # Test knowledge base (using pre-initialized retriever)
            if self.knowledge_base:
                # Test a simple search to verify it's working
                test_docs = self.knowledge_base.search("test", top_k=1)
                health_status["knowledge_base"] = True
                print("✅ Knowledge retriever working")
            else:
                health_status["knowledge_base"] = False
                print("❌ Knowledge retriever not initialized")

        except Exception as e:
            print(f"⚠️ Knowledge base health check failed: {e}")
            health_status["knowledge_base"] = False

        try:
            # Test scripts
            from src.agent.scripts import ConversationScripts

            scripts = ConversationScripts(self.scripts_file)
            health_status["scripts"] = len(scripts.get_all_script_types()) > 0

        except Exception as e:
            print(f"⚠️ Scripts health check failed: {e}")

        # Overall health
        health_status["overall"] = all(
            [
                health_status["agent"],
                health_status["memory"],
                health_status["llm"],
                health_status["knowledge_base"],
                health_status["scripts"],
            ]
        )

        return health_status


# Usage example and testing
if __name__ == "__main__":
    # Initialize agent
    agent = ViettelPayAgent()

    # Health check
    print("\n🏥 Health Check:")
    health = agent.health_check()
    for component, status in health.items():
        status_icon = "✅" if status else "❌"
        print(f"   {component}: {status_icon}")

    if not health["overall"]:
        print("\n⚠️ Some components are not healthy. Check requirements and data files.")
        exit(1)

    print(f"\n🤖 Agent ready! Workflow info: {agent.get_workflow_info()}")

    # Test multi-turn conversation with query enhancement
    test_thread = "test_conversation"

    print(
        f"\n🧪 Testing multi-turn conversation with query enhancement (thread: {test_thread}):"
    )

    test_messages = [
        "Xin chào!",
        "Mã lỗi 606 là gì?",
        "Làm sao khắc phục?",  # This should be enhanced to "làm sao khắc phục lỗi 606"
        "Còn lỗi nào khác tương tự không?",  # This should be enhanced with error context
        "Cảm ơn bạn!",
    ]

    for i, message in enumerate(test_messages, 1):
        print(f"\n--- Turn {i} ---")
        result = agent.process_message(message, test_thread)
        print(f"User: {message}")
        print(f"Bot: {result['response'][:150]}...")

        if result.get("enhanced_query") and result["enhanced_query"] != message:
            print(f"🚀 Query enhanced: {result['enhanced_query']}")

        # Show conversation history
        if i > 1:
            history = agent.get_conversation_history(test_thread)
            print(f"History length: {len(history)} messages")

    print(f"\n📜 Final conversation history:")
    history = agent.get_conversation_history(test_thread)
    for i, msg in enumerate(history, 1):
        print(f"  {i}. {msg['role']}: {msg['content'][:100]}...")