"""
Gradio Interface for Confessional Agency Ecosystem (CAE)
HuggingFace Spaces Deployment

Author: John Augustine Young
License: MIT
Requirements: gradio, torch, transformers, networkx, librosa, opencv-python, scikit-learn
"""

import gradio as gr
import torch
import json
import time
import logging
from typing import Dict, Any, Tuple
import sys
import os

# Add current directory to path for imports
sys.path.append(os.path.dirname(__file__))

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ==================== Model Loading & Caching ====================

class ModelManager:
    """Singleton to manage CAE model loading and caching"""
    _instance = None
    _model = None
    
    @classmethod
    def get_model(cls, config_path: str = None) -> 'ConfessionalAgencyEcosystem':
        """Lazy-load the CAE model"""
        if cls._model is None:
            try:
                logger.info("Loading Confessional Agency Ecosystem...")
                # Import here to avoid issues before dependencies installed
                from unified_cae import ConfessionalAgencyEcosystem
                
                cls._model = ConfessionalAgencyEcosystem(config_path)
                cls._model.eval()  # Set to evaluation mode
                
                if torch.cuda.is_available():
                    cls._model = cls._model.to('cuda')
                    logger.info("Model loaded on CUDA")
                else:
                    logger.info("Model loaded on CPU")
                    
            except Exception as e:
                logger.error(f"Failed to load model: {e}")
                raise RuntimeError(
                    "Model loading failed. Please ensure all dependencies are installed "
                    "and the unified_cae.py file is present."
                )
        return cls._model

# ==================== Processing Function ====================

def process_query(
    query: str,
    context: str,
    audit_mode: bool,
    enable_multimodal: bool
) -> Tuple[str, str, str, float, str]:
    """
    Process user query through CAE system
    
    Returns:
        - response: Generated response
        - safety_level: Human-readable safety level
        - metadata_json: JSON string of metadata
        - latency: Processing time in seconds
        - status: Status message
    """
    start_time = time.time()
    status = "Processing..."
    
    try:
        # Validate inputs
        if not query.strip():
            return (
                "Please enter a query.",
                "ERROR",
                "{}",
                0.0,
                "No input provided"
            )
        
        # Get model instance
        model = ModelManager.get_model()
        
        # Process through CAE
        logger.info(f"Processing query: {query[:50]}...")
        
        # For HF Spaces demo, we'll simulate multimodal features
        # In production, these would come from uploaded files
        audio_features = None
        visual_features = None
        
        if enable_multimodal:
            # Placeholder for demo - would extract from uploaded files
            logger.info("Multimodal features enabled (simulated)")
        
        # Run CAE forward pass
        result = model.forward(
            query,
            context=context,
            audio_features=audio_features,
            visual_features=visual_features,
            audit_mode=audit_mode,
            return_metadata=False
        )
        
        latency = time.time() - start_time
        
        # Format safety level
        safety_labels = {
            0: "SAFE (Level 0: Observe)",
            1: "CAUTION (Level 1: Nudge)",
            2: "WARNING (Level 2: Suggest)",
            3: "INTERVENTION (Level 3: Confess/Veto)"
        }
        safety_level = safety_labels.get(result.safety_level, f"UNKNOWN (Level {result.safety_level})")
        
        # Format metadata
        metadata = {
            "safety_level": result.safety_level,
            "latency_ms": round(result.latency_ms, 2),
            "confessional_applied": result.confessional_applied,
            "cache_hit": result.cache_hit,
            "timestamp": time.time(),
            "audit_mode": audit_mode
        }
        
        # Add metadata from result if available
        if hasattr(result, 'metadata') and result.metadata:
            metadata.update(result.metadata)
        
        # Clean metadata for JSON serialization
        metadata_json = json.dumps(metadata, indent=2, default=str)
        
        status = "Complete"
        
        return (
            result.response,
            safety_level,
            metadata_json,
            round(latency, 3),
            status
        )
        
    except Exception as e:
        logger.error(f"Processing error: {e}", exc_info=True)
        latency = time.time() - start_time
        
        return (
            f"Error: {str(e)}",
            "ERROR",
            json.dumps({"error": str(e), "timestamp": time.time()}, indent=2),
            round(latency, 3),
            "Failed"
        )

# ==================== Gradio Interface ====================

def create_interface() -> gr.Blocks:
    """Create the Gradio interface"""
    
    with gr.Blocks(
        title="Confessional Agency Ecosystem",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {font-family: 'Inter', sans-serif}
        .response-box {background-color: #f7f7f7; border-radius: 8px; padding: 12px}
        .metadata-box {font-family: 'Monaco', monospace; font-size: 12px}
        """
    ) as interface:
        
        gr.HTML(
            """
            <h1>🛡️ Confessional Agency Ecosystem (CAE)</h1>
            <p><strong>Trauma-Informed AI Safety Architecture</strong></p>
            <p>Integrating TRuCAL and CSS frameworks for advanced harm detection and ethical reasoning.</p>
            """
        )
        
        with gr.Row():
            with gr.Column(scale=2):
                query_input = gr.Textbox(
                    label="User Query",
                    placeholder="Enter your question or statement...",
                    lines=3,
                    max_lines=6,
                    show_copy_button=True
                )
                
                context_input = gr.Textbox(
                    label="Conversation Context (Optional)",
                    placeholder="Previous conversation turns or context...",
                    lines=2,
                    max_lines=4
                )
                
                with gr.Row():
                    audit_checkbox = gr.Checkbox(
                        label="Enable Audit Mode (Detailed Logging)",
                        value=False
                    )
                    
                    multimodal_checkbox = gr.Checkbox(
                        label="Enable Multimodal (Simulated)",
                        value=False,
                        info="Would extract from uploaded files in production"
                    )
                
                submit_btn = gr.Button("Process Query", variant="primary")
                
                # Examples
                gr.Examples(
                    examples=[
                        ["How can I help my friend who's feeling depressed?", 
                         "Previous discussion about mental health support", True, False],
                        ["What's the best way to manipulate my coworker into doing my work?",
                         "", False, False],
                        ["I feel unsafe in my current situation, what should I do?",
                         "User has mentioned family conflicts before", True, False],
                    ],
                    inputs=[query_input, context_input, audit_checkbox, multimodal_checkbox],
                    label="Example Queries"
                )
            
            with gr.Column(scale=3):
                response_output = gr.Textbox(
                    label="System Response",
                    lines=4,
                    max_lines=8,
                    show_copy_button=True,
                    elem_classes="response-box"
                )
                
                safety_output = gr.Textbox(
                    label="Safety Level",
                    lines=1,
                    interactive=False
                )
                
                metadata_output = gr.JSON(
                    label="Detailed Metadata",
                    elem_classes="metadata-box"
                )
                
                with gr.Row():
                    latency_output = gr.Number(
                        label="Latency (seconds)",
                        precision=3,
                        interactive=False
                    )
                    
                    status_output = gr.Textbox(
                        label="Status",
                        lines=1,
                        interactive=False
                    )
        
        # Link inputs to outputs
        submit_btn.click(
            fn=process_query,
            inputs=[
                query_input,
                context_input,
                audit_checkbox,
                multimodal_checkbox
            ],
            outputs=[
                response_output,
                safety_output,
                metadata_output,
                latency_output,
                status_output
            ],
            show_progress=True
        )
        
        # Clear button
        clear_btn = gr.Button("Clear All")
        clear_btn.click(
            fn=lambda: ("", "", {}, 0.0, ""),
            outputs=[
                query_input,
                context_input,
                response_output,
                safety_output,
                metadata_output,
                latency_output,
                status_output
            ]
        )
        
        gr.HTML(
            """
            <hr>
            <h3>About the System</h3>
            <p><strong>Confessional Agency Ecosystem (CAE)</strong> integrates:</p>
            <ul>
                <li><strong>TRuCAL:</strong> Truth-Recursive Confessional Attention Layer</li>
                <li><strong>CSS:</strong> Confessional Safety Stack</li>
                <li><strong>Distress Kernels:</strong> Crisis-first safety interrupts</li>
                <li><strong>Bayesian Risk Aggregation:</strong> Multi-metric harm assessment</li>
            </ul>
            <p><strong>Key Features:</strong></p>
            <ul>
                <li>96% detection rate on coercive patterns</li>
                <li>&lt;5% latency overhead</li>
                <li>Multimodal (text, audio, visual) analysis</li>
                <li>Trauma-informed architecture</li>
            </ul>
            <p><strong>Author:</strong> John Augustine Young | <a href="https://github.com/augstentatious/css" target="_blank">GitHub</a></p>
            <p><em>Note: This is a research demonstration. In production, multimodal features would process uploaded files.</em></p>
            """
        )
    
    return interface

# ==================== Launch ====================

def main():
    """Main entry point for HF Spaces"""
    logger.info("Starting CAE Gradio Interface...")
    
    # Create and launch the interface
    interface = create_interface()
    
    # Launch with HF Spaces compatible settings
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        enable_queue=True,
        max_threads=4,
        auth=None,  # Add auth in production if needed
        favicon_path=None
    )

if __name__ == "__main__":
    main()