src/interface/prompt_handlers.py

import logging
import gradio as gr
from typing import Optional
from src.interface.session_manager import SimplifiedSessionData

logger = logging.getLogger(__name__)

def format_prompt_with_html(prompt_text: str) -> str:
    """Format prompt with HTML tags for better visualization."""
    import re
    import html
    
    # Escape HTML first to prevent injection
    formatted = html.escape(prompt_text)
    
    # Highlight XML-like opening tags
    formatted = re.sub(
        r'<([a-z_]+)(?:\s+[^>]*)?>',
        r'<span style="color: #2563eb; font-weight: 600;">&lt;\1&gt;</span>',
        formatted,
        flags=re.IGNORECASE
    )
    
    # Highlight XML-like closing tags
    formatted = re.sub(
        r'&lt;/([a-z_]+)&gt;',
        r'<span style="color: #7c3aed; font-weight: 600;">&lt;/\1&gt;</span>',
        formatted,
        flags=re.IGNORECASE
    )
    
    # Highlight XML attributes
    formatted = re.sub(
        r'([a-z_]+)=&quot;([^&quot;]+)&quot;',
        r'<span style="color: #059669;">\1</span>=<span style="color: #d97706;">&quot;\2&quot;</span>',
        formatted,
        flags=re.IGNORECASE
    )
    
    # Highlight bullet points
    formatted = re.sub(
        r'^([-•]\s+.+)$',
        r'<div style="margin-left: 1.5em; color: #059669;">\1</div>',
        formatted,
        flags=re.MULTILINE
    )
    
    # Highlight numbered lists
    formatted = re.sub(
        r'^(\d+\.\s+.+)$',
        r'<div style="margin-left: 1.5em; color: #dc2626;">\1</div>',
        formatted,
        flags=re.MULTILINE
    )
    
    # Highlight important keywords
    keywords = [
        'IMPORTANT', 'CRITICAL', 'REQUIRED', 'MUST', 'SHALL',
        'WARNING', 'NOTE', 'TASK', 'GOAL', 'OUTPUT', 'ANY'
    ]
    for keyword in keywords:
        formatted = re.sub(
            f'\\b({keyword})\\b',
            r'<strong style="color: #dc2626; background-color: #fef2f2; padding: 2px 4px; border-radius: 3px;">\1</strong>',
            formatted,
            flags=re.IGNORECASE
        )
    
    # Highlight JSON/code blocks
    formatted = re.sub(
        r'(\{[^}]+\})',
        r'<code style="background-color: #f3f4f6; padding: 2px 6px; border-radius: 3px; font-family: monospace; font-size: 0.9em;">\1</code>',
        formatted
    )
    
    # Convert newlines to <br> for proper display
    formatted = formatted.replace('\n', '<br>')
    
    # Wrap in container with monospace font for code-like appearance
    formatted = f'<div style="font-family: \'Courier New\', monospace; line-height: 1.8; padding: 1em; background-color: #f9fafb; border-radius: 8px; overflow-x: auto;">{formatted}</div>'
    
    return formatted

def _prompt_name_to_agent(prompt_name: str) -> str:
    """Map UI prompt selection to internal agent/prompt key."""
    mapping = {
        "🔍 Spiritual Monitor (Classifier)": "SpiritualDistressAnalyzer",
        "🟡 Soft Spiritual Triage": "SoftSpiritualTriage",
        "📊 Triage Response Evaluator": "TriageResponseEvaluator",
        "🏥 Medical Assistant": "MedicalAssistant",
        "🩺 Soft Medical Triage": "SoftMedicalTriage",
    }
    return mapping.get(prompt_name, prompt_name)

def load_prompt(prompt_name: str, session: Optional[SimplifiedSessionData] = None):
    """Load selected prompt for editing using enhanced prompt editor."""
    try:
        from src.interface.enhanced_prompt_editor import EnhancedPromptEditor
        
        # Initialize enhanced editor
        editor = EnhancedPromptEditor()
        
        # Get session ID
        session_id = getattr(session, 'session_id', 'default_session') if session else 'default_session'
        
        # Use enhanced editor to load prompt
        prompt_content, info_html, status_html = editor.load_prompt_for_editing(prompt_name, session_id)
        
        return prompt_content, info_html, status_html
        
    except Exception as e:
        # Fallback to old system if enhanced editor fails
        logger.warning(f"Enhanced prompt editor failed, using fallback: {e}")
        
        from src.core.spiritual_monitor import SYSTEM_PROMPT_SPIRITUAL_MONITOR
        from src.core.soft_triage_manager import (
            SYSTEM_PROMPT_TRIAGE_QUESTION,
            SYSTEM_PROMPT_TRIAGE_EVALUATE
        )
        from src.config.prompts import (
            SYSTEM_PROMPT_MEDICAL_ASSISTANT,
            SYSTEM_PROMPT_SOFT_MEDICAL_TRIAGE
        )
        
        prompts = {
            "🔍 Spiritual Monitor (Classifier)": SYSTEM_PROMPT_SPIRITUAL_MONITOR,
            "🟡 Soft Spiritual Triage": SYSTEM_PROMPT_TRIAGE_QUESTION,
            "📊 Triage Response Evaluator": SYSTEM_PROMPT_TRIAGE_EVALUATE,
            "🏥 Medical Assistant": SYSTEM_PROMPT_MEDICAL_ASSISTANT,
            "🩺 Soft Medical Triage": SYSTEM_PROMPT_SOFT_MEDICAL_TRIAGE
        }
        
        prompt_text = prompts.get(prompt_name, "")
        
        info = f"""**Loaded:** {prompt_name}
**Length:** {len(prompt_text)} characters
**Status:** Fallback mode (enhanced editor unavailable)"""
        
        status = """<div style="padding: 1em; background-color: #fffbeb; border-left: 4px solid #f59e0b; border-radius: 4px;">
<h4 style="color: #d97706; margin-top: 0;">⚠️ Fallback Mode</h4>
<p style="margin-bottom: 0;">Using basic prompt editor. Enhanced features unavailable.</p>
</div>"""
        
        return prompt_text, info, status

def apply_prompt_changes(prompt_name: str, prompt_text: str, session: SimplifiedSessionData):
    """Apply custom prompt changes using enhanced prompt editor."""
    try:
        from src.interface.enhanced_prompt_editor import EnhancedPromptEditor
        
        if session is None:
            session = SimplifiedSessionData()
        
        # Initialize enhanced editor
        editor = EnhancedPromptEditor()
        
        # Get session ID
        session_id = getattr(session, 'session_id', 'default_session')
        
        # Use enhanced editor to apply changes
        status_html, success = editor.apply_prompt_changes(prompt_name, prompt_text, session_id)
        
        if success:
            # Also store in session for backward compatibility
            if not hasattr(session, 'custom_prompts'):
                session.custom_prompts = {}
            
            agent_key = _prompt_name_to_agent(prompt_name)
            session.custom_prompts[agent_key] = prompt_text
            
            # Apply to session app instance if available
            if hasattr(session, 'app_instance') and hasattr(session.app_instance, 'set_prompt_overrides'):
                session.app_instance.set_prompt_overrides(session.custom_prompts)
        
        return status_html, session
        
    except Exception as e:
        # Fallback to old system
        logger.warning(f"Enhanced prompt editor failed, using fallback: {e}")
        
        if session is None:
            session = SimplifiedSessionData()
        
        if not prompt_text.strip():
            error_html = """<div style="padding: 1em; background-color: #fef2f2; border-left: 4px solid #dc2626; border-radius: 4px;">
<h4 style="color: #dc2626; margin-top: 0;">❌ Error</h4>
<p style="margin-bottom: 0;">Prompt cannot be empty</p>
</div>"""
            return error_html, session
        
        # Store custom prompt in session (session-scoped)
        if not hasattr(session, 'custom_prompts'):
            session.custom_prompts = {}
        
        agent_key = _prompt_name_to_agent(prompt_name)
        session.custom_prompts[agent_key] = prompt_text

        status = f"""<div style="padding: 1em; background-color: #fffbeb; border-left: 4px solid #f59e0b; border-radius: 4px;">
<h4 style="color: #d97706; margin-top: 0;">⚠️ Fallback Mode - Changes Applied</h4>
<p><strong>Prompt:</strong> {prompt_name}</p>
<p><strong>Length:</strong> {len(prompt_text)} characters</p>
<p style="margin-bottom: 0;">Enhanced features unavailable, using basic session storage.</p>
</div>"""
        
        return status, session

def reset_prompt(prompt_name: str, session: SimplifiedSessionData):
    """Reset prompt to default using enhanced prompt editor."""
    try:
        from src.interface.enhanced_prompt_editor import EnhancedPromptEditor
        
        if session is None:
            session = SimplifiedSessionData()
        
        # Initialize enhanced editor
        editor = EnhancedPromptEditor()
        
        # Get session ID
        session_id = getattr(session, 'session_id', 'default_session')
        
        # Use enhanced editor to reset prompt
        prompt_content, info_html, status_html = editor.reset_prompt_to_default(prompt_name, session_id)
        
        # Also remove from session for backward compatibility
        agent_key = _prompt_name_to_agent(prompt_name)
        if hasattr(session, 'custom_prompts') and agent_key in session.custom_prompts:
            del session.custom_prompts[agent_key]

        # Apply to session app instance if available
        if hasattr(session, 'app_instance') and hasattr(session.app_instance, 'set_prompt_overrides'):
            session.app_instance.set_prompt_overrides(getattr(session, 'custom_prompts', {}))
        
        return prompt_content, info_html, status_html, session
        
    except Exception as e:
        # Fallback to old system
        logger.warning(f"Enhanced prompt editor failed, using fallback: {e}")
        
        if session is None:
            session = SimplifiedSessionData()
        
        # Remove from custom prompts
        agent_key = _prompt_name_to_agent(prompt_name)
        if hasattr(session, 'custom_prompts') and agent_key in session.custom_prompts:
            del session.custom_prompts[agent_key]
        
        # Reload default
        prompt_text, info, status = load_prompt(prompt_name, session)
        
        reset_status = """<div style="padding: 1em; background-color: #fffbeb; border-left: 4px solid #f59e0b; border-radius: 4px;">
<h4 style="color: #d97706; margin-top: 0;">🔄 Fallback Mode - Reset Complete</h4>
<p style="margin-bottom: 0;">Prompt restored using basic system. Enhanced features unavailable.</p>
</div>"""
        
        return prompt_text, info, reset_status, session

def promote_prompt_to_file(prompt_name: str, session: SimplifiedSessionData):
    """Promote session prompt override to permanent file."""
    try:
        from src.interface.enhanced_prompt_editor import EnhancedPromptEditor
        
        if session is None:
            return """<div style="padding: 1em; background-color: #fef2f2; border-left: 4px solid #dc2626; border-radius: 4px;">
<h4 style="color: #dc2626; margin-top: 0;">❌ Error</h4>
<p style="margin-bottom: 0;">No session data available</p>
</div>""", session
        
        # Initialize enhanced editor
        editor = EnhancedPromptEditor()
        
        # Get session ID
        session_id = getattr(session, 'session_id', 'default_session')
        
        # Use enhanced editor to promote prompt
        status_html, success = editor.promote_session_to_file(prompt_name, session_id)
        
        return status_html, session
        
    except Exception as e:
        logger.warning(f"Enhanced prompt editor failed: {e}")
        return f"""<div style="padding: 1em; background-color: #fef2f2; border-left: 4px solid #dc2626; border-radius: 4px;">
<h4 style="color: #dc2626; margin-top: 0;">❌ Error</h4>
<p style="margin-bottom: 0;">Failed to promote prompt: {str(e)}</p>
</div>""", session

def validate_prompt_syntax(prompt_text: str):
    """Validate prompt syntax and structure."""
    try:
        from src.interface.enhanced_prompt_editor import EnhancedPromptEditor
        
        # Initialize enhanced editor
        editor = EnhancedPromptEditor()
        
        # Use enhanced editor to validate prompt
        validation_html, is_valid = editor.validate_prompt_syntax(prompt_text)
        
        return validation_html
        
    except Exception as e:
        logger.warning(f"Enhanced prompt editor failed: {e}")
        return f"""<div style="padding: 1em; background-color: #fef2f2; border-left: 4px solid #dc2626; border-radius: 4px;">
<h4 style="color: #dc2626; margin-top: 0;">❌ Validation Error</h4>
<p style="margin-bottom: 0;">Failed to validate prompt: {str(e)}</p>
</div>"""