import os
import pandas as pd
import re
import json
from typing import List, Dict, Any, Optional
from dataclasses import dataclass

# Modern LangChain imports
try:
    from langchain_core.tools import tool
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
    from langchain_core.prompts import ChatPromptTemplate
    from langchain_core.output_parsers import JsonOutputParser
    LANGCHAIN_AVAILABLE = True
except ImportError as e:
    print(f"LangChain import error: {e}")
    LANGCHAIN_AVAILABLE = False

# Skyrelis monitoring integration
try:
    from skyrelis import observe_agent
    import skyrelis
    MONITORING_AVAILABLE = True
    print("✅ Skyrelis monitoring library loaded successfully")
    print(f"   Version: {skyrelis.__version__}")
    print(f"   LangChain support: {'langchain' in skyrelis.get_supported_frameworks()}")
except ImportError as e:
    print(f"Warning: skyrelis not installed, monitoring disabled: {e}")
    MONITORING_AVAILABLE = False


@dataclass
class AgentResponse:
    """Structured response from the agent"""
    content: str
    tool_calls: List[Dict[str, Any]] = None
    metadata: Dict[str, Any] = None


class ModernRAGAgent:
    """
    Modern LangChain 1.0+ Customer Service Agent
    
    Uses direct LLM function calling instead of deprecated AgentExecutor.
    Implements a simple message-based flow with structured tools.
    """
    
    def __init__(self, csv_path: str, api_key: str = None, monitor_url: str = None):
        self.csv_path = csv_path
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.monitor_url = monitor_url
        
        # Data storage
        self.df = None
        self.skyrelis_data = {}
        self.conversation_history = []
        
        # LLM and tools
        self.llm = None
        self.tools = []
        
        # Monitoring setup
        self._setup_monitoring()
        
        # Initialize components
        self._load_data_sources()
        if LANGCHAIN_AVAILABLE:
            self._init_llm()
            self._setup_tools()
        else:
            print("Warning: LangChain not available, using fallback mode")
            self.tools = []
            self.llm_with_tools = None

    def _setup_monitoring(self):
        """Setup Skyrelis monitoring configuration"""
        if self.monitor_url and MONITORING_AVAILABLE:
            print(f"🔒 Skyrelis monitoring enabled with URL: {self.monitor_url}")
            print(f"   Agent Name: customer_service_agent")
            print(f"   Security Level: production")
            self._monitoring_enabled = True
            
            # Store monitoring configuration for use with @observe decorator
            self._monitor_config = {
                "remote_observer_url": self.monitor_url,
                "agent_name": "customer_service_agent", 
                "security_level": "production",
                "enable_audit_logging": True
            }
        else:
            if self.monitor_url and not MONITORING_AVAILABLE:
                print("⚠️  Skyrelis monitoring URL provided but library not available")
            else:
                print("ℹ️  Skyrelis monitoring disabled - no monitor URL provided")
            self._monitoring_enabled = False
            self._monitor_config = None

    def _load_data_sources(self):
        """Load all data sources: CSV and Skyrelis corpus"""
        print("Loading data sources...")
        
        # Load CSV data
        self._load_csv()
        
        # Load Skyrelis corpus
        self._load_skyrelis_corpus()
        
        print("Data sources loaded successfully!")

    def _load_csv(self):
        """Load customer support CSV data"""
        try:
            print(f"Loading CSV data from: {self.csv_path}")
            self.df = pd.read_csv(self.csv_path)
            print(f"Loaded {len(self.df)} rows from CSV")
        except Exception as e:
            print(f"Error loading CSV: {e}")
            self.df = pd.DataFrame()  # Empty fallback

    def _load_skyrelis_corpus(self):
        """Load Skyrelis company information from corpus files"""
        self.skyrelis_data = {
            "company_info": "",
            "security_platform": "",
            "pricing_roi": "",
            "leadership_team": "",
            "contact_support": "",
            "innovation_circle": ""
        }
        
        # Look for corpus files
        data_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'data')
        jsonl_path = os.path.join(data_dir, "skyrelis_corpus.jsonl")
        txt_path = os.path.join(data_dir, "skyrelis_corpus.txt")
        
        try:
            if os.path.exists(jsonl_path):
                self._load_jsonl_corpus(jsonl_path)
            elif os.path.exists(txt_path):
                self._load_txt_corpus(txt_path)
            else:
                print(f"No Skyrelis corpus files found in {data_dir}, using fallback data...")
                self._set_fallback_skyrelis_data()
                
        except Exception as e:
            print(f"Error loading Skyrelis corpus: {e}")
            self._set_fallback_skyrelis_data()

    def _load_jsonl_corpus(self, jsonl_path: str):
        """Load structured JSONL corpus data"""
        print(f"Loading Skyrelis data from JSONL corpus: {jsonl_path}")
        
        with open(jsonl_path, 'r') as f:
            corpus_data = [json.loads(line) for line in f]
        
        # Organize data by sections
        for entry in corpus_data:
            text = entry.get('text', '')
            source = entry.get('source', '')
            
            if 'Home' in source:
                self.skyrelis_data["company_info"] += text + "\n\n"
                self.skyrelis_data["security_platform"] += text + "\n\n"
            elif 'Company' in source or 'Blog' in source:
                self.skyrelis_data["leadership_team"] += text + "\n\n"
            elif 'ROI' in source or 'Events' in source:
                self.skyrelis_data["pricing_roi"] += text + "\n\n"
            elif 'Innovation Circle' in source:
                self.skyrelis_data["innovation_circle"] += text + "\n\n"
            elif 'Contact' in source:
                self.skyrelis_data["contact_support"] += text + "\n\n"

    def _load_txt_corpus(self, txt_path: str):
        """Load simple text corpus data"""
        print(f"Loading Skyrelis data from TXT corpus: {txt_path}")
        
        with open(txt_path, 'r') as f:
            content = f.read()
        
        # Extract key information
        if "info@skyrelis.com" in content:
            self.skyrelis_data["contact_support"] = """
Contact Skyrelis:
- Email: info@skyrelis.com
- Location: San Mateo, California

For more information about our AI security platform, please reach out to our team.
"""
        
        if "Skyrelis" in content:
            self.skyrelis_data["company_info"] = """
Skyrelis is an AI security platform that provides runtime protection for autonomous AI agents.
We help enterprises secure AI agents, tools, and data with always-on behavioral monitoring 
and runtime policy enforcement—so they can scale adoption without losing control.
"""

    def _set_fallback_skyrelis_data(self):
        """Set fallback Skyrelis data if corpus files are not available"""
        self.skyrelis_data.update({
            "contact_support": """
Contact Skyrelis:
- Email: info@skyrelis.com
- Location: San Mateo, California
""",
            "company_info": """
Skyrelis is an AI security platform that provides runtime protection for autonomous AI agents.
We help enterprises secure AI agents, tools, and data with always-on behavioral monitoring 
and runtime policy enforcement—so they can scale adoption without losing control.
"""
        })

    def _init_llm(self):
        """Initialize the LLM with modern patterns"""
        try:
            print("Initializing LLM...")
            if self.api_key:
                print(f"🔑 Using OpenAI API Key: {self.api_key[:8]}...{self.api_key[-4:]}")
            else:
                print("🔑 No OpenAI API Key provided")
            
            self.llm = ChatOpenAI(
                model="gpt-4o-mini",
                temperature=0,
                api_key=self.api_key
            )
            
            print("✅ OpenAI ChatGPT (gpt-4o-mini) LLM initialized successfully")
            print(f"🤖 Model: {self.llm.model_name}")
            print(f"🌡️  Temperature: {self.llm.temperature}")
            
            # Test OpenAI connection with a simple call
            try:
                print("🔍 Testing OpenAI connection...")
                test_response = self.llm.invoke([{"role": "user", "content": "Hello, are you working?"}])
                print(f"✅ OpenAI connection confirmed! Response: {test_response.content[:50]}...")
            except Exception as test_e:
                print(f"⚠️  OpenAI connection test failed: {test_e}")
            
        except Exception as e:
            print(f"❌ Error initializing OpenAI LLM: {e}")
            self.llm = None

    def _setup_tools(self):
        """Setup tools using modern LangChain patterns"""
        print("Setting up tools...")
        
        # Create tools (always create them, even without LLM for fallback mode)
        self.tools = [
            self._create_skyrelis_info_tool(),
            self._create_general_support_tool()
        ]
        
        # Bind tools to LLM for function calling if LLM is available
        if self.llm:
            self.llm_with_tools = self.llm.bind_tools(self.tools)
            print(f"Tools setup complete: {len(self.tools)} tools available with LLM binding")
        else:
            self.llm_with_tools = None
            print(f"Tools setup complete: {len(self.tools)} tools available (fallback mode)")

    def _create_skyrelis_info_tool(self):
        """Create the Skyrelis information search tool"""
        @tool
        def skyrelis_info_search(query: str) -> str:
            """Search for information about Skyrelis company, products, pricing, features, team, or support.
            Use this when customers ask about Skyrelis specifically.
            
            Args:
                query: The user's question about Skyrelis
                
            Returns:
                Relevant Skyrelis information
            """
            query_lower = query.lower()
            relevant_info = []
            
            # Handle contact requests
            if any(word in query_lower for word in ['email', 'contact', 'reach', 'get in touch']):
                relevant_info.append(self.skyrelis_data["contact_support"])
            
            # Handle company info requests
            if any(word in query_lower for word in ['what is', 'about', 'company', 'skyrelis']) or not relevant_info:
                if self.skyrelis_data["company_info"].strip():
                    relevant_info.append(self.skyrelis_data["company_info"])
            
            # Handle team/leadership requests
            if any(word in query_lower for word in ['team', 'leadership', 'ceo', 'founder']):
                if self.skyrelis_data["leadership_team"].strip():
                    relevant_info.append(self.skyrelis_data["leadership_team"])
            
            # Handle pricing/ROI requests
            if any(word in query_lower for word in ['price', 'cost', 'roi', 'pricing']):
                if self.skyrelis_data["pricing_roi"].strip():
                    relevant_info.append(self.skyrelis_data["pricing_roi"])
            
            # Handle platform/security requests
            if any(word in query_lower for word in ['platform', 'security', 'features', 'product']):
                if self.skyrelis_data["security_platform"].strip():
                    relevant_info.append(self.skyrelis_data["security_platform"])
            
            # Default to contact if nothing specific found
            if not relevant_info:
                relevant_info.append(self.skyrelis_data["contact_support"])
            
            return "Skyrelis Information:\n\n" + "\n\n".join(relevant_info)
        
        return skyrelis_info_search

    def _create_general_support_tool(self):
        """Create the general customer support search tool"""
        @tool
        def general_support_search(query: str) -> str:
            """Search the general customer service knowledge base for support information.
            Use this for general customer service questions, billing issues, account problems, etc.
            
            Args:
                query: The customer's support question
                
            Returns:
                Relevant support information from the knowledge base
            """
            relevant_contexts = self._find_relevant_context(query)
            
            if not relevant_contexts:
                return "No relevant information found in the knowledge base for this query."
            
            return "Relevant support information:\n\n" + "\n\n".join(relevant_contexts[:3])
        
        return general_support_search

    def _find_relevant_context(self, question: str, top_k: int = 3) -> List[str]:
        """Simple keyword-based context retrieval from CSV data"""
        if self.df is None or len(self.df) == 0:
            return []
        
        # Extract keywords from question
        keywords = re.findall(r'\b\w+\b', question.lower())
        keywords = [k for k in keywords if len(k) > 3]  # Filter short words
        
        if not keywords:
            return []
        
        # Score each row based on keyword matches
        scores = []
        for idx, row in self.df.iterrows():
            row_text = ' '.join(map(str, row.values)).lower()
            score = sum(1 for keyword in keywords if keyword in row_text)
            scores.append((score, idx, row_text))
        
        # Get top-k most relevant rows
        scores.sort(reverse=True)
        relevant_contexts = []
        for score, idx, text in scores[:top_k]:
            if score > 0:  # Only include rows with some keyword matches
                relevant_contexts.append(text[:1000])  # Limit context length
        
        return relevant_contexts

    def _is_skyrelis_query(self, question: str) -> bool:
        """Check if the question is about Skyrelis"""
        skyrelis_keywords = [
            'skyrelis', 'skyr', 'contact', 'email', 'ceo', 'founder', 
            'innovation circle', 'pricing', 'cost', 'team', 'leadership',
            'company', 'platform', 'security', 'ai security', 'monitoring'
        ]
        question_lower = question.lower()
        return any(keyword in question_lower for keyword in skyrelis_keywords)

    def ask(self, user_msg: str, customer_id: str = None) -> str:
        """
        Main method to ask the agent a question using modern LangChain patterns
        
        Args:
            user_msg: The user's question
            customer_id: Optional customer identifier
            
        Returns:
            Agent's response as a string
        """
        # Input validation
        if not user_msg or len(user_msg.strip()) == 0:
            return "I'm sorry, but I need a question to help you. Could you please provide more details?"
        
        # Process the query with optional Skyrelis monitoring
        return self._process_query(user_msg, customer_id)
    
    def _process_query(self, user_msg: str, customer_id: str = None) -> str:
        """Internal method to process the query (wrapped by monitoring)"""
        # Try modern LangChain approach first
        if LANGCHAIN_AVAILABLE and self.llm_with_tools:
            try:
                return self._handle_with_modern_langchain(user_msg, customer_id)
            except Exception as e:
                print(f"Modern LangChain error: {e}")
                print("Falling back to simple mode...")
        
        # Fallback to simple keyword-based routing
        return self._handle_with_fallback(user_msg, customer_id)

    def _handle_with_modern_langchain(self, user_msg: str, customer_id: str = None) -> str:
        """Handle query using modern LangChain function calling"""
        print(f"🤖 Processing with modern LangChain + OpenAI: {user_msg}")
        
        # Create system message
        system_msg = SystemMessage(content="""You are Skyrelis Customer Support.

RULES:
- Be direct and brief. Default to ONE sentence; never exceed 3 sentences.
- Answer ONLY what was asked; no background, no speculation.
- If info isn't in context/tools: say "I don't have that information yet."
- No apologies unless there's an actual error.
- No disclaimers, no marketing language, no emojis.

Use the available tools to search for information when needed.""")
        
        # Create user message
        user_message = HumanMessage(content=user_msg)
        
        # Create message sequence
        messages = [system_msg, user_message]
        
        # Add Skyrelis monitoring if enabled
        print("📡 Calling OpenAI API...")
        if self._monitoring_enabled and MONITORING_AVAILABLE:
            print("🔒 Skyrelis monitoring active - capturing security telemetry")
            # Use function decorator for method calls
            try:
                monitored_invoke = observe_function(
                    remote_observer_url=self.monitor_url,
                    agent_name="customer_service_query"
                )(self.llm_with_tools.invoke)
                response = monitored_invoke(messages)
                print("✅ Skyrelis function monitoring successful!")
            except Exception as e:
                print(f"⚠️  Skyrelis monitoring failed: {e}")
                print("   Falling back to direct LLM call")
                response = self.llm_with_tools.invoke(messages)
        else:
            response = self.llm_with_tools.invoke(messages)
        
        print(f"✅ OpenAI API response received (type: {type(response).__name__})")
        
        # Handle tool calls if present
        if response.tool_calls:
            print(f"🔧 OpenAI requested {len(response.tool_calls)} tool call(s)")
            # Execute tool calls
            messages.append(response)
            
            for tool_call in response.tool_calls:
                tool_name = tool_call["name"]
                tool_args = tool_call["args"]
                print(f"   🛠️  Executing tool: {tool_name} with args: {tool_args}")
                
                # Find and execute the tool
                tool_result = None
                for tool in self.tools:
                    if tool.name == tool_name:
                        tool_result = tool.invoke(tool_args)
                        break
                
                if tool_result:
                    print(f"   ✅ Tool result: {str(tool_result)[:100]}...")
                    # Add tool result to messages
                    from langchain_core.messages import ToolMessage
                    messages.append(ToolMessage(
                        content=str(tool_result),
                        tool_call_id=tool_call["id"]
                    ))
            
            # Get final response after tool execution
            print("📡 Calling OpenAI API again for final response...")
            if self._monitoring_enabled and MONITORING_AVAILABLE:
                print("🔒 Skyrelis monitoring active - capturing final response telemetry")
                try:
                    monitored_final_invoke = observe_function(
                        remote_observer_url=self.monitor_url,
                        agent_name="customer_service_final_response"
                    )(self.llm.invoke)
                    final_response = monitored_final_invoke(messages)
                    print("✅ Final Skyrelis monitoring successful!")
                except Exception as e:
                    print(f"⚠️  Final Skyrelis monitoring failed: {e}")
                    final_response = self.llm.invoke(messages)
            else:
                final_response = self.llm.invoke(messages)
            print(f"✅ Final OpenAI response received")
            response_content = final_response.content
        else:
            print("ℹ️  No tool calls needed, using direct OpenAI response")
            response_content = response.content
        
        # Store conversation history
        self.conversation_history.append({
            "user": user_msg,
            "agent": response_content,
            "customer_id": customer_id
        })
        
        return response_content

    def _handle_with_fallback(self, user_msg: str, customer_id: str = None) -> str:
        """Fallback handler using simple keyword matching"""
        print(f"Processing with fallback mode: {user_msg}")
        
        response = ""
        
        if self._is_skyrelis_query(user_msg):
            user_msg_lower = user_msg.lower()
            
            # Handle specific email requests
            if 'email' in user_msg_lower and not any(word in user_msg_lower for word in ['about', 'info', 'information', 'details']):
                response = "Skyrelis email: info@skyrelis.com"
            elif any(word in user_msg_lower for word in ['email', 'contact', 'reach', 'get in touch']):
                response = f"{self.skyrelis_data['contact_support']}"
            else:
                response = f"About Skyrelis:\n\n{self.skyrelis_data['company_info']}\n\nFor more information: {self.skyrelis_data['contact_support']}"
        else:
            # Handle general customer service queries
            relevant_contexts = self._find_relevant_context(user_msg)
            
            if relevant_contexts:
                response = f"Based on your inquiry, here's what I found:\n\n{relevant_contexts[0][:500]}"
                if len(relevant_contexts) > 1:
                    response += f"\n\nAdditional information:\n{relevant_contexts[1][:300]}"
            else:
                response = "I understand you need help with that issue. For billing concerns and account-specific problems, please contact our support team directly at info@skyrelis.com with your account details so we can assist you properly."
        
        # Store conversation history
        self.conversation_history.append({
            "user": user_msg,
            "agent": response,
            "customer_id": customer_id
        })
        
        return response

    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the loaded data and agent status"""
        stats = {
            "langchain_available": LANGCHAIN_AVAILABLE,
            "monitoring_available": MONITORING_AVAILABLE,
            "monitoring_enabled": self._monitoring_enabled,
            "llm_available": self.llm is not None,
            "tools_count": len(self.tools) if self.tools else 0,
            "conversation_count": len(self.conversation_history)
        }
        
        if self.df is not None:
            stats.update({
                "csv_rows": len(self.df),
                "csv_columns": list(self.df.columns),
                "csv_memory_usage": self.df.memory_usage(deep=True).sum()
            })
        else:
            stats["csv_status"] = "Not loaded"
        
        # Skyrelis data stats
        skyrelis_sections = {k: len(v.strip()) > 0 for k, v in self.skyrelis_data.items()}
        stats["skyrelis_sections_loaded"] = skyrelis_sections
        
        return stats

    def clear_conversation_history(self):
        """Clear the conversation history"""
        self.conversation_history = []
        print("Conversation history cleared")

    def get_conversation_history(self) -> List[Dict[str, Any]]:
        """Get the conversation history"""
        return self.conversation_history.copy()


# Convenience function for backward compatibility
def create_agent(csv_path: str, api_key: str = None, monitor_url: str = None) -> ModernRAGAgent:
    """Create a new ModernRAGAgent instance"""
    return ModernRAGAgent(csv_path=csv_path, api_key=api_key, monitor_url=monitor_url)


if __name__ == "__main__":
    # Example usage
    csv_path = "../data/Bitext_Sample_Customer_Support_Training_Dataset_27K_responses-v11.csv"
    agent = create_agent(csv_path)
    
    print("Agent Stats:", agent.get_stats())
    
    # Test queries
    test_queries = [
        "What is Skyrelis?",
        "How can I contact Skyrelis?",
        "I need help with my order cancellation",
        "What's your email address?"
    ]
    
    for query in test_queries:
        print(f"\nQ: {query}")
        print(f"A: {agent.ask(query)}")
