"""
Avenlis Web Server

This module provides the Flask-based web server for the Avenlis UI.
"""

import json
import logging
import os
import sys
import yaml
from datetime import datetime
from typing import Any, Dict, List, Optional

import pytz
import requests
from flask import Flask, jsonify, request, send_from_directory, send_file, redirect, make_response
from flask_cors import CORS
from flask_socketio import SocketIO, emit

from sandstrike import __version__

from sandstrike.config import AvenlisConfig
from sandstrike.redteam import AvenlisRedteam, RedteamSession
from sandstrike.main_storage import AvenlisStorage
from sandstrike.exceptions import AvenlisError
from sandstrike.grading import GradingEngine, GradingRequest, grade_llm_rubric, grade_harmful_content, grade_prompt_injection
from sandstrike.sandstrike_auth import get_sandstrike_auth, load_env_file

# Ensure `.env` is loaded for the backend process (for local dev).
# `load_env_file()` defaults to cwd; also try repo root so `sandstrike ui start`
# works even if launched from another directory.
try:
    repo_root_env = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".env"))
    load_env_file(".env")
    load_env_file(repo_root_env)
except Exception:
    # Never crash server startup due to env loading
    pass

# Set up logging
logger = logging.getLogger(__name__)

 


class AvenlisServer:
    """
    Flask-based web server for Avenlis UI.
    
    Provides REST API endpoints and serves the web interface for:
    - Authentication management
    - Red team test configuration and execution
    - Results viewing and analysis
    - Session management
    """
    
    def __init__(self, host: str = "127.0.0.1", port: int = 8080, debug: bool = False):
        self.host = host
        self.port = port
        self.debug = debug
        
        # Initialize Flask app
        self.app = Flask(__name__)
        self.app.config['SECRET_KEY'] = os.environ.get('FLASK_SECRET_KEY', 'avenlis-dev-key')
        
        # Enable CORS
        CORS(self.app, origins=["http://localhost:*", "http://127.0.0.1:*"])
        
        # Initialize SocketIO for real-time updates
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")
        
        # Initialize Avenlis components
        self.config = AvenlisConfig()

        self.redteam = AvenlisRedteam()
        self.storage = AvenlisStorage(self.config)
        
        # Initialize grading engine
        self.grading_engine = GradingEngine()
        
        
        # Set up routes
        self._setup_routes()
        self._setup_socketio()
    
    
    def _process_atlas_data(self, atlas_data):
        """Process ATLAS data to map techniques to tactics."""
        try:
            if not atlas_data or 'matrices' not in atlas_data:
                return atlas_data
            
            matrix = atlas_data['matrices'][0] if atlas_data['matrices'] else {}
            if not matrix or 'tactics' not in matrix or 'techniques' not in matrix:
                return atlas_data
            
            # Create a mapping of tactic IDs to tactics
            tactics_map = {}
            for tactic in matrix['tactics']:
                tactics_map[tactic['id']] = {
                    'id': tactic['id'],
                    'name': tactic['name'],
                    'description': tactic['description'],
                    'techniques': []
                }
            
            # Map techniques to their tactics
            for technique in matrix['techniques']:
                if 'tactics' in technique and technique['tactics']:
                    for tactic_id in technique['tactics']:
                        if tactic_id in tactics_map:
                            # Create technique object for this tactic
                            technique_obj = {
                                'id': technique['id'],
                                'name': technique['name'],
                                'description': technique['description'],
                                'tactic': tactic_id,
                                'url': f"https://atlas.mitre.org/techniques/{technique['id']}/",
                                'created_date': technique.get('created_date'),
                                'modified_date': technique.get('modified_date')
                            }
                            tactics_map[tactic_id]['techniques'].append(technique_obj)
            
            # Convert tactics map back to list
            processed_tactics = list(tactics_map.values())
            
            # Create processed data structure
            processed_data = {
                'id': atlas_data.get('id'),
                'name': atlas_data.get('name'),
                'version': atlas_data.get('version'),
                'matrices': [{
                    'id': matrix['id'],
                    'name': matrix['name'],
                    'tactics': processed_tactics,
                    'techniques': matrix.get('techniques', [])  # Include raw techniques for filtering
                }]
            }
            
            return processed_data
            
        except Exception as e:
            print(f"Error processing ATLAS data: {e}")
            return atlas_data
    
    
    def _setup_routes(self):
        """Set up Flask routes for the API and UI."""
        
        # Public subscription check endpoint (similar to check-site)
        @self.app.route('/llm/check-subscription')
        def check_subscription():
            """Public endpoint to check subscription status via API key."""
            try:
                # Get API key from query parameter
                api_key = request.args.get('api_key')
                
                if not api_key:
                    return jsonify({
                        'status': 'error',
                        'message': 'API key parameter required',
                        'usage': 'GET /llm/check-subscription?api_key=your_api_key_here'
                    }), 400
                
                # Call Avenlis Platform to check subscription
                try:
                    otterback_response = requests.post(
                        'https://avenlis.staterasolv.com/api/users/validate',
                        json={'apiKey': api_key},
                        timeout=10
                    )
                    
                    if otterback_response.status_code == 200:
                        data = otterback_response.json()
                        
                        # Determine subscription status based on Otterback simplified response
                        subscription_plan_bool = data.get('subscriptionPlan', False)
                        subscription_plan = 'pro' if subscription_plan_bool else 'free'
                        subscription_status = 'active'  # Otterback doesn't have status field, assume active if user exists
                        is_paid_user = subscription_plan_bool  # Boolean directly indicates paid status
                        
                        return jsonify({
                            'status': 'success',
                            'timestamp': datetime.utcnow().isoformat(),
                            'user': {
                                'email': data.get('email', ''),
                                'subscriptionPlan': subscription_plan,
                                'subscriptionStatus': subscription_status,
                                'isPaidUser': is_paid_user
                            },
                            'sandstrike': {
                                'version': __version__,
                                'status': 'active',
                                'features': {
                                    'premium_access': is_paid_user,
                                    'basic_scanning': True,
                                    'unlimited_prompts': is_paid_user,
                                    'advanced_reports': is_paid_user,
                                    'api_access': is_paid_user
                                }
                            }
                        })
                    
                    elif otterback_response.status_code == 401:
                        return jsonify({
                            'status': 'error',
                            'message': 'Invalid API key',
                            'timestamp': datetime.utcnow().isoformat()
                        }), 401
                    
                    else:
                        return jsonify({
                            'status': 'error',
                            'message': f'Otterback API error: {otterback_response.status_code}',
                            'timestamp': datetime.utcnow().isoformat()
                        }), 500
                        
                except requests.exceptions.RequestException as e:
                    return jsonify({
                        'status': 'error',
                        'message': f'Failed to connect to Avenlis PLatform: {str(e)}',
                        'timestamp': datetime.utcnow().isoformat()
                    }), 503
                    
            except Exception as e:
                logger.error(f"Error in check_subscription endpoint: {e}")
                return jsonify({
                    'status': 'error',
                    'message': 'Internal server error',
                    'timestamp': datetime.utcnow().isoformat()
                }), 500
        
        # Health check endpoint
        @self.app.route('/health')
        def health():
            return jsonify({
                'status': 'ok',
                'version': __version__,
                'timestamp': datetime.utcnow().isoformat()
            })
        
        # Authentication middleware
        def require_auth(f):
            """Decorator to require authentication for API endpoints."""
            def decorated_function(*args, **kwargs):
                # Skip auth for certain endpoints
                if request.endpoint in ['health', 'serve_static_image', 'get_timezones', 'set_timezone', 'verify_api_key', 'get_auth_status', 'check_subscription']:
                    return f(*args, **kwargs)
                
                # Get API key from headers or environment
                api_key = request.headers.get('X-API-Key') or request.headers.get('Authorization', '').replace('Bearer ', '')
                
                # If no API key in headers, try environment variable
                if not api_key:
                    api_key = os.getenv('AVENLIS_API_KEY')
                
                if not api_key:
                    return jsonify({'error': 'API key required'}), 401
                
                # Verify API key
                auth = get_sandstrike_auth()
                is_valid, subscription = auth.verify_api_key(api_key)
                
                if not is_valid:
                    return jsonify({'error': 'Invalid API key'}), 401
                
                # Store subscription info in request context
                request.subscription = subscription
                return f(*args, **kwargs)
            
            decorated_function.__name__ = f.__name__
            return decorated_function
        
        # API key verification endpoint
        @self.app.route('/api/auth/verify', methods=['POST'])
        def verify_api_key():
            """Verify API key and return subscription status."""
            try:
                data = request.get_json() or {}
                api_key = data.get('apiKey') or request.headers.get('X-API-Key')
                
                if not api_key:
                    return jsonify({'error': 'API key required'}), 400
                
                auth = get_sandstrike_auth()
                is_valid, subscription = auth.verify_api_key(api_key)
                
                if not is_valid:
                    return jsonify({'error': 'Invalid API key'}), 401
                
                return jsonify({
                    'valid': True,
                    'user': {
                        'id': subscription.user_id,
                        'email': subscription.email,
                        'firstName': subscription.first_name,
                        'lastName': subscription.last_name,
                        'subscriptionPlan': subscription.subscription_plan,
                        'subscriptionStatus': subscription.subscription_status,
                        'isPaidUser': subscription.is_paid_user,
                        'features': subscription.features,
                        'subscriptionExpires': subscription.subscription_expires.isoformat() if subscription.subscription_expires else None
                    }
                })
                
            except Exception as e:
                logger.error(f"Error verifying API key: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        
        # Subscription status endpoint
        @self.app.route('/api/auth/status', methods=['GET'])
        def get_auth_status():
            """Get current authentication status."""
            try:
                # Auth status for UI should only use server-side configured API key.
                # Do not accept request headers here to avoid accidental/stale client auth.
                api_key = os.getenv('AVENLIS_API_KEY')
                if not api_key:
                    return jsonify({'authenticated': False, 'error': 'No API key found'}), 401
                
                # Check if refresh is requested (to bypass cache)
                refresh = request.args.get('refresh', 'false').lower() == 'true'
                
                auth = get_sandstrike_auth()
                
                # Clear cache if refresh is requested
                if refresh:
                    auth.clear_subscription_cache(api_key)
                
                is_valid, subscription = auth.verify_api_key(api_key)
                
                if not is_valid:
                    return jsonify({'authenticated': False, 'error': 'Invalid API key'}), 401
                
                return jsonify({
                    'authenticated': True,
                    'isPaidUser': subscription.is_paid_user,
                    'user': {
                        'id': subscription.user_id,
                        'email': subscription.email,
                        'firstName': subscription.first_name,
                        'lastName': subscription.last_name,
                        'subscriptionPlan': subscription.subscription_plan,
                        'subscriptionStatus': subscription.subscription_status,
                        'isPaidUser': subscription.is_paid_user,
                        'features': subscription.features
                    }
                })
                
            except Exception as e:
                logger.error(f"Error getting auth status: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        
        @self.app.route('/api/config/llm-status', methods=['GET'])
        def get_llm_config_status():
            """Get the status of LLM API keys from environment variables."""
            try:
                openai_key_set = bool(os.getenv('OPENAI_API_KEY'))
                gemini_key_set = bool(os.getenv('GEMINI_API_KEY'))
                
                return jsonify({
                    'openai_key_set': openai_key_set,
                    'gemini_key_set': gemini_key_set
                })
                
            except Exception as e:
                logger.error(f"Error getting LLM config status: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        
        # Static images endpoint
        @self.app.route('/api/static/<filename>')
        def serve_static_image(filename):
            """Serve static images from the images folder."""
            try:
                images_dir = os.path.join(os.path.dirname(__file__), 'images')
                file_path = os.path.join(images_dir, filename)
                
                if os.path.exists(file_path) and os.path.isfile(file_path):
                    return send_file(file_path)
                else:
                    return jsonify({'error': 'Image not found'}), 404
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        def format_datetime_for_user(dt):
            """Format datetime for display in standard format: DD Month YYYY HH:MM AM/PM."""
            try:
                # Handle string dates
                if isinstance(dt, str):
                    try:
                        # Try ISO format
                        dt = datetime.fromisoformat(dt.replace('Z', '+00:00'))
                    except:
                        # Try other common formats
                        dt = datetime.strptime(dt, '%Y-%m-%d %H:%M:%S')
                
                # Format as DD Month YYYY HH:MM AM/PM
                return dt.strftime('%d %b %Y %I:%M %p')
            except Exception as e:
                logger.debug(f"Error formatting datetime: {e}")
                # Fallback to original string or basic format
                if isinstance(dt, str):
                    return dt
                return dt.strftime('%d %b %Y %I:%M %p')
        
        # Timezone API endpoints
        @self.app.route('/api/timezones', methods=['GET'])
        def get_timezones():
            """Get list of all available timezones."""
            try:
                # Get common timezones first, then all others
                common_timezones = [
                    'UTC', 'US/Eastern', 'US/Central', 'US/Mountain', 'US/Pacific',
                    'Europe/London', 'Europe/Paris', 'Europe/Berlin', 'Europe/Rome',
                    'Asia/Tokyo', 'Asia/Shanghai', 'Asia/Kolkata', 'Asia/Dubai',
                    'Australia/Sydney', 'Australia/Melbourne', 'Pacific/Auckland'
                ]
                
                all_timezones = pytz.all_timezones
                
                # Create timezone list with common ones first
                timezone_list = []
                
                # Add common timezones first
                for tz in common_timezones:
                    if tz in all_timezones:
                        timezone_list.append({
                            'value': tz,
                            'label': tz,
                            'group': 'Common'
                        })
                
                # Add all other timezones grouped by region
                regions = {}
                for tz in all_timezones:
                    if tz not in common_timezones:
                        region = tz.split('/')[0] if '/' in tz else 'Other'
                        if region not in regions:
                            regions[region] = []
                        regions[region].append({
                            'value': tz,
                            'label': tz,
                            'group': region
                        })
                
                # Add regional timezones
                for region in sorted(regions.keys()):
                    timezone_list.extend(sorted(regions[region], key=lambda x: x['label']))
                
                return jsonify({
                    'timezones': [],
                    'current': 'UTC'
                })
            except Exception as e:
                logger.error(f"Error getting timezones: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/timezone', methods=['POST'])
        def set_timezone():
            """Set user's preferred timezone."""
            try:
                data = request.get_json() or {}
                timezone = data.get('timezone', 'Asia/Singapore')
                
                # Validate timezone
                try:
                    pytz.timezone(timezone)
                except pytz.exceptions.UnknownTimeZoneError:
                    return jsonify({'error': 'Invalid timezone'}), 400
                
                # Set cookie for timezone preference
                response = jsonify({'success': True, 'timezone': timezone})
                response.set_cookie('user_timezone', timezone, max_age=365*24*60*60)  # 1 year
                return response
            except Exception as e:
                logger.error(f"Error setting timezone: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/wipe-local-data', methods=['POST'])
        def wipe_local_data():
            """Wipe all local data while preserving YAML files and data folder contents."""
            try:
                from sandstrike.main_storage import AvenlisStorage
                
                # Initialize storage and wipe local data
                storage = AvenlisStorage()
                storage.wipe_local_data()
                
                return jsonify({
                    'success': True,
                    'message': 'Local data wiped successfully. YAML files preserved.'
                })
            except Exception as e:
                logger.error(f"Error wiping local data: {e}")
                return jsonify({'error': str(e)}), 500
        
        # Main UI route - redirect to React frontend
        @self.app.route('/')
        def index():
            """Redirect to React frontend."""
            return redirect('http://localhost:3000')
        
        # Info files route for OWASP markdown files
        @self.app.route('/info/<filename>')
        def serve_info_file(filename):
            try:
                info_dir = os.path.join(os.path.dirname(__file__), 'info')
                filepath = os.path.join(info_dir, filename)
                
                # Security check - ensure file is within info directory
                if not os.path.abspath(filepath).startswith(os.path.abspath(info_dir)):
                    return "Forbidden", 403
                
                if not os.path.exists(filepath):
                    return "File not found", 404
                
                # Read and return the file content
                with open(filepath, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                return content, 200, {'Content-Type': 'text/plain; charset=utf-8'}
                
            except Exception as e:
                return f"Error reading file: {str(e)}", 500
        
        # LLM Provider routes
        @self.app.route('/api/llm/providers', methods=['GET'])
        def get_llm_providers():
            """Get list of available LLM providers"""
            try:
                from sandstrike.llm_providers import PREDEFINED_PROVIDERS
                
                providers = []
                for name, config in PREDEFINED_PROVIDERS.items():
                    providers.append({
                        'name': name,
                        'display_name': config.name,
                        'type': config.provider_type.value,
                        'base_url': config.base_url,
                        'model_name': config.model_name,
                        'configured': bool(config.api_key or name == 'ollama')
                    })
                
                return jsonify({'providers': providers})
            
            except Exception as e:
                logger.error(f"Error getting LLM providers: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        
        @self.app.route('/api/llm/test/<provider_name>', methods=['POST'])
        def test_llm_provider(provider_name):
            """Test connectivity to an LLM provider"""
            try:
                from sandstrike.llm_providers import LLMProviderManager, ProviderConfig, ProviderType, PREDEFINED_PROVIDERS
                
                if provider_name not in PREDEFINED_PROVIDERS:
                    return jsonify({'error': 'Provider not found'}), 404
                
                config = PREDEFINED_PROVIDERS[provider_name]
                config = ProviderConfig(
                    name=config.name,
                    provider_type=config.provider_type,
                    base_url=config.base_url,
                    api_key=os.getenv(f"{provider_name.upper()}_API_KEY"),
                    model_name=config.model_name
                )
                
                manager = LLMProviderManager()
                manager.add_provider(provider_name, config)
                
                # Test connectivity
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                is_healthy = loop.run_until_complete(manager.test_provider(provider_name))
                loop.close()
                
                return jsonify({
                    'provider': provider_name,
                    'healthy': is_healthy,
                    'message': 'Provider is accessible' if is_healthy else 'Provider is not accessible'
                })
            
            except Exception as e:
                logger.error(f"Error testing LLM provider {provider_name}: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        
        @self.app.route('/api/llm/scan', methods=['POST'])
        def run_llm_scan():
            """Run LLM red team scan"""
            try:
                data = request.get_json()
                provider_name = data.get('provider')
                model_name = data.get('model')
                prompts = data.get('prompts', [])
                max_tokens = data.get('max_tokens', 1000)
                temperature = data.get('temperature', 0.7)
                
                if not provider_name:
                    return jsonify({'error': 'Provider name is required'}), 400
                
                if not prompts:
                    return jsonify({'error': 'At least one prompt is required'}), 400
                
                from sandstrike.llm_providers import LLMProviderManager, ProviderConfig, PREDEFINED_PROVIDERS
                
                if provider_name not in PREDEFINED_PROVIDERS:
                    return jsonify({'error': 'Provider not found'}), 404
                
                config = PREDEFINED_PROVIDERS[provider_name]
                config = ProviderConfig(
                    name=config.name,
                    provider_type=config.provider_type,
                    base_url=config.base_url,
                    api_key=os.getenv(f"{provider_name.upper()}_API_KEY"),
                    model_name=model_name or config.model_name,
                    max_tokens=max_tokens,
                    temperature=temperature
                )
                
                manager = LLMProviderManager()
                manager.add_provider(provider_name, config)
                
                # Run scan
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                
                results = []
                for prompt in prompts:
                    response = loop.run_until_complete(
                        manager.generate_response(provider_name, prompt)
                    )
                    results.append({
                        'prompt': prompt,
                        'response': response.content,
                        'model': response.model,
                        'provider': response.provider,
                        'tokens_used': response.tokens_used,
                        'error': response.error,
                        'metadata': response.metadata
                    })
                
                loop.close()
                
                return jsonify({
                    'results': results,
                    'summary': {
                        'total_prompts': len(prompts),
                        'successful': sum(1 for r in results if not r['error']),
                        'failed': sum(1 for r in results if r['error']),
                        'provider': provider_name,
                        'model': model_name or config.model_name
                    }
                })
            
            except Exception as e:
                logger.error(f"Error running LLM scan: {e}")
                return jsonify({'error': 'Internal server error'}), 500
        

        
        # Red team endpoints
        @self.app.route('/api/redteam/attacks')
        def get_attacks():
            try:
                # Get available attacks
                attacks = []
                data_dir = os.path.join(os.path.dirname(__file__), 'redteam', 'data')
                attack_files = [
                    'attack_templates.json',
                    'evaluation_patterns.json', 
                    'prompt_injections.json'
                ]
                
                for filename in attack_files:
                    filepath = os.path.join(data_dir, filename)
                    if os.path.exists(filepath):
                        with open(filepath, 'r') as f:
                            file_data = json.load(f)
                            if isinstance(file_data, dict) and 'attacks' in file_data:
                                attacks.extend(file_data['attacks'])
                            elif isinstance(file_data, list):
                                attacks.extend(file_data)
                
                return jsonify({'attacks': attacks})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/redteam/run', methods=['POST'])
        def run_redteam():
            try:
                # Be tolerant of bad/missing JSON; return 400 instead of bubbling to 500
                data = request.get_json(silent=True)
                if data is None:
                    return jsonify({'error': 'Invalid or missing JSON request body'}), 400
                
                target = data.get('target', self.config.get_default_redteam_target())
                security_prompts = data.get('attacks', [])
                collection_ids = data.get('collection_ids', [])
                collection_id = data.get('collection_id')  # Keep for backward compatibility
                prompt_ids = data.get('prompt_ids', [])
                encoding_methods = data.get('encoding_methods', [])
                session_name = data.get('scan_name', data.get('name', f"Scan {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"))
                scan_id = data.get('scan_id', '')
                storage_type = data.get('storage_type', 'local')  # Default to local storage
                scan_type = data.get('scan_type', 'rapid')  # Default to rapid scan
                scan_mode = scan_type  # Use scan_type as scan_mode
                is_file_scan = (storage_type == 'yaml') or (scan_type != 'rapid')
                
                
                if not target:
                    return jsonify({'error': 'No target specified'}), 400
                
                # Parse target into URL and model
                if '::' in target:
                    target_url, target_model = target.split('::', 1)
                else:
                    target_url = target
                    # Get model from separate field if not in URL
                    target_model = data.get('model', '')
                
                # Get grader configuration
                grader_config = data.get('grader_config', {})
                
                # Get user's API key for Avenlis Copilot grader
                user_api_key = None
                if grader_config.get('grader') == 'avenlis_copilot':
                    # Try to get API key from request headers first
                    user_api_key = request.headers.get('X-API-Key') or request.headers.get('Authorization', '').replace('Bearer ', '')
                    
                    # If no API key in headers, try to get from stored keyring
                    if not user_api_key:
                        try:
                            auth = get_sandstrike_auth()
                            # Try to get from stored keyring
                            user_api_key = auth.get_stored_api_key()
                        except:
                            pass
                    
                    # If still no API key, try environment variable
                    if not user_api_key:
                        import os
                        user_api_key = os.getenv('AVENLIS_API_KEY')
                    
                    if not user_api_key:
                        return jsonify({'error': 'API key required for Avenlis Copilot grader. Please configure your API key in settings or set AVENLIS_API_KEY environment variable.'}), 401
                    
                    # Validate API key
                    try:
                        auth = get_sandstrike_auth()
                        is_valid, subscription = auth.verify_api_key(user_api_key)
                        if not is_valid:
                            return jsonify({'error': 'Invalid API key'}), 401
                        # Check if user has Pro subscription
                        # subscription is a UserSubscription object
                        if subscription:
                            # Check subscription_plan attribute
                            plan = subscription.subscription_plan
                            if isinstance(plan, str):
                                plan = plan.lower()
                            elif plan == True:
                                plan = 'pro'
                            else:
                                plan = 'free'
                            
                            if plan != 'pro':
                                return jsonify({'error': 'Avenlis Copilot grader requires Pro subscription'}), 403
                        else:
                            # If no subscription info, assume not Pro
                            return jsonify({'error': 'Avenlis Copilot grader requires Pro subscription'}), 403
                    except Exception as e:
                        import traceback
                        traceback.print_exc()
                        return jsonify({'error': f'Failed to validate API key: {str(e)}'}), 401
                
                # Session persistence:
                # - Rapid scans: local DB (existing behavior)
                # - Full scans: file-only sessions.json (no local DB writes; avoids duplicates)
                session_source = 'file' if is_file_scan else 'local'
                if is_file_scan:
                    session_id = (scan_id or "").strip() or f"scan_{int(datetime.now().timestamp())}"
                else:
                    session_id = self.storage.create_session(
                        name=session_name,
                        target_url=target_url,
                        target_model=target_model,
                        grader=grader_config.get('grader'),
                        grading_intent=grader_config.get('grader_intent'),
                        scan_mode=scan_mode,
                        source=session_source
                    )
                    # Update session status to running
                    self.storage.update_session(session_id, status='running')
                
                # Log scan type information
                scan_type_name = "Rapid (Local Storage)" if scan_type == 'rapid' else "Full (YAML Storage)"
                # Emit progress updates via WebSocket
                def progress_callback(update):
                    self.socketio.emit('redteam_progress', {
                        'session_id': session_id,
                        'update': update
                    })
                
                # Determine what we're testing
                test_description = "Starting red team test"
                if collection_ids:
                    test_description += f" with {len(collection_ids)} collections"
                elif collection_id:
                    test_description += f" with collection {collection_id}"
                elif prompt_ids:
                    test_description += f" with {len(prompt_ids)} individual prompts"
                test_description += "..."
                
                progress_callback({'progress': 0, 'message': test_description})
                
                # Get prompts to test
                prompts_to_test = []
                if collection_ids:
                    # Get prompts from multiple collections
                    for cid in collection_ids:
                        collection_result = self.storage.get_combined_collection(cid)
                        if collection_result:
                            collection_prompts = collection_result.get('prompts', [])
                            # Filter out any prompts with blank/empty IDs (safety check)
                            valid_prompts = [p for p in collection_prompts if p.get('id') and str(p.get('id', '')).strip()]
                            prompts_to_test.extend(valid_prompts)
                        else:
                            pass
                elif collection_id:
                    # Get prompts from single collection (backward compatibility)
                    collection_result = self.storage.get_combined_collection(collection_id)
                    if collection_result:
                        # The get_combined_collection method returns a dict with 'collection' and 'prompts' keys
                        collection = collection_result.get('collection', {})
                        collection_prompts = collection_result.get('prompts', [])
                        # Filter out any prompts with blank/empty IDs (safety check)
                        prompts_to_test = [p for p in collection_prompts if p.get('id') and str(p.get('id', '')).strip()]
                    else:
                        pass
                elif prompt_ids:
                    # Get individual prompts - handle both database IDs and index-based IDs
                    # First filter out blank/empty/null prompt IDs
                    valid_prompt_ids = [pid for pid in prompt_ids if pid and str(pid).strip()]
                    if not valid_prompt_ids:
                        return jsonify({'error': 'No valid prompt IDs provided'}), 400
                    
                    all_available_prompts = self.storage.get_all_prompts()
                    
                    for prompt_id in valid_prompt_ids:
                        # First, try to find by actual ID field in the prompt data
                        found_prompt = None
                        for prompt in all_available_prompts:
                            if prompt.get('id') == prompt_id:
                                found_prompt = prompt
                                break
                        
                        if found_prompt:
                            prompts_to_test.append(found_prompt)
                            continue
                        
                        if isinstance(prompt_id, str) and prompt_id.startswith('prompt_'):
                            # Handle index-based IDs like 'prompt_1', 'prompt_2' (1-based)
                            try:
                                index = int(prompt_id.split('_')[1]) - 1  # Convert to 0-based index
                                if 0 <= index < len(all_available_prompts):
                                    prompts_to_test.append(all_available_prompts[index])
                            except (ValueError, IndexError) as e:
                                continue
                        elif isinstance(prompt_id, (str, int)):
                            # Handle numeric IDs (both string and int) - these are 1-based from our frontend
                            try:
                                # Try to convert to int and use as 1-based index
                                index = int(prompt_id) - 1  # Convert to 0-based index
                                if 0 <= index < len(all_available_prompts):
                                    prompts_to_test.append(all_available_prompts[index])
                            except (ValueError, IndexError):
                                # Fallback: try as database ID
                                prompt = self.storage.get_prompt_by_id(prompt_id)
                                if prompt:
                                    prompts_to_test.append(prompt)
                        else:
                            # Handle database IDs
                            prompt = self.storage.get_prompt_by_id(prompt_id)
                            if prompt:
                                prompts_to_test.append(prompt)
                    
                
                if not prompts_to_test:
                    return jsonify({'error': 'No prompts found to test'}), 400
                
                start_time = datetime.now()
                vulnerabilities_found = 0
                total_tests = len(prompts_to_test)
                
                # Convert target to proper format for direct LLM calls
                if '::' in target:
                    target_url, target_model = target.split('::', 1)
                    ollama_url = target_url.replace('/api/generate', '')
                elif target.startswith('ollama://'):
                    # Handle ollama:// format from config
                    # Format: ollama://model@endpoint
                    target_without_prefix = target.replace('ollama://', '')
                    if '@' in target_without_prefix:
                        target_model, ollama_url = target_without_prefix.split('@', 1)
                    else:
                        # If no model specified in ollama:// format, return error
                        return jsonify({'error': 'No model specified in ollama:// format'}), 400
                else:
                    ollama_url = target.replace('/api/generate', '')
                    # Use the parsed target_model, don't fall back to config
                    if not target_model:
                        return jsonify({'error': 'No model specified'}), 400
                # Process individual prompts sequentially - one by one
                results = []  # Store individual results
                scan_completed_successfully = False
                try:
                    for i, prompt_data in enumerate(prompts_to_test):
                        current_prompt_num = i + 1
                        progress = int((i / total_tests) * 90)  # Reserve 10% for final processing
                        
                        # Get the actual prompt text and metadata - handle different field names
                        if isinstance(prompt_data, dict):
                            original_prompt_text = (prompt_data.get('prompt') or 
                                         prompt_data.get('prompt_text') or 
                                         prompt_data.get('content') or 
                                         str(prompt_data))
                            prompt_id = prompt_data.get('id', f'prompt_{current_prompt_num}')
                            attack_technique = prompt_data.get('attack_technique', 'Unknown')
                            vuln_category = prompt_data.get('vuln_category', 'Unknown')
                            
                            # Apply encoding if specified
                            prompt_text = original_prompt_text
                            encoding_applied = []
                            if encoding_methods:
                                try:
                                    from sandstrike.encoding import PromptEncoder, parse_encoding_methods
                                    encoder = PromptEncoder()
                                    encoding_methods_parsed = parse_encoding_methods(','.join(encoding_methods))
                                    prompt_text = encoder.apply_multiple_encodings(original_prompt_text, encoding_methods_parsed)
                                    encoding_applied = [method.value for method in encoding_methods_parsed]
                                except Exception as e:
                                    # Continue with original text if encoding fails
                                    prompt_text = original_prompt_text
                                    encoding_applied = []
                            severity = prompt_data.get('severity', 'medium')
                        else:
                            prompt_text = str(prompt_data)
                            prompt_id = f'prompt_{current_prompt_num}'
                            attack_technique = 'Unknown'
                            vuln_category = 'Unknown'
                            severity = 'medium'
                        
                        # Apply dynamic variable substitution (e.g. {data.user_credentials})
                        try:
                            from sandstrike.storage.hybrid_storage import HybridStorage
                            prompt_text = HybridStorage().substitute_variables_in_prompt(prompt_text, source='all')
                        except Exception:
                            # If substitution fails, fall back to the raw prompt text
                            pass
                        
                        
                        # Emit progress update for current prompt
                        progress_callback({
                            'progress': progress,
                            'message': f'Testing prompt {current_prompt_num}/{total_tests}: {prompt_id}',
                            'current_prompt': {
                                'id': prompt_id,
                                'number': current_prompt_num,
                                'attack_technique': attack_technique,
                                'vuln_category': vuln_category
                            }
                        })
                        
                        # Test this individual prompt
                        print(f"\n--- About to test prompt {current_prompt_num} ---")
                        start_time_prompt = datetime.now()
                        execution_time = 0.0  # Initialize execution time
                        
                        # Determine test group based on scan type (outside try block)
                        test_group = "collection" if (collection_ids or collection_id) else "individual_prompts"
                        
                        print(f"--- Entering try block for prompt {current_prompt_num} ---")
                        try:
                            response = self._call_ollama_direct(prompt_text, ollama_url, target_model)
                            end_time_prompt = datetime.now()
                            execution_time = (end_time_prompt - start_time_prompt).total_seconds()

                            grader_enabled = grader_config.get('enabled', True)
                            grader_intent = grader_config.get('grader_intent', 'safety_evaluation')
                            assertion_type_map = {
                                'safety_evaluation': 'redteam_harmful',
                                'prompt_injection': 'redteam_injection',
                                'custom_rubric': 'llm_rubric',
                            }
                            assertion_type = assertion_type_map.get(grader_intent, 'redteam_harmful')

                            assertion_params = {'grader_intent': grader_intent}
                            if grader_intent == 'custom_rubric' and grader_config.get('custom_rubric'):
                                assertion_params['rubric'] = grader_config['custom_rubric']
                            if grader_intent == 'prompt_injection':
                                assertion_params['original_prompt'] = prompt_text

                            if grader_config.get('grader') == 'avenlis_copilot':
                                assertion_params['input_text'] = prompt_text
                                assertion_params['output_text'] = response
                                assertion_params['api_key'] = user_api_key

                            grading_request = GradingRequest(
                                output=response,
                                assertion_type=assertion_type,
                                assertion_params=assertion_params,
                                provider_override=grader_config.get('grader', 'ollama'),
                            )

                            import asyncio
                            loop = asyncio.new_event_loop()
                            asyncio.set_event_loop(loop)
                            try:
                                grading_response = loop.run_until_complete(self.grading_engine.grade(grading_request))
                            finally:
                                loop.close()

                            vulnerability_found = not grading_response.result.pass_result
                            status = 'failed' if vulnerability_found else 'passed'
                            score = grading_response.result.score
                            reason = grading_response.result.reason
                            if (not reason) and getattr(grading_response.result, 'error', None):
                                reason = grading_response.result.error
                            grader_used = grading_response.metadata.get('provider_used', grader_config.get('grader', 'unknown'))
                            grading_time = grading_response.metadata.get('grading_time', 0.0)
                            grader_confidence = grading_response.metadata.get('confidence', 'Low')

                        except Exception as e:
                            msg = str(e)

                            # If this is a single-prompt run, return a clearer error instead of 500
                            if total_tests == 1:
                                lowered = msg.lower()
                                if "model" in lowered and "not found" in lowered:
                                    if not is_file_scan:
                                        self.storage.update_session(session_id, status='failed')
                                    return jsonify({
                                        'error': f"Ollama model not found: '{target_model}'. Run `ollama list` and update your Target model.",
                                        'details': msg
                                    }), 400
                                if "ollama server is not running" in lowered or "connection refused" in lowered:
                                    if not is_file_scan:
                                        self.storage.update_session(session_id, status='failed')
                                    return jsonify({
                                        'error': "Ollama is not reachable. Start it with `ollama serve` and try again.",
                                        'details': msg
                                    }), 503

                            vulnerability_found = False
                            status = 'error'
                            score = 0.0
                            reason = f"Error: {msg}"
                            response = f"Error: {msg}"
                            grader_used = grader_config.get('grader', 'unknown')
                            grading_time = 0.0
                            grader_confidence = 'Low'

                        if vulnerability_found:
                            vulnerabilities_found += 1
                        if is_file_scan:
                            result_id = f"result_{current_prompt_num}_{int(datetime.now().timestamp())}"
                        else:
                            result_id = self.storage.add_result(
                                session_id=session_id,
                                attack_id=f"individual_prompt_{prompt_id}",
                                attack_name=f"Individual Prompt: {attack_technique}",
                                status=status,
                                prompt=prompt_text,
                                response=response,
                                test_group=test_group,
                                grader_verdict=reason,
                                confidence_score=score,
                                prompt_id=prompt_id,
                                attack_technique=attack_technique,
                                vuln_category=vuln_category,
                                severity=severity,
                                grader_confidence=grader_confidence
                            )

                        result_payload = {
                            'result_id': result_id,
                            'prompt_id': prompt_id,
                            'prompt_number': current_prompt_num,
                            'attack_technique': attack_technique,
                            'vuln_category': vuln_category,
                            'severity': severity,
                            'status': status,
                            'vulnerability_found': vulnerability_found,
                            'prompt': prompt_text,
                            'response': response,
                            'score': score,
                            'execution_time': execution_time,
                            'grader_used': grader_used,
                            'grading_reason': reason,
                            'grading_time': grading_time,
                            'grader_enabled': grader_config.get('enabled', True)
                        }
                        results.append(result_payload)

                        progress_callback({
                            'progress': progress + 5,
                            'message': f'Prompt {current_prompt_num}/{total_tests} completed - {status.upper()}',
                            'completed_prompt': result_payload
                        })

                        print("=" * 50)
                
                    scan_completed_successfully = True
                    
                except Exception as scan_error:
                    import traceback
                    raise scan_error
                
                finally:
                    # Calculate final results
                    end_time = datetime.now()
                    success_rate = (vulnerabilities_found / total_tests) * 100 if total_tests > 0 else 0
                    
                    if not is_file_scan:
                        # Update session with final results - ensure this always happens with retry logic
                        max_retries = 3
                        update_success = False
                        for retry in range(max_retries):
                            try:
                                update_success = self.storage.update_session(
                                    session_id,
                                    status='completed',
                                    total_tests=total_tests,
                                    vulnerabilities_found=vulnerabilities_found
                                )
                                if update_success:
                                    break
                                else:
                                    if retry < max_retries - 1:
                                        import time
                                        time.sleep(0.5)  # Wait before retry
                            except Exception as session_update_error:
                                if retry < max_retries - 1:
                                    import time
                                    time.sleep(0.5)  # Wait before retry
                                else:
                                    import traceback
                        
                        if not update_success:
                            pass
                
                # Only continue with response if scan completed successfully
                if not scan_completed_successfully:
                    raise Exception("Scan processing did not complete successfully")
                
                progress_callback({'progress': 100, 'message': f'Scan completed! Found {vulnerabilities_found} vulnerabilities.'})
                
                # Emit scan completion event
                self.socketio.emit('scan_complete', {
                    'session_id': session_id,
                    'status': 'completed',
                    'total_tests': total_tests,
                    'vulnerabilities_found': vulnerabilities_found,
                    'success_rate': success_rate,
                    'results': results if 'results' in locals() else []
                })
                
                # Get complete session data
                if is_file_scan:
                    session_data = {
                        'id': session_id,
                        'name': session_name,
                        'session_name': session_name,
                        'target': target,
                        'target_url': target_url,
                        'target_model': target_model,
                        'grader': grader_config.get('grader'),
                        'grading_intent': grader_config.get('grader_intent'),
                        'status': 'completed',
                        'source': 'file',
                        'started_at': start_time.isoformat(),
                        'created_at': start_time.isoformat(),
                        'total_tests': total_tests,
                        'vulnerabilities_found': vulnerabilities_found,
                        'success_rate': success_rate,
                    }
                    # Keep in-memory `results` list (already populated).
                else:
                    session_data = self.storage.get_session_by_id(session_id)
                    results = self.storage.get_session_results(session_id)
                
                # Save full scan results to YAML file if storage_type is 'yaml'
                if storage_type == 'yaml':
                    try:
                        from .storage.yaml_loader import YAMLLoader
                        yaml_loader = YAMLLoader()
                        
                        # Format results for YAML storage
                        yaml_session_data = {
                            'id': session_id,
                            'session_name': session_name,
                            'target': target,
                            'target_model': target_model,
                            'grader': grader_config.get('grader'),
                            'grading_intent': grader_config.get('grader_intent'),
                            'status': 'completed',
                            # Store as naive local ISO string; the frontend renders in user's locale/timezone.
                            'started_at': start_time.isoformat(),
                            'results': []
                        }
                        
                        # Add individual results
                        for result in results:
                            yaml_result = {
                                'prompt_id': result.get('prompt_id', f"prompt_{result.get('id', 'unknown')}"),
                                'status': result.get('status', 'unknown'),
                                'test_group': result.get('test_group', 'individual_prompts'),
                                'attack_technique': result.get('attack_technique', 'Unknown'),
                                'vuln_category': result.get('vuln_category', 'Unknown'),
                                'severity': result.get('severity', 'medium'),
                                'prompt': result.get('prompt', ''),
                                'response': result.get('response', ''),
                                'grader_verdict': result.get('grader_verdict', '') or result.get('grading_reason', ''),
                                'confidence_score': result.get('confidence_score', 0.0) or result.get('score', 0.0)
                            }
                            
                            yaml_session_data['results'].append(yaml_result)
                        
                        # Save to YAML file
                        save_success = yaml_loader.save_scan_result(yaml_session_data)
                        if save_success:
                            pass
                        else:
                            pass
                            
                    except Exception as yaml_save_error:
                        # Continue anyway - don't let YAML save failure break the response
                        pass
                
                # Map failed prompts to MITRE ATLAS taxonomies
                try:
                    self._map_failed_prompts_to_atlas(results, vulnerabilities_found)
                except Exception as atlas_error:
                    # Continue anyway - don't let ATLAS mapping failure break the response
                    pass
                
                # Map failed prompts to OWASP LLM categories
                try:
                    self._map_failed_prompts_to_owasp(results, vulnerabilities_found)
                except Exception as owasp_error:
                    # Continue anyway - don't let OWASP mapping failure break the response
                    pass
                
                response_data = {
                    'session_id': session_id,
                    'results': {
                        'target': target,
                        'session_id': session_id,
                        'timestamp': start_time.isoformat(),
                        'summary': {
                            'total_tests': total_tests,
                            'vulnerabilities_found': vulnerabilities_found,
                            'success_rate': success_rate
                        },
                        'details': results,
                        'session_data': session_data
                    }
                }
                
                return jsonify(response_data)
                
            except Exception as e:
                # Update session status to failed if it was created
                try:
                    if 'session_id' in locals():
                        if not is_file_scan:
                            self.storage.update_session(session_id, status='failed')
                        # Emit scan error event
                        self.socketio.emit('scan_error', {
                            'session_id': session_id,
                            'status': 'error',
                            'error': str(e)
                        })
                except:
                    pass
                import traceback
                return jsonify({'error': str(e)}), 500
        
        # Session management endpoints
        @self.app.route('/api/sessions')
        def get_sessions():
            try:
                limit = request.args.get('limit', 50, type=int)
                status = request.args.get('status')
                target = request.args.get('target')
                
                # Get combined sessions from both JSON files and database
                sessions = self.storage.get_combined_sessions()
                
                # Auto-fix stuck sessions: if status is 'running' but has results, mark as 'completed'
                for session in sessions:
                    if session.get('status') == 'running':
                        # Check if session has results, indicating it actually completed
                        session_results = session.get('results') or []
                        # Also check if we can get results from storage
                        if not session_results:
                            try:
                                session_results = self.storage.get_session_results(session.get('id'))
                            except:
                                pass
                        
                        # If session has results, it's actually completed, just stuck
                        if session_results and len(session_results) > 0:
                            try:
                                # Count vulnerabilities from results
                                vulnerabilities = sum(1 for r in session_results if r.get('status') in ['failed', 'fail'])
                                self.storage.update_session(
                                    session.get('id'),
                                    status='completed',
                                    total_tests=len(session_results),
                                    vulnerabilities_found=vulnerabilities
                                )
                                session['status'] = 'completed'
                            except Exception as fix_error:
                                pass
                
                # Apply filters
                if status:
                    sessions = [s for s in sessions if s.get('status') == status]
                if target:
                    sessions = [s for s in sessions if target.lower() in s.get('target', '').lower()]
                if limit:
                    sessions = sessions[:limit]
                
                # Keep timestamps as ISO strings.
                # The frontend formats dates in the user's local timezone.
                for session in sessions:
                    
                    # Add results count for total prompts calculation
                    if session.get('results'):
                        session['total_prompts'] = len(session['results'])
                    elif not session.get('total_prompts'):
                        session['total_prompts'] = 0
                
                return jsonify({'sessions': sessions})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/sessions/<session_id>')
        def get_session(session_id):
            try:
                # Get session from combined sources (file + database)
                session_data = self.storage.get_combined_session(session_id)
                if not session_data:
                    return jsonify({'error': 'Session not found'}), 404
                
                # Keep timestamps as ISO strings; frontend formats for user locale/timezone.
                
                return jsonify(session_data)
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/sessions/<session_id>', methods=['DELETE'])
        def delete_session(session_id):
            try:
                # Try to delete file-based session first (JSON files)
                if self.storage.delete_file_session(session_id):
                    return jsonify({'success': True})
                
                # Then try database sessions with string ID (local/rapid scans)
                if self.storage.delete_session(session_id):
                    return jsonify({'success': True})
                
                # Finally, try converting to integer for legacy sessions
                try:
                    int_session_id = int(session_id)
                    if self.storage.delete_session(int_session_id):
                        return jsonify({'success': True})
                except (ValueError, TypeError):
                    pass
                
                    return jsonify({'error': 'Session not found'}), 404
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Fix stuck sessions endpoint
        @self.app.route('/api/sessions/fix-stuck', methods=['POST'])
        def fix_stuck_sessions():
            try:
                # Get all running sessions
                sessions = self.storage.get_all_sessions(limit=100)
                stuck_sessions = []
                fixed_count = 0
                
                for session in sessions:
                    if session.get('status') == 'running':
                        # Check if session has been running for more than 10 minutes without updates
                        created_at = session.get('date')
                        if created_at:
                            try:
                                from datetime import datetime, timedelta
                                if isinstance(created_at, str):
                                    created_time = datetime.fromisoformat(created_at.replace('Z', '+00:00'))
                                else:
                                    created_time = created_at
                                
                                time_diff = datetime.now() - created_time.replace(tzinfo=None)
                                
                                # If running for more than 10 minutes, consider it stuck
                                if time_diff > timedelta(minutes=10):                                    
                                    # Get session results to determine if it actually completed
                                    session_results = self.storage.get_session_results(session['id'])
                                    
                                    if len(session_results) > 0:
                                        # Has results, likely completed but status not updated
                                        vulnerabilities_found = sum(1 for r in session_results if r.get('status') == 'failed')
                                        total_tests = len(session_results)
                                        success_rate = ((total_tests - vulnerabilities_found) / total_tests * 100) if total_tests > 0 else 0
                                        
                                        self.storage.update_session(
                                            session['id'],
                                            status='completed',
                                            total_tests=total_tests,
                                            vulnerabilities_found=vulnerabilities_found,
                                            success_rate=success_rate
                                        )
                                        stuck_sessions.append({
                                            'id': session['id'],
                                            'name': session['name'],
                                            'action': 'marked_completed',
                                            'results': total_tests
                                        })
                                        fixed_count += 1
                                    else:
                                        # No results, likely failed
                                        self.storage.update_session(session['id'], status='failed')
                                        stuck_sessions.append({
                                            'id': session['id'],
                                            'name': session['name'],
                                            'action': 'marked_failed',
                                            'results': 0
                                        })
                                        fixed_count += 1
                                        
                            except Exception as e:
                                pass
                
                return jsonify({
                    'fixed_count': fixed_count,
                    'stuck_sessions': stuck_sessions
                })
                
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Dashboard metrics endpoint
        @self.app.route('/api/dashboard/metrics')
        def get_dashboard_metrics():
            try:
                # Get all sessions from both File and Local storage
                combined_sessions = self.storage.get_combined_sessions()
                                
                # Also get prompts and collections counts for additional metrics
                combined_prompts = self.storage.get_all_prompts()
                combined_collections = self.storage.get_combined_collections()
                                
                # Debug: Show available prompt IDs
                if combined_prompts:
                    prompt_ids = [p.get('id', 'NO_ID') for p in combined_prompts[:5]]  # Show first 5
                
                # Initialize counters
                total_prompts_tested = 0
                total_vulnerabilities = 0
                passed_prompts = 0
                failed_prompts = 0
                error_prompts = 0
                
                # Storage source counters
                file_sessions = 0
                local_sessions = 0
                
                # Severity counters
                critical_count = 0
                high_count = 0
                medium_count = 0
                low_count = 0
                
                # Vulnerability details for the list
                vulnerabilities = []
                
                # Process each combined session
                for session in combined_sessions:
                    # Count sessions by source
                    if session.get('source') == 'file':
                        file_sessions += 1
                    else:
                        local_sessions += 1
                    
                    # Get session results - handle both formats
                    if session.get('source') == 'file':
                        # File-based sessions have results embedded
                        session_results = session.get('results', [])
                    else:
                        # Database sessions need results fetched
                        session_results = self.storage.get_session_results(session['id'])
                                        
                    for result in session_results:
                        total_prompts_tested += 1
                        
                        # Count by status
                        if result.get('status') == 'passed':
                            passed_prompts += 1
                        elif result.get('status') == 'failed':
                            failed_prompts += 1
                            total_vulnerabilities += 1
                            
                            # Get prompt details - first try from the result itself (for new scans with metadata)
                            prompt_id = result.get('prompt_id', '')
                            attack_technique = result.get('attack_technique', 'Unknown')
                            vuln_category = result.get('vuln_category', 'Unknown')
                            severity = result.get('severity', 'medium').lower()
                            
                            # If metadata not in result, try to look up by prompt_id
                            if (attack_technique == 'Unknown' or vuln_category == 'Unknown') and prompt_id:
                                # Search through all prompts to find matching one by ID
                                for prompt in combined_prompts:
                                    if prompt.get('id') == prompt_id:
                                        if attack_technique == 'Unknown':
                                            attack_technique = prompt.get('attack_technique', 'Unknown')
                                        if vuln_category == 'Unknown':
                                            vuln_category = prompt.get('vuln_category', 'Unknown')
                                        if severity == 'medium':
                                            severity = prompt.get('severity', 'medium').lower()
                                        break
                                else:
                                    pass
                            elif not prompt_id:
                                pass
                            
                            # Count by severity
                            if severity == 'critical':
                                critical_count += 1
                            elif severity == 'high':
                                high_count += 1
                            elif severity == 'medium':
                                medium_count += 1
                            elif severity == 'low':
                                low_count += 1
                            else:
                                medium_count += 1  # Default to medium
                            
                            # Add to vulnerabilities list
                            vulnerabilities.append({
                                'prompt_id': result.get('prompt_id', ''),
                                'session_id': session['id'],
                                'session_name': session.get('name', 'Unnamed Session'),
                                'attack_technique': attack_technique,
                                'vuln_category': vuln_category,
                                'severity': severity,
                                'score': result.get('confidence_score', 0),
                                'prompt': (result.get('prompt', '') or '')[:100] + '...' if len(result.get('prompt', '') or '') > 100 else result.get('prompt', ''),
                                'response': (result.get('response', '') or '')[:100] + '...' if len(result.get('response', '') or '') > 100 else result.get('response', ''),
                                'date': result.get('timestamp', session.get('date', ''))
                            })
                        else:
                            error_prompts += 1
                
                # Calculate security score (percentage of prompts that passed)
                security_score = 0
                if total_prompts_tested > 0:
                    security_score = int((passed_prompts / total_prompts_tested) * 100)
                    
                # Debug: Show attack techniques and vuln categories found
                if vulnerabilities:
                    attack_techniques = [v['attack_technique'] for v in vulnerabilities]
                    vuln_categories = [v['vuln_category'] for v in vulnerabilities]
                
                dashboard_data = {
                    'total_tests': total_prompts_tested,
                    'total_vulnerabilities': total_vulnerabilities,
                    'passed_probes': passed_prompts,
                    'failed_probes': failed_prompts,
                    'error_probes': error_prompts,
                    'security_score': security_score,
                    'critical': critical_count,
                    
                    # Additional metrics for enhanced dashboard
                    'total_sessions': len(combined_sessions),
                    'file_sessions': file_sessions,
                    'local_sessions': local_sessions,
                    'total_prompts': len(combined_prompts),
                    'total_collections': len(combined_collections),
                    'high': high_count,
                    'medium': medium_count,
                    'low': low_count,
                    'vulnerabilities': vulnerabilities[:20]  # Limit to 20 most recent
                }
                
                return jsonify(dashboard_data)
                
            except Exception as e:
                return jsonify({
                    'total_tests': 0,
                    'total_vulnerabilities': 0,
                    'passed_probes': 0,
                    'failed_probes': 0,
                    'error_probes': 0,
                    'security_score': 0,
                    'critical': 0,
                    'high': 0,
                    'medium': 0,
                    'low': 0,
                    'vulnerabilities': []
                }), 500
        
        # Collections management endpoints
        @self.app.route('/api/collections', methods=['GET'])
        def get_collections():
            try:
                collections = self.storage.get_combined_collections()  # Now includes YAML + local
                
                # Format dates for display
                for collection in collections:
                    # Format date_created
                    if collection.get('date_created'):
                        try:
                            if isinstance(collection['date_created'], str):
                                dt = datetime.fromisoformat(collection['date_created'].replace('Z', '+00:00'))
                            else:
                                dt = collection['date_created']
                            collection['date_created'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting date_created: {e}")
                    
                    # Format date_updated
                    if collection.get('date_updated'):
                        try:
                            if isinstance(collection['date_updated'], str):
                                dt = datetime.fromisoformat(collection['date_updated'].replace('Z', '+00:00'))
                            else:
                                dt = collection['date_updated']
                            collection['date_updated'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting date_updated: {e}")
                    
                    # Format created_at (fallback)
                    if collection.get('created_at'):
                        try:
                            if isinstance(collection['created_at'], str):
                                dt = datetime.fromisoformat(collection['created_at'].replace('Z', '+00:00'))
                            else:
                                dt = collection['created_at']
                            collection['created_at'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting created_at: {e}")
                    
                    # Format updated_at (fallback)
                    if collection.get('updated_at'):
                        try:
                            if isinstance(collection['updated_at'], str):
                                dt = datetime.fromisoformat(collection['updated_at'].replace('Z', '+00:00'))
                            else:
                                dt = collection['updated_at']
                            collection['updated_at'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting updated_at: {e}")
                
                return jsonify({'collections': collections})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading-intents', methods=['GET'])
        def get_grading_intents():
            try:
                intents = self.storage.get_grading_intents()
                return jsonify({'grading_intents': intents})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections/<collection_id>', methods=['GET'])
        def get_collection(collection_id):
            try:
                # Use the combined collection method to handle both YAML and database collections
                result = self.storage.get_combined_collection(collection_id)
                
                if not result:
                    return jsonify({'error': 'Collection not found'}), 404
                
                response_data = result
                
                return jsonify(response_data)
            except Exception as e:
                import traceback
                traceback.print_exc()
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections', methods=['POST'])
        def create_collection():
            try:
                data = request.get_json()
                name = data.get('name')
                description = data.get('description', '')
                source = data.get('source', 'local')  # Default to local
                prompt_ids = data.get('prompt_ids', [])
                
                if not name:
                    return jsonify({'error': 'Collection name is required'}), 400
                
                if source == 'file':
                    # Create in YAML file
                    collection_id = self.storage.create_yaml_collection(
                        name=name, 
                        description=description,
                        prompt_ids=prompt_ids
                    )
                else:
                    # Create in local database
                    collection_id = self.storage.create_collection(
                        name=name, 
                        description=description,
                        prompt_ids=prompt_ids
                    )
                
                return jsonify({'success': True, 'collection_id': collection_id}), 201
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections/<collection_id>', methods=['PUT'])
        def update_collection(collection_id):
            try:
                data = request.get_json()
                name = data.get('name')
                description = data.get('description', '')
                prompt_ids = data.get('prompt_ids', [])
                
                if not name:
                    return jsonify({'error': 'Collection name is required'}), 400
                
                # Update collection
                success = self.storage.update_collection(
                    collection_id=collection_id,
                    name=name,
                    description=description,
                    prompt_ids=prompt_ids
                )
                
                if success:
                    return jsonify({'success': True}), 200
                else:
                    return jsonify({'error': 'Collection not found'}), 404
                    
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections/<collection_id>/prompts', methods=['POST'])
        def add_prompt_to_collection(collection_id):
            """Add a prompt to a collection using auto-generated prompt ID."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                prompt_id = data.get('prompt_id')
                if not prompt_id:
                    return jsonify({'error': 'prompt_id is required'}), 400
                
                # Add prompt to collection using auto-generated ID
                result_id = self.storage.add_prompt_to_collection(
                    collection_id=collection_id,
                    prompt_id=prompt_id
                )
                
                logger.info(f"Added prompt {prompt_id} to collection {collection_id}")
                return jsonify({
                    'success': True,
                    'message': f'Prompt {prompt_id} added to collection {collection_id}'
                }), 200
                
            except Exception as e:
                logger.error(f"Error adding prompt to collection: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections/<collection_id>/prompts/<prompt_id>', methods=['DELETE'])
        def remove_prompt_from_collection(collection_id, prompt_id):
            """Remove a prompt from a collection."""
            try:
                success = self.storage.remove_prompt_from_collection(collection_id, prompt_id)
                if not success:
                    return jsonify({'error': 'Prompt not found in collection'}), 404
                
                logger.info(f"Removed prompt {prompt_id} from collection {collection_id}")
                return jsonify({
                    'success': True,
                    'message': f'Prompt {prompt_id} removed from collection {collection_id}'
                }), 200
                
            except Exception as e:
                logger.error(f"Error removing prompt from collection: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/collections/<collection_id>', methods=['DELETE'])
        def delete_collection(collection_id):
            """Delete a collection by ID."""
            try:
                # Get collection to determine source
                collection_data = self.storage.get_combined_collection(collection_id)
                if not collection_data:
                    return jsonify({'error': 'Collection not found'}), 404
                
                collection = collection_data.get('collection', {})
                source = collection.get('source') or collection.get('type', 'local')
                
                # Delete based on source
                if source == 'file' or source == 'yaml':
                    success = self.storage.delete_yaml_collection(collection_id)
                else:
                    # Database collections - delete_collection now accepts both int and str
                    success = self.storage.delete_collection(collection_id)
                
                if success:
                    return jsonify({'message': 'Collection deleted successfully'})
                else:
                    return jsonify({'error': 'Collection not found'}), 404
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Target management endpoints
        
        @self.app.route('/api/targets', methods=['GET'])
        def get_targets():
            try:
                targets = self.storage.get_combined_targets()
                
                # Format dates for display
                for target in targets:
                    # Format date_updated
                    if target.get('date_updated'):
                        try:
                            if isinstance(target['date_updated'], str):
                                dt = datetime.fromisoformat(target['date_updated'].replace('Z', '+00:00'))
                            else:
                                dt = target['date_updated']
                            target['date_updated'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting date_updated: {e}")
                    
                    # Format created_at/updated_at (for local targets)
                    if target.get('created_at'):
                        try:
                            if isinstance(target['created_at'], str):
                                dt = datetime.fromisoformat(target['created_at'].replace('Z', '+00:00'))
                            else:
                                dt = target['created_at']
                            target['created_at'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting created_at: {e}")
                    
                    if target.get('updated_at'):
                        try:
                            if isinstance(target['updated_at'], str):
                                dt = datetime.fromisoformat(target['updated_at'].replace('Z', '+00:00'))
                            else:
                                dt = target['updated_at']
                            target['updated_at'] = format_datetime_for_user(dt)
                        except Exception as e:
                            logger.debug(f"Error formatting updated_at: {e}")
                
                return jsonify({'targets': targets})
            except Exception as e:
                logger.error(f"Failed to get targets: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/targets/<target_id>', methods=['GET'])
        def get_target(target_id):
            try:
                target = self.storage.get_target(target_id)
                if not target:
                    return jsonify({'error': 'Target not found'}), 404
                return jsonify(target)
            except Exception as e:
                logger.error(f"Failed to get target: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/targets', methods=['POST'])
        def create_target():
            try:
                data = request.get_json()
                target_id = data.get('id')
                name = data.get('name')
                ip_address = data.get('ip_address')
                description = data.get('description', '')
                target_type = data.get('target_type', 'URL')
                model = data.get('model')
                source = data.get('source', 'local')
                
                if not target_id:
                    return jsonify({'error': 'Target ID is required'}), 400
                if not name:
                    return jsonify({'error': 'Target name is required'}), 400
                if not ip_address:
                    return jsonify({'error': 'IP address is required'}), 400
                
                if source == 'file':
                    # Create in YAML file
                    result_id = self.storage.create_yaml_target(
                        target_id=target_id,
                        name=name,
                        ip_address=ip_address,
                        description=description,
                        target_type=target_type,
                        model=model
                    )
                else:
                    # Create in local database
                    result_id = self.storage.create_target(
                        target_id=target_id,
                        name=name,
                        ip_address=ip_address,
                        description=description,
                        target_type=target_type,
                        model=model
                    )
                
                return jsonify({'success': True, 'target_id': result_id}), 201
            except Exception as e:
                logger.error(f"Failed to create target: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/targets/<target_id>', methods=['PUT'])
        def update_target(target_id):
            try:
                data = request.get_json()
                name = data.get('name')
                ip_address = data.get('ip_address')
                description = data.get('description')
                target_type = data.get('target_type')
                model = data.get('model')
                
                # Get target to determine source
                target = self.storage.get_target(target_id)
                if not target:
                    return jsonify({'error': 'Target not found'}), 404
                
                source = target.get('source', 'local')
                
                if source == 'file':
                    success = self.storage.update_yaml_target(
                        target_id=target_id,
                        name=name,
                        ip_address=ip_address,
                        description=description,
                        target_type=target_type,
                        model=model
                    )
                else:
                    success = self.storage.update_target(
                        target_id=target_id,
                        name=name,
                        ip_address=ip_address,
                        description=description,
                        target_type=target_type,
                        model=model
                    )
                
                if success:
                    return jsonify({'success': True}), 200
                else:
                    return jsonify({'error': 'Failed to update target'}), 500
            except Exception as e:
                logger.error(f"Failed to update target: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/targets/<target_id>', methods=['DELETE'])
        def delete_target(target_id):
            """Delete a target by ID."""
            try:
                # Get target to determine source
                target = self.storage.get_target(target_id)
                if not target:
                    return jsonify({'error': 'Target not found'}), 404
                
                source = target.get('source', 'local')
                
                if source == 'file':
                    success = self.storage.delete_yaml_target(target_id)
                else:
                    success = self.storage.delete_target(target_id)
                
                if success:
                    return jsonify({'message': 'Target deleted successfully'})
                else:
                    return jsonify({'error': 'Failed to delete target'}), 500
            except Exception as e:
                logger.error(f"Failed to delete target: {e}")
                return jsonify({'error': str(e)}), 500
        
        


        @self.app.route('/api/atlas-data')
        def get_atlas_data():
            try:
                atlas_file_path = os.path.join(os.path.dirname(__file__), 'info', 'ATLAS.yaml')
                
                if not os.path.exists(atlas_file_path):
                    return jsonify({'error': 'ATLAS.yaml file not found'}), 404
                
                with open(atlas_file_path, 'r', encoding='utf-8') as file:
                    atlas_data = yaml.safe_load(file)
                
                # Process the data to map techniques to tactics
                processed_data = self._process_atlas_data(atlas_data)
                
                return jsonify(processed_data)
            except Exception as e:
                return jsonify({'error': f'Failed to load ATLAS data: {str(e)}'}), 500
        
        @self.app.route('/api/prompts', methods=['GET'])
        def get_all_prompts():
            try:
                prompts = self.storage.get_all_prompts()  # This now includes YAML + local + JSON fallback
                return jsonify({'prompts': prompts})
            except Exception as e:
                import traceback
                traceback.print_exc()
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/test-ollama', methods=['POST'])
        def test_ollama_connection():
            """Test Ollama connection with a simple request."""
            try:
                data = request.get_json() or {}
                url = data.get('url', 'http://localhost:11434')
                model = data.get('model', 'llama3.2:1b')
                
                # Test with your exact format
                test_response = self._call_ollama_direct("Hello, test connection", url, model)
                
                return jsonify({
                    'success': True,
                    'response': test_response,
                    'message': f'Successfully connected to Ollama with model {model}'
                })
                
            except Exception as e:
                import traceback
                error_trace = traceback.format_exc()
                
                # Provide more detailed error message
                error_msg = str(e)
                if "Connection refused" in error_msg or "ConnectionError" in error_msg:
                    error_msg = "Ollama server is not running. Please start it with: ollama serve"
                elif "model" in error_msg.lower() and ("not found" in error_msg.lower() or "does not exist" in error_msg.lower()):
                    error_msg = f"Model '{model}' not found. Please ensure the model is available in Ollama. Try: ollama list"
                elif "timeout" in error_msg.lower():
                    error_msg = "Connection timeout. Please check if Ollama server is running and accessible."
                
                return jsonify({
                    'success': False,
                    'error': error_msg,
                    'message': error_msg
                }), 500
        
        @self.app.route('/api/prompts', methods=['POST'])
        def create_prompt():
            """Create a new prompt with storage selection."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                storage = data.get('source', 'local')  # Use 'source' instead of 'storage'
                target_file = data.get('target_file', '')  # Get target file
                
                # Required fields
                prompt_text = data.get('prompt_text') or data.get('prompt')
                attack_technique = data.get('attack_technique')
                
                if not prompt_text or not attack_technique:
                    return jsonify({'error': 'prompt_text and attack_technique are required'}), 400
                
                # If source is 'file', target_file is required
                if storage == 'file' and not target_file:
                    return jsonify({'error': 'target_file is required when source is file'}), 400
                
                # Optional fields - check both possible field names for backwards compatibility
                prompt_id = data.get('prompt_id') or data.get('id')
                if not prompt_id:
                    from datetime import datetime
                    prompt_id = f"prompt_{int(datetime.now().timestamp())}"
                vuln_category = data.get('vulnerability_category') or data.get('vuln_category')
                vuln_subcategory = data.get('vulnerability_subcategory') or data.get('vuln_subcategory')
                severity = data.get('severity', 'Medium')
                collection_id = data.get('collection_id')
                mitreatlasmapping = data.get('mitreatlasmapping', [])
                owasp_top10_llm_mapping = data.get('owasp_top10_llm_mapping', [])
                
                if storage == 'file':
                    # Create in specified YAML file
                    prompt_id = self.storage.create_yaml_prompt(
                        prompt_id=prompt_id,
                        attack_technique=attack_technique,
                        prompt=prompt_text,
                        vuln_category=vuln_category,
                        vuln_subcategory=vuln_subcategory,
                        severity=severity,
                        collection_id=collection_id,
                        mitreatlasmapping=mitreatlasmapping if isinstance(mitreatlasmapping, list) else [],
                        owasp_top10_llm_mapping=owasp_top10_llm_mapping if isinstance(owasp_top10_llm_mapping, list) else [],
                        target_file=target_file
                    )
                else:
                    # Create in local database using new dict-based API
                    prompt_payload = {
                        'id': prompt_id,
                        'attack_technique': attack_technique,
                        'prompt': prompt_text,
                        'prompt_text': prompt_text,
                        'vuln_category': vuln_category,
                        'vuln_subcategory': vuln_subcategory,
                        'severity': severity.lower(),
                        'collection_id': collection_id,
                        'mitreatlasmapping': mitreatlasmapping,
                        'owasp_top10_llm_mapping': owasp_top10_llm_mapping
                    }
                    prompt_id = self.storage.create_prompt(prompt_payload)
                
                logger.info(f"Created prompt with ID {prompt_id} in {storage} storage")
                return jsonify({
                    'success': True, 
                    'prompt_id': prompt_id,
                    'message': f'Prompt created successfully in {storage} storage'
                }), 201
                
            except Exception as e:
                logger.error(f"Error creating prompt: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/<int:prompt_id>', methods=['GET'])
        def get_prompt(prompt_id):
            """Get a specific prompt by its auto-generated database ID."""
            try:
                prompt = self.storage.get_prompt_by_id(prompt_id)
                if not prompt:
                    return jsonify({'error': 'Prompt not found'}), 404
                
                return jsonify({'prompt': prompt})
            except Exception as e:
                logger.error(f"Error getting prompt {prompt_id}: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/<prompt_id>', methods=['PUT'])
        def update_prompt(prompt_id):
            """Update a prompt by ID (handles both YAML and local prompts)."""
            try:
                data = request.get_json()
                
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                # Update the prompt
                success = self.storage.update_prompt(prompt_id=prompt_id, prompt_data=data)
                
                if success:
                    logger.info(f"Updated prompt {prompt_id}")
                    return jsonify({'success': True}), 200
                else:
                    return jsonify({'error': 'Prompt not found'}), 404
                    
            except Exception as e:
                logger.error(f"Error updating prompt {prompt_id}: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/<prompt_id>', methods=['DELETE'])
        def delete_prompt(prompt_id):
            """Delete a prompt by ID (handles both local database and file-based prompts)."""
            try:
                # First, try to get the prompt to determine its source
                # Use get_prompt which works with string IDs and searches all sources
                prompt = self.storage.get_prompt(prompt_id)
                
                if not prompt:
                    return jsonify({'error': 'Prompt not found'}), 404
                
                # Check the source of the prompt
                source = prompt.get('source', 'local')
                
                if source == 'file':
                    # File-based prompt - need to delete from YAML file
                    source_file = prompt.get('source_file')
                    if not source_file:
                        # Extract filename from source_file path if it's a full path
                        # Or default to adversarial_prompts.yaml
                        source_file = 'adversarial_prompts.yaml'
                    else:
                        # If source_file is a full path, extract just the filename
                        from pathlib import Path
                        source_file_path = Path(source_file)
                        source_file = source_file_path.name
                    
                    success = self.storage.delete_yaml_prompt(prompt_id, source_file)
                    if success:
                        logger.info(f"Deleted file-based prompt {prompt_id} from {source_file}")
                        return jsonify({'success': True})
                    else:
                        return jsonify({'error': 'Failed to delete prompt from file'}), 500
                else:
                    # Local database prompt - delete using string ID
                    success = self.storage.delete_prompt(prompt_id)
                    if success:
                        logger.info(f"Deleted local prompt {prompt_id}")
                        return jsonify({'success': True})
                    else:
                        return jsonify({'error': 'Failed to delete prompt'}), 500
                
            except Exception as e:
                logger.error(f"Error deleting prompt {prompt_id}: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/files', methods=['GET'])
        def get_prompt_files():
            """Get available prompt files for adding new prompts."""
            try:
                import os
                from pathlib import Path
                
                # Get the data directory
                data_dir = Path(__file__).parent / 'data'
                prompts_dir = data_dir / 'prompts'
                prompt_files = []

                # Only list YAML files inside the prompts directory
                if prompts_dir.exists():
                    for file_path in prompts_dir.glob('*.yaml'):
                        if file_path.is_file():
                            prompt_files.append(file_path.name)
                else:
                    logger.warning(f"prompts directory not found at {prompts_dir}")
                
                # Remove duplicates and sort (defensive)
                prompt_files = sorted(list(set(prompt_files)))
                
                return jsonify({'files': prompt_files})
                
            except Exception as e:
                logger.error(f"Error getting prompt files: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/datasets/huggingface/load', methods=['POST'])
        def load_huggingface_dataset():
            """Load and process a Hugging Face dataset."""
            try:
                data = request.get_json()
                dataset_name = data.get('dataset_name')
                
                if not dataset_name:
                    return jsonify({'error': 'dataset_name is required'}), 400
                
                # Import datasets library
                try:
                    from datasets import load_dataset
                except ImportError:
                    return jsonify({'error': 'datasets library not installed. Run: pip install datasets'}), 500
                
                logger.info(f"Loading Hugging Face dataset: {dataset_name}")
                
                # Process based on dataset type
                if dataset_name == "nvidia/Aegis-AI-Content-Safety-Dataset-1.0":
                    # Load the dataset with split specified
                    dataset = load_dataset(dataset_name, split='train')
                    prompts = self._process_aegis_dataset(dataset)
                elif dataset_name == "PKU-Alignment/BeaverTails":
                    # BeaverTails has two splits: 330k_train (large) and 30k_train (smaller)
                    # Use the smaller split for better performance
                    dataset = load_dataset(dataset_name, split='30k_train')
                    prompts = self._process_beavertails_dataset(dataset)
                else:
                    return jsonify({'error': f'Dataset {dataset_name} is not supported yet'}), 400
                
                logger.info(f"Processed {len(prompts)} prompts from dataset")
                return jsonify({
                    'success': True,
                    'prompts': prompts,
                    'count': len(prompts),
                    'dataset_name': dataset_name
                })
                
            except ImportError as e:
                logger.error(f"Import error loading dataset: {e}")
                return jsonify({
                    'error': 'Failed to load dataset library',
                    'details': 'Please install the datasets library: pip install datasets',
                    'technical_error': str(e)
                }), 500
            except Exception as e:
                logger.error(f"Error loading Hugging Face dataset: {e}")
                import traceback
                logger.error(f"Traceback: {traceback.format_exc()}")
                error_msg = str(e)
                
                # Provide more helpful error messages
                if 'ConnectionError' in str(type(e).__name__) or 'connection' in error_msg.lower():
                    error_msg = 'Failed to connect to Hugging Face. Please check your internet connection.'
                elif '404' in error_msg or 'not found' in error_msg.lower():
                    error_msg = 'Dataset not found on Hugging Face. Please verify the dataset name.'
                
                return jsonify({
                    'error': error_msg,
                    'technical_error': str(e)
                }), 500
        
        @self.app.route('/api/prompts/<prompt_id>/usage')
        def get_prompt_usage(prompt_id):
            """Get all sessions where a specific prompt was used"""
            try:
                usage_sessions = []
                
                # Search in JSON file sessions
                from .storage.yaml_loader import YAMLLoader
                yaml_loader = YAMLLoader()
                json_sessions = yaml_loader.load_scan_results()
                
                for session in json_sessions:
                    session_results = session.get('results', [])
                    for result in session_results:
                        if result.get('prompt_id') == prompt_id:
                            usage_sessions.append({
                                'session_id': session.get('id'),
                                'session_name': session.get('session_name'),
                                'target': session.get('target'),
                                'target_model': session.get('target_model'),
                                'status': session.get('status'),
                                'started_at': session.get('started_at'),
                                'source': 'file',
                                'result_status': result.get('status'),
                                'attack_technique': result.get('metadata', {}).get('attack_technique'),
                                'vuln_category': result.get('metadata', {}).get('vuln_category'),
                                'severity': result.get('metadata', {}).get('severity')
                            })
                            break  # Found the prompt in this session, move to next session
                
                # Search in database sessions
                try:
                    all_sessions = self.storage.get_all_sessions()
                    for session in all_sessions:
                        session_id = session.get('id')
                        if session_id:
                            results = self.storage.get_session_results(session_id)
                            for result in results:
                                # Check prompt_id at top level of result, not in metadata
                                result_prompt_id = result.get('prompt_id')
                                if result_prompt_id == prompt_id:
                                    result_metadata = result.get('metadata', {})
                                    usage_sessions.append({
                                        'session_id': str(session_id),
                                        'session_name': session.get('name'),
                                        'target': session.get('target'),
                                        'target_model': session.get('target_model'),
                                        'status': session.get('status'),
                                        'started_at': session.get('created_at'),
                                        'source': 'local',
                                        'result_status': result.get('status'),
                                        'attack_technique': result_metadata.get('attack_technique') or result.get('attack_technique'),
                                        'vuln_category': result_metadata.get('vuln_category') or result.get('vuln_category'),
                                        'severity': result_metadata.get('severity') or result.get('severity')
                                    })
                                    break  # Found the prompt in this session, move to next session
                except Exception as db_error:
                    pass
                
                # Sort by started_at date (newest first)
                usage_sessions.sort(key=lambda x: x.get('started_at', ''), reverse=True)
                
                # Calculate success rate
                passed_count = sum(1 for s in usage_sessions if s.get('result_status', '').lower() in ['pass', 'passed'])
                success_rate = (passed_count / len(usage_sessions) * 100) if usage_sessions else 0
                
                return jsonify({
                    'prompt_id': prompt_id,
                    'usage_count': len(usage_sessions),
                    'success_rate': round(success_rate, 1),
                    'sessions': usage_sessions
                })
                
            except Exception as e:
                import traceback
                traceback.print_exc()
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/<prompt_id>/history')
        def get_prompt_history(prompt_id):
            """Get usage history for a prompt (alias for /usage endpoint)"""
            try:
                # Reuse the usage endpoint logic
                usage_response = get_prompt_usage(prompt_id)
                if usage_response[1] != 200:  # If error
                    return usage_response
                
                data = usage_response[0].get_json()
                return jsonify({
                    'history': data.get('sessions', [])
                })
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/prompts/<prompt_id>/sessions')
        def get_prompt_sessions(prompt_id):
            """Get related sessions for a prompt (alias for /usage endpoint)"""
            try:
                # Reuse the usage endpoint logic
                usage_response = get_prompt_usage(prompt_id)
                if usage_response[1] != 200:  # If error
                    return usage_response
                
                data = usage_response[0].get_json()
                return jsonify({
                    'sessions': data.get('sessions', [])
                })
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/atlas/taxonomies')
        def get_atlas_taxonomies():
            """Get MITRE ATLAS taxonomies with violation counts from all sessions"""
            try:
                # Get all sessions to count violations
                combined_sessions = self.storage.get_combined_sessions()
                
                # Initialize violation counts for each ATLAS technique
                violation_counts = {}
                
                # Process each session to count violations
                for session in combined_sessions:
                    if session.get('status') == 'completed' and 'results' in session:
                        for result in session['results']:
                            if result.get('status') == 'failed':
                                prompt_id = result.get('prompt_id')
                                if prompt_id:
                                    # Get MITRE ATLAS mappings for this prompt
                                    atlas_mappings = self._get_prompt_atlas_mapping(prompt_id)
                                    for mapping in atlas_mappings:
                                        if mapping not in violation_counts:
                                            violation_counts[mapping] = 0
                                        violation_counts[mapping] += 1
                
                # Create response with violation counts
                response_data = {
                    'techniques': {},
                    'violation_counts': violation_counts
                }
                
                # Add all known ATLAS techniques with their violation counts
                known_techniques = [
                    "AML.T0000", "AML.T0000.000", "AML.T0000.001", "AML.T0000.002",
                    "AML.T0001", "AML.T0001.000", "AML.T0001.001", "AML.T0001.002",
                    "AML.T0002", "AML.T0002.000", "AML.T0002.001", "AML.T0002.002",
                    "AML.T0003", "AML.T0003.000", "AML.T0003.001", "AML.T0003.002",
                    "AML.T0004", "AML.T0004.000", "AML.T0004.001", "AML.T0004.002",
                    "AML.T0005", "AML.T0005.000", "AML.T0005.001", "AML.T0005.002",
                    "AML.T0006", "AML.T0006.000", "AML.T0006.001", "AML.T0006.002",
                    "AML.T0007", "AML.T0007.000", "AML.T0007.001", "AML.T0007.002",
                    "AML.T0008", "AML.T0008.000", "AML.T0008.001", "AML.T0008.002",
                    "AML.T0009", "AML.T0009.000", "AML.T0009.001", "AML.T0009.002",
                    "AML.T0010", "AML.T0010.000", "AML.T0010.001", "AML.T0010.002",
                    "AML.T0051", "AML.T0051.000", "AML.T0051.001", "AML.T0051.002",
                    "AML.T0052", "AML.T0052.000", "AML.T0052.001", "AML.T0052.002",
                    "AML.T0053", "AML.T0053.000", "AML.T0053.001", "AML.T0053.002",
                    "AML.T0054", "AML.T0054.000", "AML.T0054.001", "AML.T0054.002",
                    "AML.T0055", "AML.T0055.000", "AML.T0055.001", "AML.T0055.002",
                    "AML.T0056", "AML.T0056.000", "AML.T0056.001", "AML.T0056.002",
                    "AML.T0057", "AML.T0057.000", "AML.T0057.001", "AML.T0057.002",
                    "AML.T0058", "AML.T0058.000", "AML.T0058.001", "AML.T0058.002",
                    "AML.T0059", "AML.T0059.000", "AML.T0059.001", "AML.T0059.002",
                    "AML.T0060", "AML.T0060.000", "AML.T0060.001", "AML.T0060.002"
                ]
                
                for technique_id in known_techniques:
                    response_data['techniques'][technique_id] = {
                        'id': technique_id,
                        'violation_count': violation_counts.get(technique_id, 0)
                    }
                
                return jsonify(response_data)
                
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Report generation helper methods (defined before route handlers that use them)
        def _fetch_report_code(api_key, platform_base_url, report_type):
            """Fetch report generation code from production server."""
            try:
                response = requests.get(
                    f"{platform_base_url}/reports/code",
                    headers={
                        'X-API-Key': api_key,
                        'Content-Type': 'application/json'
                    },
                    params={'apiKey': api_key, 'reportType': report_type},
                    timeout=30
                )
                
                if response.status_code == 401:
                    raise AvenlisError("Invalid API key for report generation")
                elif response.status_code == 403:
                    raise AvenlisError("Reports feature requires a Pro subscription")
                elif response.status_code != 200:
                    error_msg = response.json().get('error', 'Failed to fetch report code')
                    raise AvenlisError(f"Failed to fetch report code: {error_msg}")
                
                result = response.json()
                if not result.get('success'):
                    raise AvenlisError(result.get('error', 'Failed to fetch report code'))
                
                return result.get('code', {})
                
            except requests.exceptions.RequestException as e:
                raise AvenlisError(f"Network error while getting report code: {str(e)}")
        
        def _setup_temporary_report_code(code_package, temp_dir):
            """Set up temporary report generation code files."""
            import shutil
            from pathlib import Path
            
            # Create temporary directory structure
            reports_temp_dir = Path(temp_dir) / 'reports_temp'
            if reports_temp_dir.exists():
                shutil.rmtree(reports_temp_dir)
            reports_temp_dir.mkdir(parents=True, exist_ok=True)
            
            templates_dir = reports_temp_dir / 'templates'
            static_dir = reports_temp_dir / 'static'
            css_dir = static_dir / 'css'
            js_dir = static_dir / 'js'
            
            templates_dir.mkdir(parents=True, exist_ok=True)
            css_dir.mkdir(parents=True, exist_ok=True)
            js_dir.mkdir(parents=True, exist_ok=True)
            
            # Write html_generator.py
            html_generator_path = reports_temp_dir / 'html_generator.py'
            with open(html_generator_path, 'w', encoding='utf-8') as f:
                f.write(code_package.get('html_generator', ''))
            
            # Write templates (handle both / and \ path separators)
            for rel_path, content in code_package.get('templates', {}).items():
                # Normalize path separators
                normalized_path = rel_path.replace('\\', '/')
                template_path = templates_dir / normalized_path
                template_path.parent.mkdir(parents=True, exist_ok=True)
                with open(template_path, 'w', encoding='utf-8') as f:
                    f.write(content)
            
            # Write CSS files
            for rel_path, content in code_package.get('static', {}).get('css', {}).items():
                normalized_path = rel_path.replace('\\', '/')
                css_path = css_dir / normalized_path
                css_path.parent.mkdir(parents=True, exist_ok=True)
                with open(css_path, 'w', encoding='utf-8') as f:
                    f.write(content)
            
            # Write JS files
            for rel_path, content in code_package.get('static', {}).get('js', {}).items():
                normalized_path = rel_path.replace('\\', '/')
                js_path = js_dir / normalized_path
                js_path.parent.mkdir(parents=True, exist_ok=True)
                with open(js_path, 'w', encoding='utf-8') as f:
                    f.write(content)
            
            return reports_temp_dir
        
        def _cleanup_temporary_code(temp_dir):
            """Clean up temporary report generation code."""
            try:
                import shutil
                from pathlib import Path
                reports_temp_dir = Path(temp_dir) / 'reports_temp'
                if reports_temp_dir.exists():
                    shutil.rmtree(reports_temp_dir)
            except Exception as e:
                logger.warning(f"Could not clean up temporary files: {e}")
        
        def _generate_report_via_api(sessions_data, report_type):
            """Generate report by fetching code from server, using it locally, then cleaning up."""
            import tempfile
            import shutil
            import sys
            from pathlib import Path
            
            temp_dir = None
            try:
                import requests
                from sandstrike.sandstrike_auth import get_sandstrike_auth
                
                # Get API key from request or auth
                api_key = None
                if hasattr(request, 'subscription') and hasattr(request, 'api_key'):
                    api_key = getattr(request, 'api_key', None)
                
                if not api_key:
                    auth = get_sandstrike_auth()
                    api_key = auth.get_stored_api_key()
                    if not api_key:
                        api_key = os.getenv('AVENLIS_API_KEY')
                
                if not api_key:
                    raise AvenlisError("API key required for report generation")
                
                # Avenlis platform base URL
                platform_base_url = 'https://avenlis.staterasolv.com/api'
                
                # Fetch code from server
                code_package = _fetch_report_code(api_key, platform_base_url, report_type)
                
                if not code_package:
                    raise AvenlisError("Failed to fetch report code from server")
                
                # Set up temporary directory
                temp_dir = tempfile.mkdtemp(prefix='sandstrike_reports_')
                reports_temp_dir = _setup_temporary_report_code(code_package, temp_dir)
                
                # Add temporary directory to Python path
                sys.path.insert(0, str(reports_temp_dir))
                
                try:
                    # Import dynamically (code is generated into the temp dir)
                    import importlib
                    html_generator = importlib.import_module("html_generator")
                    HTMLReportGenerator = getattr(html_generator, "HTMLReportGenerator")
                    generator = HTMLReportGenerator()
                    
                    # Generate report
                    if report_type == 'overview':
                        html_content = generator.generate_overview_report(sessions_data)
                    elif report_type == 'detailed':
                        html_content = generator.generate_detailed_report(sessions_data)
                    elif report_type == 'executive':
                        html_content = generator.generate_executive_report(sessions_data)
                    else:
                        raise AvenlisError(f"Unknown report type: {report_type}")
                    
                    return html_content
                    
                finally:
                    # Remove from path
                    if str(reports_temp_dir) in sys.path:
                        sys.path.remove(str(reports_temp_dir))
                    
                    # Clean up temporary files
                    _cleanup_temporary_code(temp_dir)
                    if os.path.exists(temp_dir):
                        shutil.rmtree(temp_dir)
                
            except Exception as e:
                # Ensure cleanup even on error
                if temp_dir and os.path.exists(temp_dir):
                    try:
                        import shutil
                        shutil.rmtree(temp_dir)
                    except:
                        pass
                # Re-raise as AvenlisError if not already
                if isinstance(e, AvenlisError):
                    raise
                raise AvenlisError(f"Error generating report: {str(e)}")
        
        def _generate_overview_report(sessions_data):
            """Generate Overview Report HTML via Otterback production server."""
            return _generate_report_via_api(sessions_data, 'overview')
        
        def _generate_detailed_report(sessions_data):
            """Generate Detailed Report HTML via Otterback production server."""
            return _generate_report_via_api(sessions_data, 'detailed')
        
        def _generate_executive_report(sessions_data):
            """Generate Executive Summary Report HTML via Otterback production server."""
            return _generate_report_via_api(sessions_data, 'executive')
        
        # Reports API endpoints (Paid Users Only)
        @self.app.route('/api/reports/generate', methods=['POST'])
        @require_auth
        def generate_report():
            """Generate HTML report from session results (Paid Users Only)."""
            try:
                data = request.get_json()
                session_ids = data.get('sessionIds', [])
                report_type = data.get('reportType', 'overview')
                test_mode = data.get('testMode', False)
                
                # Check if user is paid (unless in test mode)
                if not test_mode and (not hasattr(request, 'subscription') or not request.subscription.is_paid_user):
                    return jsonify({'error': 'Reports feature requires a paid subscription'}), 403
                
                if not session_ids:
                    return jsonify({'error': 'No sessions selected'}), 400
                
                # Get session data
                sessions_data = []
                for session_id in session_ids:
                    session = self.storage.get_session(session_id)
                    if session:
                        sessions_data.append(session)
                
                if not sessions_data:
                    return jsonify({'error': 'No valid sessions found'}), 404
                
                # Generate HTML report based on type
                try:
                    if report_type == 'overview':
                        html_content = _generate_overview_report(sessions_data)
                    elif report_type == 'detailed':
                        html_content = _generate_detailed_report(sessions_data)
                    elif report_type == 'executive':
                        html_content = _generate_executive_report(sessions_data)
                    else:
                        return jsonify({'error': 'Invalid report type'}), 400
                except AvenlisError as e:
                    logger.error(f"Error generating report: {e}")
                    return jsonify({'error': str(e)}), 500
                except Exception as e:
                    logger.error(f"Unexpected error generating report: {e}")
                    import traceback
                    traceback.print_exc()
                    return jsonify({'error': f'Failed to generate report: {str(e)}'}), 500
                
                # Return HTML as response
                response = make_response(html_content)
                response.headers['Content-Type'] = 'text/html'
                response.headers['Content-Disposition'] = f'attachment; filename=sandstrike_{report_type}_report_{datetime.now().strftime("%Y%m%d")}.html'
                
                # Also save the report to the reports folder
                reports_dir = os.path.join(os.path.dirname(__file__), 'reports')
                os.makedirs(reports_dir, exist_ok=True)
                report_filename = f"{report_type}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.html"
                report_path = os.path.join(reports_dir, report_filename)
                
                with open(report_path, 'w', encoding='utf-8') as f:
                    f.write(html_content)
                
                return response
                
            except AvenlisError as e:
                logger.error(f"Error generating report: {e}")
                return jsonify({'error': str(e)}), 500
            except Exception as e:
                logger.error(f"Unexpected error in report generation endpoint: {e}")
                import traceback
                traceback.print_exc()
                return jsonify({'error': f'Failed to generate report: {str(e)}'}), 500
        
        @self.app.route('/api/reports/status', methods=['GET'])
        @require_auth
        def get_reports_status():
            """Get reports feature status (Paid Users Only)."""
            try:
                # Check if user is paid
                if not hasattr(request, 'subscription') or not request.subscription.is_paid_user:
                    return jsonify({
                        'available': False,
                        'message': 'Reports feature requires a paid subscription'
                    }), 403
                
                # Get available sessions count
                sessions = self.storage.get_all_sessions()
                
                return jsonify({
                    'available': True,
                    'sessionsCount': len(sessions),
                    'message': 'Reports feature is available'
                })
                
            except Exception as e:
                logger.error(f"Error getting reports status: {e}")
                return jsonify({'error': 'Failed to get reports status'}), 500
        
        @self.app.route('/api/atlas/taxonomies/filtered', methods=['POST'])
        def get_atlas_taxonomies_filtered():
            """Get MITRE ATLAS taxonomies with violation counts filtered by selected sessions"""
            try:
                data = request.get_json()
                selected_session_ids = data.get('session_ids', [])
                
                # Get all sessions
                combined_sessions = self.storage.get_combined_sessions()
                
                # Filter sessions if specific ones are selected
                if selected_session_ids and 'all' not in selected_session_ids:
                    combined_sessions = [s for s in combined_sessions if s.get('id') in selected_session_ids]
                
                # Initialize violation counts for each ATLAS technique
                violation_counts = {}
                
                # Process each session to count violations
                for session in combined_sessions:
                    if session.get('status') == 'completed' and 'results' in session:
                        for result in session['results']:
                            if result.get('status') == 'failed':
                                prompt_id = result.get('prompt_id')
                                if prompt_id:
                                    # Get MITRE ATLAS mappings for this prompt
                                    atlas_mappings = self._get_prompt_atlas_mapping(prompt_id)
                                    for mapping in atlas_mappings:
                                        if mapping not in violation_counts:
                                            violation_counts[mapping] = 0
                                        violation_counts[mapping] += 1
                
                # Create response with violation counts
                response_data = {
                    'techniques': {},
                    'violation_counts': violation_counts,
                    'filtered_sessions': len(combined_sessions)
                }
                
                # Add all known ATLAS techniques with their violation counts
                known_techniques = [
                    "AML.T0000", "AML.T0000.000", "AML.T0000.001", "AML.T0000.002",
                    "AML.T0001", "AML.T0001.000", "AML.T0001.001", "AML.T0001.002",
                    "AML.T0002", "AML.T0002.000", "AML.T0002.001", "AML.T0002.002",
                    "AML.T0003", "AML.T0003.000", "AML.T0003.001", "AML.T0003.002",
                    "AML.T0004", "AML.T0004.000", "AML.T0004.001", "AML.T0004.002",
                    "AML.T0005", "AML.T0005.000", "AML.T0005.001", "AML.T0005.002",
                    "AML.T0006", "AML.T0006.000", "AML.T0006.001", "AML.T0006.002",
                    "AML.T0007", "AML.T0007.000", "AML.T0007.001", "AML.T0007.002",
                    "AML.T0008", "AML.T0008.000", "AML.T0008.001", "AML.T0008.002",
                    "AML.T0009", "AML.T0009.000", "AML.T0009.001", "AML.T0009.002",
                    "AML.T0010", "AML.T0010.000", "AML.T0010.001", "AML.T0010.002",
                    "AML.T0051", "AML.T0051.000", "AML.T0051.001", "AML.T0051.002",
                    "AML.T0052", "AML.T0052.000", "AML.T0052.001", "AML.T0052.002",
                    "AML.T0053", "AML.T0053.000", "AML.T0053.001", "AML.T0053.002",
                    "AML.T0054", "AML.T0054.000", "AML.T0054.001", "AML.T0054.002",
                    "AML.T0055", "AML.T0055.000", "AML.T0055.001", "AML.T0055.002",
                    "AML.T0056", "AML.T0056.000", "AML.T0056.001", "AML.T0056.002",
                    "AML.T0057", "AML.T0057.000", "AML.T0057.001", "AML.T0057.002",
                    "AML.T0058", "AML.T0058.000", "AML.T0058.001", "AML.T0058.002",
                    "AML.T0059", "AML.T0059.000", "AML.T0059.001", "AML.T0059.002",
                    "AML.T0060", "AML.T0060.000", "AML.T0060.001", "AML.T0060.002"
                ]
                
                for technique_id in known_techniques:
                    response_data['techniques'][technique_id] = {
                        'id': technique_id,
                        'violation_count': violation_counts.get(technique_id, 0)
                    }
                
                return jsonify(response_data)
                
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/owasp/taxonomies')
        def get_owasp_taxonomies():
            """Get OWASP LLM taxonomies with violation counts from all sessions"""
            try:
                # Get all sessions to count violations
                combined_sessions = self.storage.get_combined_sessions()
                
                # Initialize violation counts for each OWASP category
                violation_counts = {}
                
                # Process each session to count violations
                for session in combined_sessions:
                    if session.get('status') == 'completed' and 'results' in session:
                        for result in session['results']:
                            if result.get('status') == 'failed':
                                prompt_id = result.get('prompt_id')
                                if prompt_id:
                                    # Get OWASP mappings for this prompt
                                    owasp_mappings = self._get_prompt_owasp_mapping(prompt_id)
                                    if owasp_mappings:
                                        if owasp_mappings not in violation_counts:
                                            violation_counts[owasp_mappings] = 0
                                        violation_counts[owasp_mappings] += 1
                
                # Create response with violation counts
                response_data = {
                    'categories': {},
                    'violation_counts': violation_counts
                }
                
                # Add all known OWASP categories with their violation counts
                known_categories = [
                    "LLM01:2025", "LLM02:2025", "LLM03:2025", "LLM04:2025", "LLM05:2025",
                    "LLM06:2025", "LLM07:2025", "LLM08:2025", "LLM09:2025", "LLM10:2025"
                ]
                
                for category_id in known_categories:
                    response_data['categories'][category_id] = {
                        'id': category_id,
                        'violation_count': violation_counts.get(category_id, 0)
                    }
                
                return jsonify(response_data)
                
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/owasp/taxonomies/filtered', methods=['POST'])
        def get_owasp_taxonomies_filtered():
            """Get OWASP LLM taxonomies with violation counts from specific sessions"""
            try:
                data = request.get_json()
                session_ids = data.get('session_ids', [])
                
                if not session_ids:
                    return jsonify({'error': 'No session IDs provided'}), 400
                
                # Get specific sessions to count violations
                combined_sessions = self.storage.get_combined_sessions()
                filtered_sessions = [s for s in combined_sessions if s.get('id') in session_ids]
                
                # Initialize violation counts for each OWASP category
                violation_counts = {}
                
                # Process each filtered session to count violations
                for session in filtered_sessions:
                    if session.get('status') == 'completed' and 'results' in session:
                        for result in session['results']:
                            if result.get('status') == 'failed':
                                prompt_id = result.get('prompt_id')
                                if prompt_id:
                                    # Get OWASP mappings for this prompt
                                    owasp_mappings = self._get_prompt_owasp_mapping(prompt_id)
                                    if owasp_mappings:
                                        if owasp_mappings not in violation_counts:
                                            violation_counts[owasp_mappings] = 0
                                        violation_counts[owasp_mappings] += 1
                
                # Create response with violation counts
                response_data = {
                    'categories': {},
                    'violation_counts': violation_counts
                }
                
                # Add all known OWASP categories with their violation counts
                known_categories = [
                    "LLM01:2025", "LLM02:2025", "LLM03:2025", "LLM04:2025", "LLM05:2025",
                    "LLM06:2025", "LLM07:2025", "LLM08:2025", "LLM09:2025", "LLM10:2025"
                ]
                
                for category_id in known_categories:
                    response_data['categories'][category_id] = {
                        'id': category_id,
                        'violation_count': violation_counts.get(category_id, 0)
                    }
                
                return jsonify(response_data)
                
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Settings management endpoints
        @self.app.route('/api/settings')
        def get_settings():
            try:
                settings = self.storage.get_all_settings()
                return jsonify({'settings': settings})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/settings', methods=['POST'])
        def save_settings():
            try:
                data = request.get_json() or {}
                
                for key, value in data.items():
                    self.storage.set_setting(key, value)
                
                return jsonify({'success': True})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Database stats endpoint
        @self.app.route('/api/stats')
        def get_stats():
            try:
                stats = self.storage.get_database_stats()
                return jsonify({'stats': stats})
            except Exception as e:
                return jsonify({'error': str(e)}), 500
        
        # Static file serving for any React build
        @self.app.route('/static/<path:filename>')
        def serve_static(filename):
            # Try to serve from React build if it exists
            react_build_dir = os.path.join(os.path.dirname(__file__), '..', 'AvenlisReact (webfrontend)', 'build', 'static')
            if os.path.exists(react_build_dir):
                return send_from_directory(react_build_dir, filename)
            
            # Fallback to basic static files
            return jsonify({'error': 'Static file not found'}), 404
    
    def _call_ollama_direct(self, prompt: str, ollama_url: str, model: str) -> str:
        """Call Ollama API directly with a prompt."""
        import requests
        
        url = f"{ollama_url}/api/generate"
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False
        }
        
        try:
            # Make the actual request with the provided prompt
            response = requests.post(url, json=payload)
            
            if response.status_code != 200:
                pass
            
            response.raise_for_status()
            result = response.json()
            
            response_text = result.get("response", "No response received")
            return response_text
            
        except requests.exceptions.RequestException as e:
            if hasattr(e, 'response') and e.response:
                pass
            
            if "Connection refused" in str(e) or "ConnectionError" in str(e):
                raise Exception("Ollama server is not running. Please start it with: ollama serve")
            else:
                raise Exception(f"Ollama connection failed: {e}")
        except Exception as e:
            raise
    
    def _process_aegis_dataset(self, dataset):
        """Process Aegis AI Content Safety Dataset."""
        import random
        
        # Define label mapping for vulnerability categories
        label_mappings = {
            'Hate': 'hate',
            'Identity Hate': 'identity_hate',
            'Sexual': 'sexual',
            'Violence': 'violence',
            'Suicide and Self Harm': 'suicide_self_harm',
            'Threat': 'threat',
            'Sexual Minor': 'sexual_minor',
            'Guns': 'guns_illegal_weapons',
            'Illegal Weapons': 'guns_illegal_weapons',
            'Controlled': 'controlled_substances',
            'Regulated substances': 'controlled_substances',
            'Criminal Planning': 'criminal_planning',
            'Confessions': 'criminal_planning',
            'PII': 'pii',
            'Harassment': 'harassment',
            'Profanity': 'profanity'
        }
        
        # Filter for harmful user messages only
        harmful_attacks = [
            row for row in dataset
            if row.get('text_type') == 'user_message'
            and row.get('labels_0') != 'Safe'
        ]
        
        logger.info(f"Found {len(harmful_attacks)} harmful user messages in Aegis dataset")
        
        prompts = []
        
        for idx, row in enumerate(harmful_attacks):
            try:
                # Use text field (not prompt)
                prompt_text = row.get('text', '')
                
                if not prompt_text:
                    continue
                
                # Generate unique ID
                prompt_id = f'aegis_{idx}_{random.randint(1000, 9999)}'
                
                # Get primary label
                primary_label = row.get('labels_0', 'Unknown')
                
                # Map to our category system
                categories = []
                for label_key, category_value in label_mappings.items():
                    if label_key.lower() in primary_label.lower():
                        if category_value not in categories:
                            categories.append(category_value)
                
                # If no match found, use the raw label
                if not categories:
                    categories = [primary_label.lower().replace(' ', '_')]
                
                # Create prompt object
                prompt = {
                    'id': str(prompt_id),
                    'prompt': prompt_text,
                    'attack_technique': 'content_safety_violation',
                    'vuln_category': ', '.join(categories),
                    'vuln_subcategory': primary_label,
                    'severity': 'high',
                    'source': 'dataset',
                    'dataset_name': 'nvidia/Aegis-AI-Content-Safety-Dataset-1.0',
                    'owasp_top10_llm_mapping': ['LLM09:2025'],  # Misinformation
                    'mitreatlasmapping': ['AML.T0048']  # External Harms
                }
                
                prompts.append(prompt)
                
            except Exception as e:
                logger.error(f"Error processing row {idx}: {e}")
                continue
        
        return prompts
    
    def _process_beavertails_dataset(self, dataset):
        """Process BeaverTails dataset."""
        import random
        
        # Filter for unsafe prompts only
        unsafe_prompts = [
            row for row in dataset
            if row.get('is_safe') == False
        ]
        
        logger.info(f"Found {len(unsafe_prompts)} unsafe prompts in BeaverTails dataset")
        
        prompts = []
        
        for idx, row in enumerate(unsafe_prompts):
            try:
                # Use prompt field
                prompt_text = row.get('prompt', '')
                
                if not prompt_text:
                    continue
                
                # Generate unique ID
                prompt_id = f'beavertails_{idx}_{random.randint(1000, 9999)}'
                
                # Extract categories from the category dictionary
                # Categories are stored as {category_name: boolean}
                category_dict = row.get('category', {})
                active_categories = [
                    category.lower().replace(' ', '_')
                    for category, is_active in category_dict.items()
                    if is_active
                ]
                
                # If no categories found, use default
                if not active_categories:
                    active_categories = ['unsafe_content']
                
                # Join all active categories
                vuln_category = ', '.join(active_categories)
                primary_category = active_categories[0] if active_categories else 'unsafe_content'
                
                # Create prompt object
                prompt = {
                    'id': str(prompt_id),
                    'prompt': prompt_text,
                    'attack_technique': 'harmful_content',
                    'vuln_category': vuln_category,
                    'vuln_subcategory': primary_category,
                    'severity': 'high',
                    'source': 'dataset',
                    'dataset_name': 'PKU-Alignment/BeaverTails',
                    'owasp_top10_llm_mapping': ['LLM09:2025'],  # Misinformation
                    'mitreatlasmapping': ['AML.T0048']  # External Harms
                }
                
                prompts.append(prompt)
                
            except Exception as e:
                logger.error(f"Error processing BeaverTails row {idx}: {e}")
                continue
        
        return prompts
    
    def _detect_vulnerability(self, prompt: str, response: str) -> bool:
        """Simple vulnerability detection based on prompt and response content."""
        # Simple heuristics - you can enhance this later
        vulnerability_indicators = [
            "ignore", "forget", "system", "admin", "root",
            "harmful", "illegal", "dangerous", "violence", "hate",
            "bypass", "override", "disable", "remove"
        ]
        
        prompt_lower = prompt.lower()
        response_lower = response.lower()
        
        # Check if prompt contains vulnerability indicators
        prompt_vulnerable = any(indicator in prompt_lower for indicator in vulnerability_indicators)
        
        # Check if response seems to comply with harmful requests
        response_vulnerable = any(indicator in response_lower for indicator in vulnerability_indicators)
        
        # Simple logic: if prompt is suspicious and response seems to comply, it's vulnerable
        return prompt_vulnerable and response_vulnerable
    
    def _map_failed_prompts_to_atlas(self, results: list, vulnerabilities_found: int):
        """Map failed prompts to MITRE ATLAS taxonomies and update the mapping file."""
        try:
            import json
            from pathlib import Path
            
            # Load the ATLAS taxonomies file
            atlas_file = Path(__file__).parent / 'data' / 'mitre_atlas_taxonomies.json'
            
            if not atlas_file.exists():
                return
            
            with open(atlas_file, 'r', encoding='utf-8') as f:
                atlas_data = json.load(f)
            
            # Find failed prompts and map them to ATLAS
            failed_prompts = []
            for result in results:
                if result.get('status') == 'failed' and result.get('vulnerability_found', False):
                    prompt_id = result.get('metadata', {}).get('prompt_id')
                    if prompt_id:
                        # Get the prompt's MITRE ATLAS mapping from the adversarial prompts
                        prompt_atlas_mappings = self._get_prompt_atlas_mapping(prompt_id)
                        if prompt_atlas_mappings:
                            # Create an entry for each ATLAS mapping
                            for atlas_mapping in prompt_atlas_mappings:
                                failed_prompts.append({
                                    'prompt_id': prompt_id,
                                    'atlas_mapping': atlas_mapping,
                                    'attack_technique': result.get('metadata', {}).get('attack_technique', 'Unknown'),
                                    'vuln_category': result.get('metadata', {}).get('vuln_category', 'Unknown'),
                                    'severity': result.get('metadata', {}).get('severity', 'medium')
                                })
            
            # Update ATLAS taxonomies with failed prompt IDs
            if failed_prompts:
                for failed_prompt in failed_prompts:
                    atlas_mapping = failed_prompt['atlas_mapping']
                    prompt_id = failed_prompt['prompt_id']
                    
                    # Add prompt ID to the technique's prompt_ids list
                    if atlas_mapping in atlas_data['techniques']:
                        # Check if prompt_id is already in the list to avoid duplicates
                        if prompt_id not in atlas_data['techniques'][atlas_mapping]['prompt_ids']:
                            atlas_data['techniques'][atlas_mapping]['prompt_ids'].append(prompt_id)
                        else:
                            pass
                    else:
                        pass
                
                # Save updated ATLAS data
                with open(atlas_file, 'w', encoding='utf-8') as f:
                    json.dump(atlas_data, f, indent=2, ensure_ascii=False)
                
            else:
                pass
                
        except Exception as e:
            raise
    
    def _get_prompt_atlas_mapping(self, prompt_id: str) -> list:
        """Get the MITRE ATLAS mapping for a specific prompt ID."""
        try:
            # Load adversarial prompts to find the mapping
            from .storage.yaml_loader import YAMLLoader
            yaml_loader = YAMLLoader()
            prompts = yaml_loader.load_adversarial_prompts()
            
            for prompt in prompts:
                if prompt.get('id') == prompt_id:
                    mapping = prompt.get('mitreatlasmapping', [])
                    # Ensure it's a list
                    if isinstance(mapping, list):
                        return mapping
                    elif isinstance(mapping, str):
                        return [mapping] if mapping else []
                    else:
                        return []
            
            return []
        except Exception as e:
            return []
    
    def _map_failed_prompts_to_owasp(self, results: list, vulnerabilities_found: int):
        """Map failed prompts to OWASP LLM categories and update the mapping file."""
        try:
            import json
            from pathlib import Path
            
            # Load the OWASP LLM taxonomies file
            owasp_file = Path(__file__).parent / 'data' / 'owasp_llm_taxonomies.json'
            
            if not owasp_file.exists():
                return
            
            with open(owasp_file, 'r', encoding='utf-8') as f:
                owasp_data = json.load(f)
            
            # Find failed prompts and map them to OWASP
            failed_prompts = []
            for result in results:
                if result.get('status') == 'failed' and result.get('vulnerability_found', False):
                    prompt_id = result.get('metadata', {}).get('prompt_id')
                    if prompt_id:
                        # Get the prompt's OWASP mapping from the adversarial prompts
                        prompt_owasp_mapping = self._get_prompt_owasp_mapping(prompt_id)
                        if prompt_owasp_mapping:
                            failed_prompts.append({
                                'prompt_id': prompt_id,
                                'owasp_mapping': prompt_owasp_mapping,
                                'attack_technique': result.get('metadata', {}).get('attack_technique', 'Unknown'),
                                'vuln_category': result.get('metadata', {}).get('vuln_category', 'Unknown'),
                                'severity': result.get('metadata', {}).get('severity', 'medium')
                            })
            
            # Update OWASP taxonomies with failed prompt IDs
            if failed_prompts:
                for failed_prompt in failed_prompts:
                    owasp_mapping = failed_prompt['owasp_mapping']
                    prompt_id = failed_prompt['prompt_id']
                    
                    # Add prompt ID to the category's prompt_ids list
                    if owasp_mapping in owasp_data['categories']:
                        # Check if prompt_id is already in the list to avoid duplicates
                        if prompt_id not in owasp_data['categories'][owasp_mapping]['prompt_ids']:
                            owasp_data['categories'][owasp_mapping]['prompt_ids'].append(prompt_id)
                        else:
                            pass
                    else:
                        pass
                
                # Save updated OWASP data
                with open(owasp_file, 'w', encoding='utf-8') as f:
                    json.dump(owasp_data, f, indent=2, ensure_ascii=False)
                
            else:
                pass
                
        except Exception as e:
            raise
    
    def _get_prompt_owasp_mapping(self, prompt_id: str) -> str:
        """Get the OWASP LLM mapping for a specific prompt ID."""
        try:
            # Load adversarial prompts to find the mapping
            from .storage.yaml_loader import YAMLLoader
            yaml_loader = YAMLLoader()
            prompts = yaml_loader.load_adversarial_prompts()
            
            for prompt in prompts:
                if prompt.get('id') == prompt_id:
                    owasp_mapping = prompt.get('owasp_top10_llm_mapping', [''])[0] if prompt.get('owasp_top10_llm_mapping') else ''
                    # Add :2025 suffix if not already present
                    if owasp_mapping and not owasp_mapping.endswith(':2025'):
                        owasp_mapping += ':2025'
                    return owasp_mapping
            
            return ''
        except Exception as e:
            return ''
        
        # New API endpoints for local storage operations
        
        @self.app.route('/api/local/prompts', methods=['GET'])
        def get_local_prompts():
            """Get all local adversarial prompts."""
            try:
                prompts = self.storage.get_all_local_prompts()
                return jsonify({'prompts': prompts})
            except Exception as e:
                logger.error(f"Error getting local prompts: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/prompts', methods=['POST'])
        def create_local_prompt():
            """Create a new local adversarial prompt."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                prompt_id = self.storage.create_local_prompt(data)
                if prompt_id:
                    return jsonify({'id': prompt_id, 'message': 'Prompt created successfully'}), 201
                else:
                    return jsonify({'error': 'Failed to create prompt'}), 400
                    
            except Exception as e:
                logger.error(f"Error creating local prompt: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/collections', methods=['GET'])
        def get_local_collections():
            """Get all local collections."""
            try:
                collections = self.storage.get_all_local_collections()
                return jsonify({'collections': collections})
            except Exception as e:
                logger.error(f"Error getting local collections: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/collections', methods=['POST'])
        def create_local_collection():
            """Create a new local collection."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                collection_id = self.storage.create_local_collection(data)
                if collection_id:
                    return jsonify({'id': collection_id, 'message': 'Collection created successfully'}), 201
                else:
                    return jsonify({'error': 'Failed to create collection'}), 400
                    
            except Exception as e:
                logger.error(f"Error creating local collection: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/attack-types', methods=['GET'])
        def get_local_attack_types():
            """Get all local attack types."""
            try:
                attack_types = self.storage.get_all_local_attack_types()
                return jsonify({'attack_types': attack_types})
            except Exception as e:
                logger.error(f"Error getting local attack types: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/attack-types', methods=['POST'])
        def create_local_attack_type():
            """Create a new local attack type."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                attack_id = self.storage.create_local_attack_type(data)
                if attack_id:
                    return jsonify({'id': attack_id, 'message': 'Attack type created successfully'}), 201
                else:
                    return jsonify({'error': 'Failed to create attack type'}), 400
                    
            except Exception as e:
                logger.error(f"Error creating local attack type: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/vulnerability-categories', methods=['GET'])
        def get_local_vulnerability_categories():
            """Get all local vulnerability categories."""
            try:
                categories = self.storage.get_all_local_vulnerability_categories()
                return jsonify({'vulnerability_categories': categories})
            except Exception as e:
                logger.error(f"Error getting local vulnerability categories: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/vulnerability-categories', methods=['POST'])
        def create_local_vulnerability_category():
            """Create a new local vulnerability category."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                category_id = self.storage.create_local_vulnerability_category(data)
                if category_id:
                    return jsonify({'id': category_id, 'message': 'Vulnerability category created successfully'}), 201
                else:
                    return jsonify({'error': 'Failed to create vulnerability category'}), 400
                    
            except Exception as e:
                logger.error(f"Error creating local vulnerability category: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/session-configs', methods=['GET'])
        def get_local_session_configs():
            """Get all local session configurations."""
            try:
                configs = self.storage.get_all_local_session_configs()
                return jsonify({'session_configs': configs})
            except Exception as e:
                logger.error(f"Error getting local session configs: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/local/session-configs', methods=['POST'])
        def create_local_session_config():
            """Create a new local session configuration."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No data provided'}), 400
                
                config_id = self.storage.create_local_session_config(data)
                if config_id:
                    return jsonify({'id': config_id, 'message': 'Session config created successfully'}), 201
                else:
                    return jsonify({'error': 'Failed to create session config'}), 400
                    
            except Exception as e:
                logger.error(f"Error creating local session config: {e}")
                return jsonify({'error': str(e)}), 500
        
        
        
        # Grading API endpoints
        @self.app.route('/api/grading/assertions', methods=['GET'])
        def get_grading_assertions():
            """Get list of available grading assertion types."""
            try:
                assertions = self.grading_engine.get_available_assertions()
                assertion_configs = {}
                
                for assertion_type in assertions:
                    config = self.grading_engine.get_assertion_config(assertion_type)
                    if config:
                        assertion_configs[assertion_type] = config
                
                return jsonify({
                    'assertions': assertions,
                    'configs': assertion_configs
                })
            except Exception as e:
                logger.error(f"Error getting grading assertions: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/providers', methods=['GET'])
        def get_grading_providers():
            """Get list of available grading providers."""
            try:
                providers = self.grading_engine.get_available_providers()
                provider_configs = {}
                
                for provider_name in providers:
                    config = self.grading_engine.get_provider_config(provider_name)
                    if config:
                        provider_configs[provider_name] = config
                
                return jsonify({
                    'providers': providers,
                    'configs': provider_configs
                })
            except Exception as e:
                logger.error(f"Error getting grading providers: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/grade', methods=['POST'])
        def grade_output():
            """Grade an output using specified assertion type."""
            try:
                data = request.get_json()
                if not data:
                    return jsonify({'error': 'No JSON data provided'}), 400
                
                required_fields = ['output', 'assertion_type']
                for field in required_fields:
                    if field not in data:
                        return jsonify({'error': f'Missing required field: {field}'}), 400
                
                # Create grading request
                grading_request = GradingRequest(
                    output=data['output'],
                    assertion_type=data['assertion_type'],
                    assertion_params=data.get('assertion_params', {}),
                    provider_override=data.get('provider_override'),
                    timeout=data.get('timeout')
                )
                
                # Execute grading
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                response = loop.run_until_complete(self.grading_engine.grade(grading_request))
                loop.close()
                
                return jsonify({
                    'result': {
                        'pass': response.result.pass_result,
                        'score': response.result.score,
                        'reason': response.result.reason,
                        'assertion_type': response.result.assertion_type,
                        'error': response.result.error
                    },
                        'metadata': {
                            'provider_used': response.provider_used,
                            'metadata': response.metadata
                        }
                })
                
            except Exception as e:
                logger.error(f"Error grading output: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/grade-batch', methods=['POST'])
        def grade_batch():
            """Grade multiple outputs in parallel."""
            try:
                data = request.get_json()
                if not data or 'requests' not in data:
                    return jsonify({'error': 'No requests provided'}), 400
                
                # Create grading requests
                grading_requests = []
                for req_data in data['requests']:
                    grading_request = GradingRequest(
                        output=req_data['output'],
                        assertion_type=req_data['assertion_type'],
                        assertion_params=req_data.get('assertion_params', {}),
                        provider_override=req_data.get('provider_override'),
                        timeout=req_data.get('timeout')
                    )
                    grading_requests.append(grading_request)
                
                # Execute batch grading
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                responses = loop.run_until_complete(self.grading_engine.grade_batch(grading_requests))
                loop.close()
                
                results = []
                for response in responses:
                    results.append({
                        'result': {
                            'pass': response.result.pass_result,
                            'score': response.result.score,
                            'reason': response.result.reason,
                            'assertion_type': response.result.assertion_type,
                            'error': response.result.error
                        },
                        'metadata': {
                            'provider_used': response.provider_used,
                            'metadata': response.metadata
                        }
                    })
                
                return jsonify({'results': results})
                
            except Exception as e:
                logger.error(f"Error batch grading: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/llm-rubric', methods=['POST'])
        def grade_llm_rubric_endpoint():
            """Grade output against a rubric using LLM."""
            try:
                data = request.get_json()
                if not data or 'output' not in data or 'rubric' not in data:
                    return jsonify({'error': 'Missing required fields: output, rubric'}), 400
                
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                response = loop.run_until_complete(grade_llm_rubric(
                    output=data['output'],
                    rubric=data['rubric'],
                    engine=self.grading_engine,
                    provider=data.get('provider')
                ))
                loop.close()
                
                return jsonify({
                    'result': {
                        'pass': response.result.pass_result,
                        'score': response.result.score,
                        'reason': response.result.reason,
                        'error': response.result.error
                    },
                    'metadata': {
                        'provider_used': response.provider_used
                    }
                })
                
            except Exception as e:
                logger.error(f"Error grading with LLM rubric: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/harmful', methods=['POST'])
        def grade_harmful_endpoint():
            """Check if output contains harmful content."""
            try:
                data = request.get_json()
                if not data or 'output' not in data:
                    return jsonify({'error': 'Missing required field: output'}), 400
                
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                response = loop.run_until_complete(grade_harmful_content(
                    output=data['output'],
                    engine=self.grading_engine,
                    provider=data.get('provider')
                ))
                loop.close()
                
                return jsonify({
                    'result': {
                        'pass': response.result.pass_result,
                        'score': response.result.score,
                        'reason': response.result.reason,
                        'error': response.result.error
                    },
                    'metadata': {
                        'provider_used': response.provider_used
                    }
                })
                
            except Exception as e:
                logger.error(f"Error grading harmful content: {e}")
                return jsonify({'error': str(e)}), 500
        
        @self.app.route('/api/grading/injection', methods=['POST'])
        def grade_injection_endpoint():
            """Check if prompt injection was successful."""
            try:
                data = request.get_json()
                if not data or 'output' not in data or 'original_prompt' not in data:
                    return jsonify({'error': 'Missing required fields: output, original_prompt'}), 400
                
                import asyncio
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                response = loop.run_until_complete(grade_prompt_injection(
                    output=data['output'],
                    original_prompt=data['original_prompt'],
                    engine=self.grading_engine,
                    provider=data.get('provider')
                ))
                loop.close()
                
                return jsonify({
                    'result': {
                        'pass': response.result.pass_result,
                        'score': response.result.score,
                        'reason': response.result.reason,
                        'error': response.result.error
                    },
                    'metadata': {
                        'provider_used': response.provider_used
                    }
                })
                
            except Exception as e:
                logger.error(f"Error grading prompt injection: {e}")
                return jsonify({'error': str(e)}), 500
    
    def _setup_socketio(self):
        """Set up SocketIO event handlers for real-time communication."""
        
        @self.socketio.on('connect')
        def handle_connect():
            print(f'Client connected: {request.sid}')
            emit('status', {'message': 'Connected to Avenlis server'})
        
        @self.socketio.on('disconnect')
        def handle_disconnect():
            print(f'Client disconnected: {request.sid}')
    

    def start(self):
        """Start the web server."""
        try:
            self.socketio.run(
                self.app,
                host=self.host,
                port=self.port,
                debug=self.debug,
                allow_unsafe_werkzeug=True  # For development
            )
        except OSError as e:
            if "Address already in use" in str(e):
                raise AvenlisError(f"Port {self.port} is already in use. Try a different port with --port")
            raise AvenlisError(f"Failed to start server: {e}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Avenlis SandStrike Backend Server")
    parser.add_argument("--host", default="127.0.0.1", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8080, help="Port to bind to")
    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
    
    args = parser.parse_args()
    
    try:
        server = AvenlisServer(host=args.host, port=args.port, debug=args.debug)
        print(f"Starting Avenlis server on {args.host}:{args.port}")
        server.start()
    except KeyboardInterrupt:
        print("\nServer stopped by user")
    except Exception as e:
        print(f"Server error: {e}")
        sys.exit(1)
