#!/usr/bin/env python3
"""
Zep Cloud Client
This module provides a client for interacting with the Zep Cloud API
using the official Zep Cloud SDK.
"""

import os
import logging
from typing import Dict, List, Optional, Any, Union
import json
import requests

from dotenv import load_dotenv

# Import the Zep Cloud SDK
try:
    from zep_cloud.client import Zep
    from zep_cloud import User
except ImportError:
    raise ImportError(
        "zep-cloud SDK not found. Install with: pip install zep-cloud"
    )

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

class ZepCloudClient:
    """Client for interacting with the Zep Cloud API"""
    
    def __init__(self):
        """Initialize the Zep Cloud client"""
        self.api_key = os.getenv("ZEP_API_KEY")
        self.fallback_mode = False
        
        if not self.api_key:
            logger.error("ZEP_API_KEY environment variable not set. Running in fallback mode.")
            self.fallback_mode = True
            return
        
        # Initialize the client
        try:
            self.client = Zep(api_key=self.api_key)
            logger.info("Zep Cloud client initialized successfully")
            # Test the connection by trying to list users
            self._test_connection()
        except Exception as e:
            logger.error(f"Failed to initialize Zep Cloud client: {str(e)}")
            logger.warning("Running in fallback mode.")
            self.fallback_mode = True
            
    def _test_connection(self):
        """Test the connection to Zep Cloud by making a simple API call"""
        try:
            # Try to list users as a connection test
            self.client.user.list_ordered()
            logger.info("✅ Connected to Zep Cloud API")
            self.fallback_mode = False
        except Exception as e:
            logger.warning(f"❌ Failed to connect to Zep Cloud API: {str(e)}")
            logger.warning("⚠️ Running in fallback mode")
            self.fallback_mode = True
        
    def list_users(self, limit: int = 100, cursor: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List all users
        
        Returns:
            List[Dict[str, Any]]: List of user objects
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. User listing simulated.")
            return []
            
        try:
            user_response = self.client.user.list_ordered()
            users = user_response.users or []
            
            # Convert User objects to dictionaries
            user_dicts = []
            for user in users:
                user_dict = {
                    "user_id": user.user_id,
                    "metadata": user.metadata if user.metadata else {}
                }
                user_dicts.append(user_dict)
                
            return user_dicts
            
        except Exception as e:
            logger.error(f"Error listing users: {str(e)}")
            return []
            
    def get_user(self, user_id: str) -> Optional[Dict[str, Any]]:
        """
        Get a user by ID
        
        Args:
            user_id (str): The user ID
            
        Returns:
            Optional[Dict[str, Any]]: User object if found, None otherwise
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. User retrieval simulated.")
            return {"user_id": user_id, "success": True, "fallback": True}
            
        try:
            user = self.client.user.get(user_id=user_id)
            
            if user:
                return {
                    "user_id": user.user_id,
                    "metadata": user.metadata if user.metadata else {},
                    "first_name": user.first_name,
                    "last_name": user.last_name,
                    "email": user.email
                }
            return None
            
        except Exception as e:
            logger.error(f"Error getting user {user_id}: {str(e)}")
            return None
            
    def create_user(self, user_id: str, metadata: Optional[Dict[str, Any]] = None, first_name: Optional[str] = None, last_name: Optional[str] = None, email: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Create a new user
        
        Args:
            user_id (str): The user ID
            metadata (Optional[Dict[str, Any]]): User metadata
            first_name (Optional[str]): User's first name
            last_name (Optional[str]): User's last name
            email (Optional[str]): User's email address
            
        Returns:
            Optional[Dict[str, Any]]: Created user object if successful, None otherwise
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. User creation simulated.")
            return {
                "user_id": user_id, 
                "metadata": metadata or {}, 
                "first_name": first_name, 
                "last_name": last_name, 
                "email": email, 
                "success": True, 
                "fallback": True
            }
            
        try:
            # Ensure metadata is not None
            metadata_dict = metadata if metadata is not None else {}
            
            # Create the user
            user = self.client.user.add(
                user_id=user_id,
                metadata=metadata_dict,
                first_name=first_name,
                last_name=last_name,
                email=email
            )
            
            return {
                "user_id": user.user_id,
                "metadata": user.metadata if user.metadata else {},
                "first_name": user.first_name,
                "last_name": user.last_name,
                "email": user.email
            }
            
        except Exception as e:
            logger.error(f"Error creating user {user_id}: {str(e)}")
            return None
            
    def update_user(self, user_id: str, metadata: Optional[Dict[str, Any]] = None, first_name: Optional[str] = None, last_name: Optional[str] = None, email: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Update a user's metadata and profile information
        
        Args:
            user_id (str): The user ID
            metadata (Optional[Dict[str, Any]]): User metadata
            first_name (Optional[str]): User's first name
            last_name (Optional[str]): User's last name 
            email (Optional[str]): User's email address
            
        Returns:
            Optional[Dict[str, Any]]: Updated user object if successful, None otherwise
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. User update simulated.")
            return {
                "user_id": user_id, 
                "metadata": metadata or {}, 
                "first_name": first_name, 
                "last_name": last_name, 
                "email": email, 
                "success": True, 
                "fallback": True
            }
            
        try:
            # Ensure metadata is not None
            metadata_dict = metadata if metadata is not None else {}
            
            # Update the user
            user = self.client.user.update(
                user_id=user_id,
                metadata=metadata_dict,
                first_name=first_name,
                last_name=last_name,
                email=email
            )
            
            return {
                "user_id": user.user_id,
                "metadata": user.metadata if user.metadata else {},
                "first_name": user.first_name,
                "last_name": user.last_name,
                "email": user.email
            }
            
        except Exception as e:
            logger.error(f"Error updating user {user_id}: {str(e)}")
            return None
            
    def delete_user(self, user_id: str) -> bool:
        """
        Delete a user
        
        Args:
            user_id (str): The user ID
            
        Returns:
            bool: True if successful, False otherwise
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. User deletion simulated.")
            return True
            
        try:
            self.client.user.delete(user_id=user_id)
            return True
            
        except Exception as e:
            logger.error(f"Error deleting user {user_id}: {str(e)}")
            return False
            
    def search_graph(self, user_id: str, query: str, limit: int = 10) -> Optional[Dict[str, Any]]:
        """
        Search the user's graph with a query
        
        Args:
            user_id (str): The user ID
            query (str): The search query
            limit (int): The maximum number of results to return
            
        Returns:
            Optional[Dict[str, Any]]: Search results if successful, None otherwise
        """
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. Graph search simulated.")
            return {
                "query": query, 
                "user_id": user_id, 
                "limit": limit,
                "edges": [],
                "nodes": [],
                "results": [],  
                "success": True,
                "summary": "No results found for query (fallback mode)",
                "fallback": True
            }
            
        try:
            # Keep queries concise as recommended by docs
            if len(query) > 8000:
                logger.warning(f"Search query exceeds recommended length. Truncating to 8000 characters.")
                query = query[:8000]
                
            # Call the graph search API
            search_results = self.client.graph.search(
                user_id=user_id,
                query=query,
                limit=limit
            )
            
            # Create a digestible result that can be serialized to JSON
            results = {
                "query": query,
                "user_id": user_id,
                "limit": limit,
                "success": True,
                "results": []  # Keep generic results array for backward compatibility
            }
            
            # Add edges (facts) if they exist
            if hasattr(search_results, 'edges') and search_results.edges:
                results["edges"] = []
                for edge in search_results.edges:
                    edge_data = {
                        "id": edge.id if hasattr(edge, 'id') else None,
                        "fact": edge.fact if hasattr(edge, 'fact') else None,
                        "created_at": str(edge.created_at) if hasattr(edge, 'created_at') else None,
                        "updated_at": str(edge.updated_at) if hasattr(edge, 'updated_at') else None,
                        "score": edge.score if hasattr(edge, 'score') else None
                    }
                    results["edges"].append(edge_data)
                    # Also add to generic results for backward compatibility
                    results["results"].append(edge_data)
            
            # Add nodes if they exist
            if hasattr(search_results, 'nodes') and search_results.nodes:
                results["nodes"] = []
                for node in search_results.nodes:
                    node_data = {
                        "id": node.id if hasattr(node, 'id') else None,
                        "label": node.label if hasattr(node, 'label') else None,
                        "attributes": node.attributes if hasattr(node, 'attributes') else {},
                        "score": node.score if hasattr(node, 'score') else None
                    }
                    results["nodes"].append(node_data)
                    # Also add to generic results for backward compatibility
                    results["results"].append(node_data)
            
            # Add a summary field to help Claude understand the results
            if len(results["results"]) > 0:
                results["summary"] = f"Found {len(results['results'])} results for query '{query}'"
                if "nodes" in results and len(results["nodes"]) > 0:
                    results["summary"] += f", including {len(results['nodes'])} nodes"
                if "edges" in results and len(results["edges"]) > 0:
                    results["summary"] += f", including {len(results['edges'])} edges/facts"
            else:
                results["summary"] = f"No results found for query '{query}'"
            
            return results
            
        except Exception as e:
            logger.error(f"Error searching graph for user {user_id}: {str(e)}")
            return None
            
    def add_graph_data(self, user_id, data, data_type="text"):
        """
        Add data to a user's graph in Zep Cloud.
        
        Args:
            user_id: The unique identifier for the user
            data: The data to add to the graph (string or dict)
            data_type: The type of data, can be "text", "json", or "message"
            
        Returns:
            A dictionary with information about the added data
        """
        if not user_id:
            return {"error": "User ID is required", "success": False}
        
        if data is None:
            return {"error": "Data is required", "success": False}
        
        valid_types = ["text", "json", "message"]
        if data_type not in valid_types:
            return {
                "error": f"Invalid data type: {data_type}. Must be one of {valid_types}",
                "success": False
            }
            
        if self.fallback_mode:
            logger.warning("⚠️ Running in fallback mode. Graph data addition simulated.")
            return {
                "success": True,
                "user_id": user_id,
                "data_type": data_type,
                "data_length": len(str(data)),
                "fallback": True,
                "response": {
                    "uuid": "simulated-uuid",
                    "content": "Simulated content (fallback mode)",
                    "created_at": "simulated-timestamp",
                    "processed": True
                }
            }
        
        # Convert dict to JSON string if needed
        if data_type == "json" and isinstance(data, dict):
            try:
                data = json.dumps(data)
                logging.info(f"Converted Python dict to JSON string: {type(data)}")
            except Exception as e:
                return {
                    "error": f"Error converting dict to JSON: {str(e)}",
                    "success": False
                }
        
        # Check data size limits (after possible conversion)
        data_len = len(data) if isinstance(data, str) else len(str(data))
        if data_len > 100000:  # 100KB limit
            return {
                "error": f"Data size of {data_len} bytes exceeds 100KB limit",
                "success": False
            }
        
        # Additional validation for JSON data
        if data_type == "json":
            try:
                # Verify it's valid JSON
                if isinstance(data, str):
                    try:
                        json_data = json.loads(data)
                    except json.JSONDecodeError as e:
                        # Check for JSON with extra quotes
                        if data.startswith('"') and data.endswith('"') and len(data) > 2:
                            try:
                                # Try to parse the inner content without outer quotes
                                inner_data = data[1:-1]
                                if inner_data.lstrip().startswith(('{', '[')):
                                    json_data = json.loads(inner_data)
                                    data = inner_data
                                    logging.info(f"Fixed JSON by removing outer quotes: {type(json_data)}")
                                else:
                                    raise e
                            except Exception:
                                # If that didn't work, re-raise the original error
                                raise e
                        else:
                            raise e
                else:
                    # If already a dict or other Python object, convert to JSON string
                    json_data = data
                    data = json.dumps(data)
                
                # If we get here, data is valid JSON
                logging.info(f"Validated JSON data: {type(json_data)}")
            except json.JSONDecodeError as e:
                return {
                    "error": f"Invalid JSON data: {str(e)}",
                    "success": False
                }
            except Exception as e:
                return {
                    "error": f"Error processing JSON data: {str(e)}",
                    "success": False
                }
        
        # Use the Zep SDK to add the data to the graph
        try:
            # Call the graph add API
            response = self.client.graph.add(
                user_id=user_id,
                type=data_type,
                data=data
            )
            
            # Create a return value that can be serialized to JSON
            result = {
                "success": True,
                "user_id": user_id,
                "data_type": data_type,
                "data_length": len(data),
                "response": {
                    "uuid": response.uuid if hasattr(response, 'uuid') else None,
                    "content": response.content if hasattr(response, 'content') else None,
                    "created_at": str(response.created_at) if hasattr(response, 'created_at') else None,
                    "processed": response.processed if hasattr(response, 'processed') else None
                }
            }
            
            return result
            
        except Exception as e:
            error_message = f"Exception adding graph data: {str(e)}"
            logging.error(error_message)
            return {
                "error": error_message,
                "success": False
            } 