import requests
import json
import logging
import uuid
import time
import os
from typing import List, Optional, Union, Callable, Any, Dict
from types import SimpleNamespace
from concurrent.futures import ThreadPoolExecutor, as_completed
from .config import DEFAULT_API_URL, CHATBOT_ENDPOINT, EVALUATION_ENDPOINT, SAFETY_EVALUATION_ENDPOINT, TRACE_PROJECT_NAME_ENDPOINT
from .conversation import ConversationHistory
from .model import EvaluationResult, ChatResponse, BatchEvaluationResult, BatchChatResult

logger = logging.getLogger(__name__)

class Client:
    """
    User-friendly client for InsightFinder AI SDK.
    
    This client provides easy-to-use methods for:
    - Single and batch chatting with streaming support and conversation history
    - Evaluation of prompts and responses with automatic project name generation
    - Safety evaluation for prompts
    
    The client automatically generates project names for evaluations by calling the API
    with the session_name and appending "-Prompt" to the result.
    """
    
    # Maximum context window size in characters before trimming old messages
    MAX_CONTEXT_WINDOW_SIZE = 8000

    def __init__(self, session_name: str, url: Optional[str] = None, username: Optional[str] = None, api_key: Optional[str] = None, enable_chat_evaluation: bool = True):
        """
        Initialize the client with user credentials and project settings.

        Args:
            session_name (str): Session name for chat requests and used to generate project name automatically
            url (str, optional): Custom API URL (defaults to https://ai.insightfinder.com)
            username (str, optional): Username for authentication (can be set via INSIGHTFINDER_USERNAME env var)
            api_key (str, optional): API key for authentication (can be set via INSIGHTFINDER_API_KEY env var)
            enable_chat_evaluation (bool): Whether to display evaluation and safety results in chat responses (default: True)
        
        Note:
            - session_name is used for both chat operations and to automatically generate the project_name for evaluation operations
            - The project name is automatically generated by calling the API with session_name and appending "-Prompt"
        
        Environment Variables:
            INSIGHTFINDER_USERNAME: Username for authentication
            INSIGHTFINDER_API_KEY: API key for authentication
        
        Example:
            # Using parameters
            client = Client(
                session_name="llm-eval-test", 
                username="john_doe",
                api_key="your_api_key_here",
                enable_chat_evaluation=True
            )
            
            # Using environment variables
            # export INSIGHTFINDER_USERNAME="john_doe"
            # export INSIGHTFINDER_API_KEY="your_api_key_here"
            client = Client(
                session_name="llm-eval-test",
                enable_chat_evaluation=True
            )
        """
        
        # Get credentials from parameters or environment variables
        self.username = username or os.getenv('INSIGHTFINDER_USERNAME')
        self.api_key = api_key or os.getenv('INSIGHTFINDER_API_KEY')
        
        if not self.username:
            raise ValueError("Username must be provided either as parameter or INSIGHTFINDER_USERNAME environment variable")
        if not self.api_key:
            raise ValueError("API key must be provided either as parameter or INSIGHTFINDER_API_KEY environment variable")
        if not session_name:
            raise ValueError("Session name cannot be empty")
        
        self.session_name = session_name
        self.enable_evaluations = enable_chat_evaluation
        
        # Set base URL with default fallback
        self.base_url = url if url else DEFAULT_API_URL
        if not self.base_url.endswith('/'):
            self.base_url += '/'
            
        # Construct API URLs
        self.chat_url = self.base_url + CHATBOT_ENDPOINT
        self.evaluation_url = self.base_url + EVALUATION_ENDPOINT  
        self.safety_url = self.base_url + SAFETY_EVALUATION_ENDPOINT
        self.trace_project_name_url = self.base_url + TRACE_PROJECT_NAME_ENDPOINT
        
        # Initialize conversation history
        self.conversation = ConversationHistory()
        
        # Generate project name dynamically
        self.project_name = self._get_project_name()

    def _get_headers(self) -> dict:
        """Get authentication headers."""
        return {
            'X-Api-Key': self.api_key,
            'X-User-Name': self.username,
            'Content-Type': 'application/json'
        }

    def _generate_trace_id(self) -> str:
        """Generate a unique trace ID."""
        return str(uuid.uuid4())

    def _get_timestamp(self) -> int:
        """Get current timestamp in milliseconds."""
        return int(time.time() * 1000)

    def _trim_context_if_needed(self):
        """Trim conversation history if it exceeds MAX_CONTEXT_WINDOW_SIZE."""
        while self.conversation.messages:
            # Calculate current context size
            current_context = self.conversation.to_string()
            if len(current_context) <= self.MAX_CONTEXT_WINDOW_SIZE:
                break
            
            # Remove the oldest message (first message)
            self.conversation.messages.pop(0)
            
            # If we only have one message left and it's still too long, we have a problem
            if len(self.conversation.messages) <= 1:
                break

    def _get_project_name(self, session_name: Optional[str] = None) -> str:
        """
        Get the project name by calling the trace project name API and appending '-Prompt'.
        
        Args:
            session_name (Optional[str]): Session name to use. If None, uses the default session name.
        
        Returns:
            str: The generated project name for evaluations
        """
        # Use provided session_name or fall back to default
        effective_session_name = session_name or self.session_name
        
        data = {
            "userCreatedModelName": effective_session_name
        }
        
        try:
            response = requests.post(
                self.trace_project_name_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= response.status_code < 300):
                raise ValueError(f"Trace project name API error {response.status_code}: {response.text}")
            
            # The API returns raw text, not JSON
            trace_project_name = response.text.strip()
            
            # Append "-Prompt" to the trace project name
            return f"{trace_project_name}-Prompt"
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get trace project name: {str(e)}")

    def chat(self, messages: Union[str, List[Dict[str, str]]], stream: bool = False, chat_history: bool = False, session_name: Optional[str] = None) -> ChatResponse:
        """
        Send a chat message and get response. Supports both simple strings and conversation history.
        
        Args:
            messages (Union[str, List[Dict[str, str]]]): Your message/question or conversation history
                - String: Simple prompt like "What is the capital of France?"
                - List: Conversation history like [{"role": "user", "content": "Hello"}, {"role": "system", "content": "Hi there!"}]
            stream (bool): Whether to show streaming response (default: False)
            chat_history (bool): Whether to store this conversation in client history for future context (default: False)
            session_name (Optional[str]): Session name to use for this request. If None, uses the default session name.
        
        Returns:
            ChatResponse: Response object with formatted display including evaluations (if enabled)
            
        Examples:
            # Simple chat without storing history
            response = client.chat("What is the capital of France?")
            
            # Chat with conversation storage enabled
            client.chat("Hello", chat_history=True)  # Stored
            client.chat("How are you?", chat_history=True)  # Uses previous context
            
            # Chat with custom session name
            response = client.chat("Hello", session_name="custom-session")
            
            # Provide explicit conversation history
            response = client.chat([
                {"role": "user", "content": "knock knock."},
                {"role": "system", "content": "Who's there?"},
                {"role": "user", "content": "Orange."}
            ])
        """
        # Handle different input types
        if isinstance(messages, str):
            if not messages.strip():
                raise ValueError("Prompt cannot be empty")
            
            # If store=True and we have conversation history, add to existing conversation
            if chat_history and self.conversation.messages:
                # Add user message to existing conversation
                self.conversation.add_user_message(messages)
                # Trim context if needed
                self._trim_context_if_needed()
                prompt_for_api = self.conversation.to_string()
                prompt_for_display = messages  # Current prompt only
            else:
                # Simple string prompt (no context from history)
                prompt_for_api = f'{{"role": "user", "content": "{messages}"}}'
                prompt_for_display = messages
                
                # If store=True, start new conversation
                if chat_history:
                    self.conversation.clear()  # Start fresh
                    self.conversation.add_user_message(messages)
                    
        elif isinstance(messages, list):
            if not messages:
                raise ValueError("Messages list cannot be empty")
            
            # Validate message format
            for msg in messages:
                if not isinstance(msg, dict) or 'role' not in msg or 'content' not in msg:
                    raise ValueError("Each message must be a dict with 'role' and 'content' keys")
            
            # Convert to string format for API (as per your existing test format)
            prompt_for_api = ', '.join([json.dumps(msg, separators=(',', ':')) for msg in messages])
            
            # Extract the last user message as the current prompt for display
            last_user_message = None
            for msg in reversed(messages):
                if msg.get('role') == 'user':
                    last_user_message = msg.get('content', '')
                    break
            prompt_for_display = last_user_message or messages[-1].get('content', '') if messages else ''
            
            # If store=True, replace conversation history with provided messages
            if chat_history:
                self.conversation.messages = messages.copy()
                self._trim_context_if_needed()
        else:
            raise ValueError("Messages must be either a string or a list of message dictionaries")
            
        # Use provided session_name or fall back to default
        effective_session_name = session_name or self.session_name
        
        # Prepare request data (using 'prompt' as per the original API)
        data = {
            'prompt': prompt_for_api,
            'userCreatedModelName': effective_session_name,
        }
        
        try:
            response = requests.post(
                self.chat_url,
                headers=self._get_headers(),
                json=data,
                stream=True
            )
            
            if not (200 <= response.status_code < 300):
                raise ValueError(f"API error {response.status_code}: {response.text}")
            
            # Process streaming response
            results = []
            stitched_response = ""
            evaluations = None
            trace_id = None
            model = None
            evaluation_buffer = ""  # Buffer to accumulate evaluation JSON
            in_evaluation_block = False
            
            for line in response.iter_lines(decode_unicode=True):    
                if line and line.startswith('data:'):
                    json_part = line[5:].strip()
                    if json_part and json_part != '[START]':
                        try:
                            chunk = json.loads(json_part)
                            results.append(chunk)
                            
                            # Extract metadata
                            if not model and "model" in chunk:
                                model = chunk["model"]
                            if "id" in chunk:
                                trace_id = chunk["id"]
                                
                            # Process content
                            if "choices" in chunk:
                                for choice in chunk["choices"]:
                                    delta = choice.get("delta", {})
                                    content = delta.get("content", "")
                                    
                                    # Check if we're starting an evaluation block
                                    if content and content.startswith("{") and "evaluations" in content:
                                        # Check if it's a complete JSON in one chunk
                                        if content.endswith("}"):
                                            try:
                                                eval_obj = json.loads(content)
                                                evaluations = eval_obj.get("evaluations")
                                                trace_id = eval_obj.get("traceId", trace_id)
                                            except json.JSONDecodeError:
                                                # If parsing fails, treat as start of multi-chunk evaluation
                                                in_evaluation_block = True
                                                evaluation_buffer = content
                                        else:
                                            # Start of multi-chunk evaluation
                                            in_evaluation_block = True
                                            evaluation_buffer = content
                                    elif in_evaluation_block:
                                        # We're in an evaluation block, accumulate content
                                        evaluation_buffer += content
                                        
                                        # Try to parse the accumulated JSON
                                        try:
                                            eval_obj = json.loads(evaluation_buffer)
                                            evaluations = eval_obj.get("evaluations")
                                            trace_id = eval_obj.get("traceId", trace_id)
                                            in_evaluation_block = False
                                            evaluation_buffer = ""
                                        except json.JSONDecodeError:
                                            # Not complete JSON yet, continue accumulating
                                            pass
                                    else:
                                        # Regular response content
                                        stitched_response += content
                                        if stream and content:
                                            print(content, end='', flush=True)
                        except:
                            pass
            
            # Save system response to history if store=True
            if chat_history and stitched_response:
                self.conversation.add_assistant_message(stitched_response)
                # Trim context again after adding response if needed
                self._trim_context_if_needed()
            
            # Create response object
            chat_response = ChatResponse(
                response=stitched_response,
                prompt=prompt_for_display,
                evaluations=evaluations if self.enable_evaluations else None,
                trace_id=trace_id,
                model=model,
                raw_chunks=results,
                enable_evaluations=self.enable_evaluations,
                history=self.conversation.get_messages() if chat_history else []
            )
            
            return chat_response
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Request failed: {str(e)}")

    def get_context_info(self) -> Dict[str, Any]:
        """
        Get information about the current conversation context.
        
        Returns:
            Dict[str, Any]: Context information including size, message count, and whether it's near the limit
        """
        current_context = self.conversation.to_string()
        message_count = len(self.conversation.messages)
        context_size = len(current_context)
        
        return {
            "message_count": message_count,
            "context_size_chars": context_size,
            "max_context_size": self.MAX_CONTEXT_WINDOW_SIZE,
            "remaining_capacity": max(0, self.MAX_CONTEXT_WINDOW_SIZE - context_size),
            "usage_percentage": round((context_size / self.MAX_CONTEXT_WINDOW_SIZE) * 100, 1),
            "near_limit": context_size > (self.MAX_CONTEXT_WINDOW_SIZE * 0.8),  # 80% threshold
            "context_preview": current_context[:200] + "..." if len(current_context) > 200 else current_context
        }

    # Conversation History Management Methods
    
    def get_conversation_history(self) -> List[Dict[str, str]]:
        """
        Get the current conversation history.
        
        Returns:
            List[Dict[str, str]]: List of messages in the conversation
        """
        return self.conversation.get_messages()
    
    def clear_chat_history(self):
        """Clear the conversation history to start fresh."""
        self.conversation.clear()
    
    def add_to_conversation(self, role: str, content: str):
        """
        Manually add a message to the conversation history.
        
        Args:
            role (str): Either "user" or "system"
            content (str): The message content
        """
        self.conversation.add_message(role, content)
    
    def set_chat_history(self, data: Union[List[Dict[str, str]], Dict[str, Any]]):
        """
        Set the conversation history from either a list of messages or loaded chat history data.
        
        Args:
            data (Union[List[Dict[str, str]], Dict[str, Any]]): Either:
                - List of message dictionaries with 'role' and 'content' keys
                - Full loaded chat history data (will extract 'conversation' key)
        
        Examples:
            # Set from message list
            client.set_chat_history([
                {"role": "user", "content": "Hello"},
                {"role": "system", "content": "Hi there!"}
            ])
            
            # Set from loaded data
            data = client.load_chat_history("my_conversation.json")
            client.set_chat_history(data)  # Automatically uses data["conversation"]
            
            # Or explicitly use conversation
            client.set_chat_history(data["conversation"])
        """
        # Handle different input types
        if isinstance(data, list):
            # Direct list of messages
            messages = data
        elif isinstance(data, dict):
            # Check if it's loaded chat history data with 'conversation' key
            if 'conversation' in data:
                messages = data['conversation']
            else:
                # Assume it's a single message dict, wrap in list
                messages = [data]
        else:
            raise ValueError("Data must be either a list of messages or a dict with 'conversation' key")
        
        # Validate message format
        if not isinstance(messages, list):
            raise ValueError("Messages must be a list")
            
        for i, msg in enumerate(messages):
            if not isinstance(msg, dict):
                raise ValueError(f"Message {i} must be a dict")
            if 'role' not in msg or 'content' not in msg:
                raise ValueError(f"Message {i} must have 'role' and 'content' keys")
        
        self.conversation.messages = messages.copy()

    def retrieve_chat_history(self) -> List[Dict[str, str]]:
        """
        Retrieve the current conversation chat history.
        
        This is an alias for get_conversation_history() with a more intuitive name.
        
        Returns:
            List[Dict[str, str]]: List of messages in the conversation history
            
        Example:
            history = client.retrieve_chat_history()
            for msg in history:
                print(f"[{msg['role'].upper()}] {msg['content']}")
        """
        return self.conversation.get_messages()

    def save_chat_history(self, filename: Optional[str] = None) -> str:
        """
        Save the current conversation chat history to a JSON file.
        
        Args:
            filename (Optional[str]): The filename to save to. If None, generates a unique filename.
        
        Returns:
            str: The actual filename used for saving
            
        Example:
            # Save with custom filename
            saved_file = client.save_chat_history("my_conversation.json")
            
            # Save with auto-generated filename
            saved_file = client.save_chat_history()
            print(f"Conversation saved to: {saved_file}")
        """
        import json
        from datetime import datetime
        import os
        
        # Get current conversation history
        history = self.retrieve_chat_history()
        
        # Generate filename if not provided
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            session_safe = "".join(c for c in self.session_name if c.isalnum() or c in ('-', '_'))
            filename = f"chat_history_{session_safe}_{timestamp}.json"
        
        # Ensure filename has .json extension
        if not filename.endswith('.json'):
            filename += '.json'
        
        # Create the data structure to save
        save_data = {
            "session_name": self.session_name,
            "timestamp": datetime.now().isoformat(),
            "message_count": len(history),
            "conversation": history
        }
        
        # Save to file in current working directory
        filepath = os.path.join(os.getcwd(), filename)
        
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(save_data, f, indent=2, ensure_ascii=False)
            
            return filename
            
        except Exception as e:
            raise ValueError(f"Failed to save chat history to {filename}: {str(e)}")

    def load_chat_history(self, filename: str) -> Dict[str, Any]:
        """
        Load conversation chat history from a JSON file.
        
        Args:
            filename (str): The filename to load from
        
        Returns:
            Dict[str, Any]: The loaded chat history data including metadata
            
        Example:
            # Load and restore conversation
            data = client.load_chat_history("my_conversation.json")
            client.set_conversation_history(data["conversation"])
            print(f"Loaded {data['message_count']} messages from {data['timestamp']}")
        """
        import json
        import os
        
        # Add .json extension if not present
        if not filename.endswith('.json'):
            filename += '.json'
        
        # Try to find the file in current directory first
        filepath = os.path.join(os.getcwd(), filename)
        if not os.path.exists(filepath):
            # If not found, try the filename as-is (might be full path)
            filepath = filename
        
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # Validate the data structure
            if not isinstance(data, dict) or 'conversation' not in data:
                raise ValueError("Invalid chat history file format")
            
            return data
            
        except FileNotFoundError:
            raise ValueError(f"Chat history file not found: {filename}")
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in chat history file: {str(e)}")
        except Exception as e:
            raise ValueError(f"Failed to load chat history from {filename}: {str(e)}")

    def batch_chat(self, prompts: List[str], stream: bool = False, max_workers: int = 3, session_name: Optional[str] = None, enable_history: bool = False) -> BatchChatResult:
        """
        Send multiple chat messages in parallel or sequentially with conversation history.
        
        Args:
            prompts (List[str]): List of messages/questions
            stream (bool): Whether to show progress updates (default: False)
            max_workers (int): Number of parallel requests (default: 3, ignored when enable_history=True)
            session_name (Optional[str]): Session name to use for all requests. If None, uses the default session name.
            enable_history (bool): Whether to process prompts sequentially with conversation history (default: False)
                - If True: Processes prompts one by one, maintaining conversation context between them
                - If False: Processes prompts in parallel with no shared context
        
        Returns:
            BatchChatResult: Batch chat result object with summary statistics
            
        Example:
            # Parallel processing (default)
            prompts = ["Hello!", "What's the weather?", "Tell me a joke"]
            responses = client.batch_chat(prompts)
            
            # Sequential processing with conversation history
            prompts = ["Hello!", "What's my name?", "Tell me about our conversation"]
            responses = client.batch_chat(prompts, enable_history=True)
            
            # With custom session name
            responses = client.batch_chat(prompts, session_name="custom-session")
        """
        if not prompts:
            raise ValueError("Prompts list cannot be empty")
        
        if enable_history:
            # Sequential processing with conversation history
            sequential_results: List[ChatResponse] = []
            
            # Save current conversation state
            original_conversation = self.conversation.get_messages().copy()
            
            try:
                # Clear conversation to start fresh for batch
                self.conversation.clear()
                
                for i, prompt in enumerate(prompts):
                    if stream:
                        print(f"\n--- Processing prompt {i+1}/{len(prompts)} ---")
                        print(f"Prompt: {prompt}")
                    
                    # Process with conversation history enabled
                    response = self.chat(prompt, stream=stream, chat_history=True, session_name=session_name)
                    sequential_results.append(response)
                    
                    # Check context size and trim if needed
                    context_info = self.get_context_info()
                    if context_info['near_limit']:
                        if stream:
                            print(f"Context approaching limit ({context_info['usage_percentage']}%), trimming...")
                        self._trim_context_if_needed()
                
                return BatchChatResult(sequential_results)
                
            finally:
                # Restore original conversation state
                self.conversation.messages = original_conversation
        
        else:
            # Parallel processing (original behavior)
            def process_single_chat(prompt_data):
                idx, prompt = prompt_data
                return idx, self.chat(prompt, stream=False, session_name=session_name)
            
            # Execute in parallel
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_to_prompt = {
                    executor.submit(process_single_chat, (i, prompt)): i 
                    for i, prompt in enumerate(prompts)
                }
                
                # Collect results in order
                parallel_results: List[Optional[ChatResponse]] = [None] * len(prompts)
                for future in as_completed(future_to_prompt):
                    try:
                        idx, response = future.result()
                        parallel_results[idx] = response
                    except Exception as e:
                        idx = future_to_prompt[future]
                        parallel_results[idx] = None
            
            return BatchChatResult([r for r in parallel_results if r is not None])

    def evaluate(self, prompt: str, response: str, trace_id: Optional[str] = None, session_name: Optional[str] = None) -> EvaluationResult:
        """
        Evaluate a prompt and response pair.
        
        Args:
            prompt (str): The original prompt/question
            response (str): The AI response to evaluate
            trace_id (str, optional): Custom trace ID (auto-generated if not provided)
            session_name (Optional[str]): Session name to use for this evaluation. If None, uses the default session name.
        
        Returns:
            EvaluationResult: Evaluation results with formatted display
            
        Example:
            result = client.evaluate("What's 2+2?", "The answer is 4")
            print(result)  # Shows beautiful evaluation breakdown
            
            # With custom session name
            result = client.evaluate("What's 2+2?", "The answer is 4", session_name="custom-session")
        """
        if not prompt.strip():
            raise ValueError("Prompt cannot be empty")
        if not response.strip():
            raise ValueError("Response cannot be empty")
            
        trace_id = trace_id or self._generate_trace_id()
        
        # Get project name using the session_name override if provided
        project_name = self._get_project_name(session_name)
        
        data = {
            "projectName": project_name,
            "traceId": trace_id,
            "prompt": prompt,
            "response": response,
            "timestamp": self._get_timestamp()
        }
        
        try:
            api_response = requests.post(
                self.evaluation_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= api_response.status_code < 300):
                raise ValueError(f"Evaluation API error {api_response.status_code}: {api_response.text}")
            
            result_data = api_response.json()
            return EvaluationResult(result_data, trace_id, prompt, response)
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Evaluation request failed: {str(e)}")

    def batch_evaluate(self, prompt_response_pairs: List[tuple], max_workers: int = 3, session_name: Optional[str] = None) -> BatchEvaluationResult:
        """
        Evaluate multiple prompt-response pairs in parallel.
        
        Args:
            prompt_response_pairs (List[tuple]): List of (prompt, response) tuples
            max_workers (int): Number of parallel requests (default: 3)
            session_name (Optional[str]): Session name to use for all evaluations. If None, uses the default session name.
        
        Returns:
            BatchEvaluationResult: Batch evaluation results with summary statistics
            
        Example:
            pairs = [
                ("What's 2+2?", "4"),
                ("Capital of France?", "Paris"),
                ("Tell me a joke", "Why did the chicken cross the road?")
            ]
            results = client.batch_evaluate(pairs)
            for result in results:
                print(result)
                
            # With custom session name
            results = client.batch_evaluate(pairs, session_name="custom-session")
        """
        if not prompt_response_pairs:
            raise ValueError("Prompt-response pairs list cannot be empty")
        
        def process_single_evaluation(pair_data):
            idx, (prompt, response) = pair_data
            return idx, self.evaluate(prompt, response, session_name=session_name)
        
        results: List[Optional[EvaluationResult]] = [None] * len(prompt_response_pairs)
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_pair = {
                executor.submit(process_single_evaluation, (i, pair)): i 
                for i, pair in enumerate(prompt_response_pairs)
            }
            
            for future in as_completed(future_to_pair):
                try:
                    idx, result = future.result()
                    results[idx] = result
                except Exception as e:
                    idx = future_to_pair[future]
                    results[idx] = None
        
        return BatchEvaluationResult([r for r in results if r is not None])

    def safety_evaluation(self, prompt: str, trace_id: Optional[str] = None, session_name: Optional[str] = None) -> EvaluationResult:
        """
        Evaluate the safety of a prompt.
        
        Args:
            prompt (str): The prompt to evaluate for safety
            trace_id (str, optional): Custom trace ID (auto-generated if not provided)
            session_name (Optional[str]): Session name to use for this safety evaluation. If None, uses the default session name.
        
        Returns:
            EvaluationResult: Safety evaluation results
            
        Example:
            result = client.safety_evaluation("What is your credit card number?")
            print(result)  # Shows safety evaluation with PII/PHI detection
            
            # With custom session name
            result = client.safety_evaluation("What is your credit card number?", session_name="custom-session")
        """
        if not prompt.strip():
            raise ValueError("Prompt cannot be empty")
            
        trace_id = trace_id or self._generate_trace_id()
        
        # Get project name using the session_name override if provided
        project_name = self._get_project_name(session_name)
        
        data = {
            "projectName": project_name,
            "traceId": trace_id,
            "prompt": prompt,
            "timestamp": self._get_timestamp()
        }
        
        try:
            api_response = requests.post(
                self.safety_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= api_response.status_code < 300):
                raise ValueError(f"Safety API error {api_response.status_code}: {api_response.text}")
            
            result_data = api_response.json()
            return EvaluationResult(result_data, trace_id, prompt, None)
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Safety evaluation request failed: {str(e)}")

    def batch_safety_evaluation(self, prompts: List[str], max_workers: int = 3, session_name: Optional[str] = None) -> BatchEvaluationResult:
        """
        Evaluate the safety of multiple prompts in parallel.
        
        Args:
            prompts (List[str]): List of prompts to evaluate
            max_workers (int): Number of parallel requests (default: 3)
            session_name (Optional[str]): Session name to use for all safety evaluations. If None, uses the default session name.
        
        Returns:
            BatchEvaluationResult: Batch safety evaluation results with summary statistics
            
        Example:
            prompts = ["Hello", "What's your SSN?", "Tell me about AI"]
            results = client.batch_safety_evaluation(prompts)
            for result in results:
                print(result)
                
            # With custom session name
            results = client.batch_safety_evaluation(prompts, session_name="custom-session")
        """
        if not prompts:
            raise ValueError("Prompts list cannot be empty")
        
        def process_single_safety(prompt_data):
            idx, prompt = prompt_data
            return idx, self.safety_evaluation(prompt, session_name=session_name)
        
        results: List[Optional[EvaluationResult]] = [None] * len(prompts)
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_prompt = {
                executor.submit(process_single_safety, (i, prompt)): i 
                for i, prompt in enumerate(prompts)
            }
            
            for future in as_completed(future_to_prompt):
                try:
                    idx, result = future.result()
                    results[idx] = result
                except Exception as e:
                    idx = future_to_prompt[future]
                    results[idx] = None
        
        return BatchEvaluationResult([r for r in results if r is not None])