import requests
import json
import logging
import uuid
import time
import os
from typing import List, Optional, Union, Callable, Any, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
from .config import (
    DEFAULT_API_URL, 
    CHATBOT_ENDPOINT,
    SET_SYSTEM_PROMPT_ENDPOINT,
    APPLY_SYSTEM_PROMPT_ENDPOINT,
    CLEAR_SYSTEM_PROMPT_ENDPOINT,
    NEW_CHAT_SESSION_ENDPOINT,
    EVALUATION_ENDPOINT, 
    SAFETY_EVALUATION_ENDPOINT, 
    TRACE_PROJECT_NAME_ENDPOINT, 
    MODEL_INFO_ENDPOINT,
    MODEL_INFO_LIST_ENDPOINT,
    ORG_TOKEN_USAGE_ENDPOINT
)
from .model import (
    EvaluationResult,
    ChatResponse,
    BatchEvaluationResult,
    BatchChatResult,
    BatchComparisonResult,
    SessionTokenUsage,
    SessionList,
    SessionMetadata,
    UsageStats
)

import threading

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
    - Evaluation of prompts and responses with automatic project name generation
    - Safety evaluation for prompts
    - System prompt management
    
    The client automatically generates project names for evaluations by calling the API
    with the session_name and appending "-Prompt" to the result.
    
    Upon initialization, the client automatically clears the context and system prompt
    to ensure a clean starting state.
    """

    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.
        
        Upon initialization, the client automatically clears the context and system prompt
        to ensure a clean starting state for the session.

        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"
            - Context and system prompt are automatically cleared upon initialization
        
        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.set_system_prompt_url = self.base_url + SET_SYSTEM_PROMPT_ENDPOINT
        self.apply_system_prompt_url = self.base_url + APPLY_SYSTEM_PROMPT_ENDPOINT
        self.clear_system_prompt_url = self.base_url + CLEAR_SYSTEM_PROMPT_ENDPOINT
        self.new_chat_session_url = self.base_url + NEW_CHAT_SESSION_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
        self.model_info_url = self.base_url + MODEL_INFO_ENDPOINT
        self.model_info_list_url = self.base_url + MODEL_INFO_LIST_ENDPOINT
        self.org_token_usage_url = self.base_url + ORG_TOKEN_USAGE_ENDPOINT
        
        # Cache for project names to avoid repeated API calls
        self._project_name_cache: Dict[str, str] = {}
        
        # Cache for model info to avoid repeated API calls
        self._model_info_cache: Dict[str, Dict[str, str]] = {}
        
        # Generate project name dynamically
        self.project_name = self._get_project_name()
        
        # Clear context and system prompt immediately after initialization
        # Clear context and system prompt in background threads (do not block init)
        threading.Thread(target=lambda: self._safe_clear_context(), daemon=True).start()
        threading.Thread(target=lambda: self._safe_clear_system_prompt(), daemon=True).start()

    def _safe_clear_context(self):
        """Safely clear context in a background thread without raising exceptions."""
        try:
            self.clear_context()
        except Exception as e:
            logger.debug(f"Failed to clear context during initialization: {e}")
    
    def _safe_clear_system_prompt(self):
        """Safely clear system prompt in a background thread without raising exceptions."""
        try:
            self.clear_system_prompt()
        except Exception as e:
            logger.debug(f"Failed to clear system prompt during initialization: {e}")

    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 _get_project_name(self, session_name: Optional[str] = None) -> str:
        """
        Get the project name by calling the trace project name API and appending '-Prompt'.
        Uses caching to avoid repeated API calls for the same session name.
        
        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
        
        # Check cache first
        if effective_session_name in self._project_name_cache:
            return self._project_name_cache[effective_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
            project_name = f"{trace_project_name}-Prompt"
            
            # Cache the result
            self._project_name_cache[effective_session_name] = project_name
            
            return project_name
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get trace project name: {str(e)}")

    def clear_project_name_cache(self):
        """
        Clear the cached project names. 
        Use this if you need to force refresh project names from the API.
        """
        self._project_name_cache.clear()
    
    def get_cached_project_names(self) -> Dict[str, str]:
        """
        Get a copy of the currently cached project names.
        
        Returns:
            Dict[str, str]: Dictionary mapping session names to project names
        """
        return self._project_name_cache.copy()

    def list_sessions(self) -> SessionList:

        try:
            response = requests.post(
                self.model_info_list_url,
                headers=self._get_headers(),
            )

            if not (200 <= response.status_code < 300):
                raise ValueError(f"Model info list API error {response.status_code}: {response.text}")

            result_data = response.json()

            result_list = SessionList()
            for session_name in result_data:

                session_data = SessionMetadata(
                    name=session_name,
                    model_type=result_data[session_name]['modelType'],
                    model_version=result_data[session_name]['modelVersion'],
                    token_usage=SessionTokenUsage(result_data[session_name]['inputTokens'],result_data[session_name]['outputTokens'])
                )

                result_list.sessions.append(session_data)


            return result_list

        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get session info: {str(e)}")


    def token_usage(self, session_name: Optional[str] = None) -> SessionTokenUsage:
        session_name = session_name or self.session_name
        data = {
            "userCreatedModelName": session_name
        }

        try:
            response = requests.post(
                self.model_info_url,
                headers=self._get_headers(),
                json=data
            )

            if not (200 <= response.status_code < 300):
                raise ValueError(f"Model info API error {response.status_code}: {response.text}")

            result_data = response.json()

            return SessionTokenUsage(result_data['inputTokens'],result_data['outputTokens'])

        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get session info: {str(e)}")

    def usage_stats(self):
        try:
            response = requests.post(
                self.org_token_usage_url,
                headers=self._get_headers(),
            )

            if not (200 <= response.status_code < 300):
                raise ValueError(f"Organization token usage API error {response.status_code}: {response.text}")

            result_data = response.json()

            return UsageStats(result_data['totalInputTokens'],result_data['totalOutputTokens'],result_data['totalTokenLimit'])

        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get organization token usage info: {str(e)}")



    def _get_model_info(self, session_name: Optional[str] = None) -> Dict[str, str]:
        """
        Get model information for a given session name.
        
        Args:
            session_name (Optional[str]): Session name to get model info for. 
                                        If None, uses the default session name.
        
        Returns:
            Dict[str, str]: Dictionary with 'modelType' and 'modelVersion' keys
            
        Raises:
            ValueError: If the API request fails
        """
        effective_session_name = session_name or self.session_name
        
        # Check cache first
        if effective_session_name in self._model_info_cache:
            return self._model_info_cache[effective_session_name]
        
        data = {
            "userCreatedModelName": effective_session_name
        }
        
        try:
            response = requests.post(
                self.model_info_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= response.status_code < 300):
                raise ValueError(f"Model info API error {response.status_code}: {response.text}")
            
            result_data = response.json()
            model_info = {
                'modelType': result_data.get('modelType', 'Unknown'),
                'modelVersion': result_data.get('modelVersion', 'Unknown')
            }
            
            # Cache the result
            self._model_info_cache[effective_session_name] = model_info
            
            return model_info
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to get model info: {str(e)}")

    def clear_model_info_cache(self):
        """
        Clear the cached model information. 
        Use this if you need to force refresh model info from the API.
        """
        self._model_info_cache.clear()
    
    def get_cached_model_info(self) -> Dict[str, Dict[str, str]]:
        """
        Get a copy of the currently cached model information.
        
        Returns:
            Dict[str, Dict[str, str]]: Dictionary mapping session names to model info dictionaries
        """
        return self._model_info_cache.copy()

    def chat(self, messages: Union[str, List[Dict[str, str]]], stream: bool = False, chat_history: bool = True, session_name: Optional[str] = None) -> ChatResponse:
        """
        Send a chat message and get response. Supports both simple strings and conversation history.
        
        The chatbot API handles history based on the withHistory parameter.
        
        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 use conversation history (default: True)
                - If True: Uses conversation history (withHistory=true in API)
                - If False: Disables conversation history (withHistory=false in API)
            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 with history enabled (default)
            response = client.chat("What is the capital of France?")
            
            # Chat without history - clears context first
            response = client.chat("What is the capital of France?", chat_history=False)
            
            # 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")
            
            prompt_for_api = messages
            prompt_for_display = 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 ''
        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
        
        # Get project name for model version extraction
        project_name = None
        try:
            project_name = self._get_project_name(effective_session_name)
        except Exception:
            # If project name retrieval fails, continue without it
            pass
        
        # Prepare request data (using 'prompt' as per the original API)
        data = {
            'prompt': prompt_for_api,
            'userCreatedModelName': effective_session_name,
            'withHistory': chat_history  # Add withHistory parameter based on chat_history
        }
        
        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
            evaluation_buffer = ""  # Buffer to accumulate evaluation JSON
            in_evaluation_block = False
            prompt_token = 0
            response_token = 0
            
            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 "id" in chunk:
                                trace_id = chunk["id"]

                            # Extract prompt / response token usage
                            if 'inputOutputTokenPair' in chunk:
                                prompt_token = chunk['inputOutputTokenPair']['inputTokens']
                                response_token = chunk['inputOutputTokenPair']['outputTokens']
                                
                            # 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
            
            # Get model info from dedicated API
            try:
                model_info = self._get_model_info(effective_session_name)
                model_type = model_info.get('modelType', 'Unknown')
                model_version = model_info.get('modelVersion', 'Unknown')
            except Exception as e:
                logger.warning(f"Failed to get model info: {e}")
                model_type = 'Unknown'
                model_version = 'Unknown'
            
            # 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_type,
                model_version=model_version,
                raw_chunks=results,
                enable_evaluations=self.enable_evaluations,
                project_name=project_name,
                session_name=effective_session_name,
                prompt_token=prompt_token,
                response_token=response_token
            )
            
            return chat_response
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Request failed: {str(e)}")

    def batch_chat(self, prompts: List[str], stream: bool = False, max_workers: int = 3, session_name: Optional[str] = None) -> BatchChatResult:
        """
        Send multiple chat messages in parallel using multithreading with the regular chat API.
        
        Note: Chat history is not supported in batch mode. All prompts are processed independently 
        with withHistory=false to ensure no cross-contamination between parallel requests.
        
        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)
            session_name (Optional[str]): Session name to use for all requests. If None, uses the default session name.
        
        Returns:
            BatchChatResult: Batch chat result object with summary statistics
            
        Example:
            # Parallel processing of independent prompts
            prompts = ["Hello!", "What's the weather?", "Tell me a joke"]
            responses = client.batch_chat(prompts)
            
            # With custom session name
            responses = client.batch_chat(prompts, session_name="custom-session")
        """
        if not prompts:
            raise ValueError("Prompts list cannot be empty")
        
        def process_single_chat(prompt_data):
            idx, prompt = prompt_data
            return idx, self.chat(prompt, stream=False, chat_history=False, session_name=session_name)
        
        # Execute in parallel using ThreadPoolExecutor
        results: List[Optional[ChatResponse]] = [None] * len(prompts)
        
        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)
            }
            
            for future in as_completed(future_to_prompt):
                try:
                    idx, response = future.result()
                    results[idx] = response
                    if stream:
                        print(f"Completed prompt {idx + 1}/{len(prompts)}")
                except Exception as e:
                    idx = future_to_prompt[future]
                    results[idx] = None
                    # Print the error to console
                    print(f"Error processing prompt {idx + 1}: {str(e)}")
        
        # Filter out None results and return BatchChatResult
        valid_results = [r for r in results if r is not None]
        return BatchChatResult(valid_results)

    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}")
            
            # Get model info for the session
            try:
                model_info = self._get_model_info(session_name)
                model_type = model_info.get('modelType', 'Unknown')
                model_version = model_info.get('modelVersion', 'Unknown')
            except Exception as e:
                logger.warning(f"Failed to get model info for evaluation: {e}")
                model_type = 'Unknown'
                model_version = 'Unknown'
            
            result_data = api_response.json()
            return EvaluationResult(result_data, trace_id, prompt, response, model_type, model_version)
            
        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
                    # Print the error to console
                    prompt, response = prompt_response_pairs[idx]
                    print(f"Error evaluating prompt {idx + 1} ('{prompt[:50]}...'): {str(e)}")
        
        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}")
            
            # Get model info for the session
            try:
                model_info = self._get_model_info(session_name)
                model_type = model_info.get('modelType', 'Unknown')
                model_version = model_info.get('modelVersion', 'Unknown')
            except Exception as e:
                logger.warning(f"Failed to get model info for safety evaluation: {e}")
                model_type = 'Unknown'
                model_version = 'Unknown'
            
            result_data = api_response.json()
            return EvaluationResult(result_data, trace_id, prompt, None, model_type, model_version)
            
        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
                    # Print the error to console
                    prompt = prompts[idx]
                    print(f"Error evaluating safety for prompt {idx + 1} ('{prompt[:50]}...'): {str(e)}")
        
        return BatchEvaluationResult([r for r in results if r is not None])

    def compare_models(self, session1_name: str, session2_name: str, prompts: List[str], 
                         stream: bool = False, max_workers: int = 3) -> BatchComparisonResult:
        """
        Compare two sessions/models by running the same set of prompts on both and comparing results.
        
        Args:
            session1_name (str): Name of the first session/model to compare
            session2_name (str): Name of the second session/model to compare  
            prompts (List[str]): List of prompts to test on both sessions
            stream (bool): Whether to show progress updates (default: False)
            max_workers (int): Number of parallel requests per session (default: 3)
        
        Returns:
            BatchComparisonResult: Comparison result object with side-by-side display
            
        Example:
            # Compare two different models/sessions
            prompts = [
                "What is artificial intelligence?",
                "Explain machine learning", 
                "Tell me a joke"
            ]
            
            comparison = client.compare_models(
                session1_name="gpt-4",
                session2_name="claude-3", 
                prompts=prompts
            )
            
            # Print side-by-side comparison
            comparison.print()
            
            # Access individual results
            session1_results = comparison.session1
            session2_results = comparison.session2
            
            # Check which performed better
            if comparison.comparison_summary['better_performing_model'] != 'tie':
                print(f"Better model: {comparison.comparison_summary['better_performing_model']}")
        """
        if not prompts:
            raise ValueError("Prompts list cannot be empty")
        if not session1_name or not session2_name:
            raise ValueError("Both session names must be provided")
        if session1_name == session2_name:
            raise ValueError("Session names must be different for comparison")
        
        # Run batch_chat for both sessions in parallel using ThreadPoolExecutor
        def run_batch_for_session(session_name: str) -> BatchChatResult:
            return self.batch_chat(
                prompts=prompts,
                stream=False,  # Disable streaming for parallel execution
                max_workers=max_workers,
                session_name=session_name
            )
        
        if stream:
            print(f"Starting comparison between '{session1_name}' and '{session2_name}'...")
            print(f"Testing {len(prompts)} prompts on both sessions...")
        
        # Execute both batch operations in parallel
        with ThreadPoolExecutor(max_workers=2) as executor:
            # Submit both tasks
            future1 = executor.submit(run_batch_for_session, session1_name)
            future2 = executor.submit(run_batch_for_session, session2_name)
            
            if stream:
                print("Running batch chats in parallel...")
            
            # Wait for both to complete
            model1_result = future1.result()
            model2_result = future2.result()
        
        if stream:
            print("Comparison completed! Generating results...")
        
        # Create and return comparison result
        return BatchComparisonResult(
            model1_result=model1_result,
            model2_result=model2_result,
            session1_name=session1_name,
            session2_name=session2_name,
            prompts=prompts
        )

    def set_system_prompt(self, prompt: str, session_name: Optional[str] = None) -> ChatResponse:
        """
        Set a system prompt using the API with streaming response.
        
        Args:
            prompt (str): The system prompt to set
            session_name (Optional[str]): Session name to use. If None, uses the default session name.
        
        Returns:
            ChatResponse: Response object with generated system prompt and evaluation results
            
        Example:
            response = client.set_system_prompt("generate a prompt to summarize a provided text into one sentence")
            print(response)
            if response.system_prompt_applied:
                print("System prompt was applied successfully.")
            else:
                print("System prompt was not applied due to evaluation issues.")
        """
            
        effective_session_name = session_name or self.session_name
        
        data = {
            "prompt": prompt,
            "userCreatedModelName": effective_session_name
        }
        
        try:
            response = requests.post(
                self.set_system_prompt_url,
                headers=self._get_headers(),
                json=data,
                stream=True
            )
            
            if not (200 <= response.status_code < 300):
                raise ValueError(f"Set system prompt API error {response.status_code}: {response.text}")
            
            # Parse streaming response similar to chat
            collected_response = ""
            evaluations = []
            trace_id = None
            system_prompt_applied = False
            model_name = effective_session_name
            model_version = "Unknown"
            raw_chunks = []
            
            for line in response.iter_lines(decode_unicode=True):
                if not line or line.strip() == "":
                    continue
                    
                # Skip event lines
                if line.startswith("event:"):
                    continue
                    
                if line.startswith("data:"):
                    data_content = line[5:].strip()  # Remove "data:" prefix
                    
                    # Skip status messages
                    if data_content in ["[START]", "[END]"]:
                        continue
                        
                    try:
                        chunk_data = json.loads(data_content)
                        raw_chunks.append(chunk_data)
                        
                        if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
                            delta = chunk_data['choices'][0].get('delta', {})
                            content = delta.get('content', '')
                            finish_reason = chunk_data['choices'][0].get('finish_reason')
                            
                            # Extract model info
                            if 'model' in chunk_data:
                                model_name = chunk_data['model']
                            
                            # Check for evaluation results in content
                            if finish_reason == "stop-and-evaluation":
                                try:
                                    eval_data = json.loads(content)
                                    evaluations = eval_data.get('evaluations', [])
                                    trace_id = eval_data.get('traceId')
                                    # System prompt is applied only if evaluations are empty or None
                                    system_prompt_applied = not evaluations or len(evaluations) == 0
                                except json.JSONDecodeError:
                                    # If content is not JSON, it's regular content
                                    collected_response += content
                            else:
                                # Regular content chunk
                                collected_response += content
                                
                    except json.JSONDecodeError:
                        # If it's not JSON, treat it as plain text
                        collected_response += data_content
            
            # Get model info
            try:
                model_info = self._get_model_info(effective_session_name)
                model_version = model_info.get('modelVersion', 'Unknown')
            except:
                model_version = "Unknown"
            
            # Create ChatResponse object
            chat_response = ChatResponse(
                response=collected_response.strip(),
                prompt=prompt,
                evaluations=evaluations,
                trace_id=trace_id,
                model=model_name,
                model_version=model_version,
                raw_chunks=raw_chunks,
                enable_evaluations=self.enable_evaluations,
                session_name=effective_session_name
            )
            
            # Set system_prompt_applied as a dynamic attribute only for set_system_prompt responses
            setattr(chat_response, 'system_prompt_applied', system_prompt_applied)
            
            return chat_response
            
        except requests.exceptions.RequestException as e:
            raise ValueError(f"Failed to set system prompt: {str(e)}")

    def apply_system_prompt(self, prompt: str, session_name: Optional[str] = None) -> bool:
        """
        Directly apply a system prompt without evaluation using the apply-sysprompt endpoint.
        
        This function bypasses the evaluation process and directly applies the system prompt.
        Use this when you want to force apply a system prompt without waiting for evaluation results.
        
        Args:
            prompt (str): The system prompt to apply directly
            session_name (Optional[str]): Session name to use. If None, uses the default session name.
        
        Returns:
            bool: True if the system prompt was applied successfully, False otherwise
            
        Example:
            success = client.apply_system_prompt("You are a helpful assistant that always responds in JSON format.")
            if success:
                print("System prompt applied successfully.")
            else:
                print("Failed to apply system prompt.")
                
            # With custom session name
            success = client.apply_system_prompt("You are a helpful assistant.", session_name="custom-session")
        """
        if not prompt.strip():
            raise ValueError("Prompt cannot be empty")
            
        effective_session_name = session_name or self.session_name
        
        # Get model info for the session
        try:
            model_info = self._get_model_info(effective_session_name)
            model_type = model_info.get('modelType', 'Unknown')
            model_version = model_info.get('modelVersion', 'Unknown')
        except Exception as e:
            logger.warning(f"Failed to get model info for apply_system_prompt: {e}")
            model_type = 'Unknown'
            model_version = 'Unknown'
        
        data = {
            "prompt": prompt,
            "modelType": model_type,
            "modelVersion": model_version,
            "userName": self.username,
            "userCreatedModelName": effective_session_name
        }
        
        try:
            response = requests.post(
                self.apply_system_prompt_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= response.status_code < 300):
                logger.warning(f"Apply system prompt API error {response.status_code}: {response.text}")
                return False
            
            # The API returns true or false as response
            result = response.text.strip().lower()
            return result == 'true'
            
        except requests.exceptions.RequestException as e:
            logger.warning(f"Failed to apply system prompt: {str(e)}")
            return False

    def clear_system_prompt(self, session_name: Optional[str] = None) -> bool:
        """
        Clear the system prompt by calling the clear API endpoint and setting it to an empty string.
        
        Args:
            session_name (Optional[str]): Session name to use. If None, uses the default session name.
        
        Returns:
            ChatResponse: Response object indicating if the system prompt was cleared successfully
            
        Example:
            response = client.clear_system_prompt()
            if getattr(response, 'system_prompt_applied', False):
                print("System prompt cleared successfully.")
            else:
                print("Failed to clear system prompt.")
        """
        effective_session_name = session_name or self.session_name
        
        # Call the clear system prompt API endpoint
        data = {
            "userCreatedModelName": effective_session_name
        }
        
        try:
            clear_response = requests.post(
                self.clear_system_prompt_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= clear_response.status_code < 300):
                # Debug print when not in 200 range
                logger.debug(f"Clear system prompt API returned status {clear_response.status_code}: {clear_response.text}")
            
        except requests.exceptions.RequestException as e:
            # Debug print for request failures
            logger.debug(f"Clear system prompt API request failed: {str(e)}")
        
        # Also call set_system_prompt with empty string to get the response
        return True

    def clear_context(self, session_name: Optional[str] = None) -> bool:
        """
        Clear the chat context/history by starting a new chat session.
        
        Args:
            session_name (Optional[str]): Session name to use. If None, uses the default session name.
        
        Returns:
            bool: True if successful, False otherwise
            
        Example:
            success = client.clear_context()
            if success:
                print("Context cleared successfully.")
            else:
                print("Failed to clear context.")
        """
        effective_session_name = session_name or self.session_name
        
        data = {
            "userCreatedModelName": effective_session_name
        }
        
        try:
            response = requests.post(
                self.new_chat_session_url,
                headers=self._get_headers(),
                json=data
            )
            
            if not (200 <= response.status_code < 300):
                # If the API endpoint doesn't exist, return True as fallback
                # This is a fallback for development/testing
                logger.warning(f"Clear context API not available (error {response.status_code}), assuming success")
                return True
            
            # The API may return a boolean value or a plain text message
            text = response.text.strip().lower()
            # Consider success if response contains 'true' or 'cleared'
            return text == 'true' or 'cleared' in text
            
        except requests.exceptions.RequestException as e:
            # If request fails, return True as fallback
            logger.warning(f"Clear context API request failed: {str(e)}, assuming success")
            return True