#!/usr/bin/env python3
"""
app.py — BLUX-cA Ideology Demo (Gradio)

Enhanced version based on https://github.com/Outer-Void/blux-ca

Features:
- Constitution spine: concise rules guiding every reply
- Enhanced discernment compass with BLUX-cA patterns
- Explainable routing: shows "why this response" trace
- Advanced redaction: PI/phone/email/SSN scrubbing
- Append-only audit: JSONL log per session (downloadable)
- Conversational memory with context tracking
- BLUX-cA aligned response strategies
- Improved error handling and validation

Run:
  pip install gradio>=4.44.0
  python app.py
"""

from __future__ import annotations
import json, re, time, uuid, logging, tempfile, shutil
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Any
from enum import Enum

import gradio as gr

# --- Configure logging ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- BLUX-cA: Enhanced Constitution (BLUX-cA aligned) ------------------------
class ConstitutionRule(Enum):
    INTEGRITY = "Integrity over approval; truth over comfort."
    STRATEGY = "Strategy: enlighten, not humiliate."
    SAFETY = "Defaults: do-no-harm, de-escalate, safeguard minors."
    TRANSPARENCY = "Transparency: explain routing & decisions."
    AUTONOMY = "Autonomy & dignity: offer choices, never coerce."
    CONTEXT = "Context-aware: adapt to user's stated needs and constraints."

CONSTITUTION = [rule.value for rule in ConstitutionRule]

# --- Enhanced redaction patterns (BLUX-cA security standards) ----------------
class RedactionPatterns:
    EMAIL = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
    PHONE = re.compile(r"\+?\d[\d\-\s\(\)]{7,}\d")
    IP = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")
    SSN = re.compile(r"\d{3}-\d{2}-\d{4}")
    CREDIT_CARD = re.compile(r"\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}")
    
    @classmethod
    def get_patterns(cls) -> Dict[str, re.Pattern]:
        return {
            "emails": cls.EMAIL,
            "phones": cls.PHONE,
            "ips": cls.IP,
            "ssn": cls.SSN,
            "credit_cards": cls.CREDIT_CARD,
        }

def redact(text: str) -> Tuple[str, Dict[str, List[str]]]:
    """Enhanced redaction with comprehensive pattern matching."""
    found = {}
    redacted_text = text
    
    for category, pattern in RedactionPatterns.get_patterns().items():
        matches = pattern.findall(redacted_text)
        found[category] = matches
        if matches:
            redacted_text = pattern.sub(f"[redacted:{category}]", redacted_text)
    
    return redacted_text, found

# --- Enhanced Discernment Compass (BLUX-cA patterns) -------------------------
@dataclass
class CompassResult:
    state: str           # "struggler" | "indulger" | "seeker" | "neutral" | "crisis"
    signals: List[str]   # matched cues
    tactic: str          # response strategy
    confidence: float    # 0.0 to 1.0
    risk_level: str      # "low" | "medium" | "high"

class UserState(Enum):
    STRUGGLER = "struggler"      # Hopelessness, overwhelm
    INDULGER = "indulger"        # Blame, deflection
    SEEKER = "seeker"           # Help-seeking, growth-oriented
    NEUTRAL = "neutral"         # Information request
    CRISIS = "crisis"           # Urgent safety concerns

# Enhanced cue patterns based on BLUX-cA research
CRISIS_CUES = [
    "kill myself", "end it all", "suicide", "harm myself", "don't want to live",
    "can't go on", "life isn't worth", "better off dead"
]

STRUGGLER_CUES = [
    "it's hopeless", "i give up", "worthless", "i'm broken", "failure",
    "nobody cares", "can't change", "no control", "overwhelmed", "stuck",
    "nothing works", "always fails"
]

INDULGER_CUES = [
    "their fault", "they made me", "always them", "not on me", "because of them",
    "if only they", "it's all their", "they should", "they need to change"
]

SEEKER_CUES = [
    "how can i", "what should i", "help me", "advice", "suggestions",
    "looking for guidance", "trying to improve", "want to learn",
    "how to handle", "ways to", "strategies for"
]

def discern(message: str) -> CompassResult:
    """Enhanced discernment with confidence scoring and risk assessment."""
    m = message.lower()
    signals = []
    risk_level = "low"
    
    # Crisis detection (highest priority)
    if any(k in m for k in CRISIS_CUES):
        signals.append("crisis_indicators")
        risk_level = "high"
        return CompassResult(
            state=UserState.CRISIS.value,
            signals=signals,
            tactic="crisis_protocol",
            confidence=0.95,
            risk_level=risk_level
        )
    
    # Other signal detection
    if any(k in m for k in STRUGGLER_CUES):
        signals.append("hopelessness/overwhelm")
        risk_level = "medium" if risk_level == "low" else risk_level
    
    if any(k in m for k in INDULGER_CUES):
        signals.append("deflection/blame")
        risk_level = "medium" if risk_level == "low" else risk_level
    
    if any(k in m for k in SEEKER_CUES):
        signals.append("help_seeking/growth")
    
    # State determination with confidence scoring
    confidence = min(0.3 + (len(signals) * 0.2), 0.9)  # Base confidence
    
    if "hopelessness/overwhelm" in signals and "deflection/blame" not in signals:
        return CompassResult(
            state=UserState.STRUGGLER.value,
            signals=signals,
            tactic="validate+tools",
            confidence=confidence,
            risk_level=risk_level
        )
    
    if "deflection/blame" in signals and "help_seeking/growth" not in signals:
        return CompassResult(
            state=UserState.INDULGER.value,
            signals=signals,
            tactic="boundary+off-ramp", 
            confidence=confidence,
            risk_level=risk_level
        )
    
    if "help_seeking/growth" in signals:
        return CompassResult(
            state=UserState.SEEKER.value,
            signals=signals,
            tactic="inform+coach+expand",
            confidence=confidence,
            risk_level=risk_level
        )
    
    return CompassResult(
        state=UserState.NEUTRAL.value,
        signals=signals or ["information_request"],
        tactic="inform+coach",
        confidence=0.3,
        risk_level=risk_level
    )

# --- Enhanced Tactics Library (BLUX-cA aligned) ------------------------------
class ResponseTactics:
    @staticmethod
    def validate_tools(gentle: bool = True) -> str:
        base = (
            "I hear the weight in your words. Let's pick one small, doable step you can own today. "
            "Name a tiny target (5–10 minutes). I'll help break it down and set a checkpoint. "
        )
        if gentle:
            base += "We'll move at your pace—what feels manageable right now?"
        else:
            base += "What's the smallest action you can commit to?"
        return base
    
    @staticmethod
    def boundary_off_ramp(gentle: bool = True) -> str:
        base = (
            "I can't co-sign blame, but I can help you find what you control. "
            "If you want movement, choose one action within your power. "
        )
        if gentle:
            base += "We can map the situation together, or pause until you're ready to focus on your agency."
        else:
            base += "Let's focus on what you can influence directly."
        return base
    
    @staticmethod
    def inform_coach(gentle: bool = True) -> str:
        base = (
            "Tell me the goal, current constraints, and your time window. "
            "I'll outline options, risks, and a minimal path forward."
        )
        if gentle:
            base += " No pressure—we'll find what works for your situation."
        return base
    
    @staticmethod
    def inform_coach_expand(gentle: bool = True) -> str:
        base = (
            "Great that you're seeking growth. Let's explore this systematically. "
            "What's your main objective? I'll provide frameworks and actionable steps."
        )
        if gentle:
            base += " We can adapt based on what resonates with you."
        return base
    
    @staticmethod
    def crisis_protocol(gentle: bool = True) -> str:
        return (
            "🚨 **Urgent Support Needed** 🚨\n\n"
            "Your safety is the absolute priority. Please contact emergency services now:\n"
            "• **Emergency**: 911 (US) or your local emergency number\n"
            "• **Crisis Text Line**: Text HOME to 741741\n"
            "• **National Suicide Prevention Lifeline**: 988 (US)\n\n"
            "These services are available 24/7 with trained professionals who can help immediately. "
            "Please reach out now—you don't have to face this alone."
        )

def respond_with_tactic(user_input: str, tactic: str, gentle: bool, compass: CompassResult) -> str:
    """Enhanced response generation with context awareness."""
    tactics = ResponseTactics()
    
    if tactic == "validate+tools":
        return tactics.validate_tools(gentle)
    elif tactic == "boundary+off-ramp":
        return tactics.boundary_off_ramp(gentle)
    elif tactic == "inform+coach+expand":
        return tactics.inform_coach_expand(gentle)
    elif tactic == "crisis_protocol":
        return tactics.crisis_protocol(gentle)
    else:  # inform+coach (default)
        return tactics.inform_coach(gentle)

# --- Enhanced Audit System ---------------------------------------------------
@dataclass
class AuditEvent:
    id: str
    ts: float
    user_input: str
    user_redactions: Dict[str, List[str]]
    compass: Dict[str, Any]
    tactic: str
    constitution: List[str]
    assistant: str
    session_id: str
    version: str = "blux-ca-1.1"

class SessionManager:
    def __init__(self):
        self.session_id = str(uuid.uuid4())
        self.start_time = time.time()
        self.temp_dir = Path(tempfile.mkdtemp(prefix="blux_ca_"))
        logger.info(f"Session temp directory: {self.temp_dir}")
    
    def get_session_info(self) -> Dict[str, Any]:
        return {
            "session_id": self.session_id,
            "start_time": self.start_time,
            "duration": round(time.time() - self.start_time, 2),
            "temp_dir": str(self.temp_dir)
        }
    
    def create_temp_audit_file(self) -> Path:
        """Create a temporary audit file that Gradio can safely serve."""
        audit_file = self.temp_dir / f"blux_ca_audit_{self.session_id}.jsonl"
        audit_file.touch()
        return audit_file
    
    def cleanup(self):
        """Clean up temporary files."""
        try:
            if self.temp_dir.exists():
                shutil.rmtree(self.temp_dir)
                logger.info(f"Cleaned up temp directory: {self.temp_dir}")
        except Exception as e:
            logger.warning(f"Failed to clean up temp directory: {e}")

session_manager = SessionManager()

def make_audit_event(
    user_text: str, 
    redactions: Dict[str, List[str]], 
    compass: CompassResult, 
    reply: str
) -> AuditEvent:
    """Create a comprehensive audit event."""
    return AuditEvent(
        id=str(uuid.uuid4()),
        ts=time.time(),
        user_input=user_text,
        user_redactions=redactions,
        compass=asdict(compass),
        tactic=compass.tactic,
        constitution=CONSTITUTION,
        assistant=reply,
        session_id=session_manager.session_id,
    )

# --- Core Chat Handler -------------------------------------------------------
def ca_chat(
    history: List[Dict[str, str]],
    message: str,
    enable_redaction: bool,
    gentle_tone: bool,
    use_temp_logging: bool  # New parameter to control logging location
) -> Tuple[List[Dict[str, str]], Dict[str, Any], Optional[str]]:
    """
    Enhanced chat handler with proper Gradio messages format and error handling.
    """
    try:
        if not message or not message.strip():
            logger.warning("Empty message received")
            return history, {}, None

        raw_message = message.strip()
        logger.info(f"Processing message: {raw_message[:50]}...")

        # Redaction phase
        if enable_redaction:
            clean_message, found_redactions = redact(raw_message)
            logger.info(f"Redacted {sum(len(v) for v in found_redactions.values())} items")
        else:
            clean_message, found_redactions = raw_message, {"emails": [], "phones": [], "ips": [], "ssn": [], "credit_cards": []}

        # Discernment phase
        compass_result = discern(clean_message)
        logger.info(f"Discernment: state={compass_result.state}, tactic={compass_result.tactic}")

        # Response generation
        reply = respond_with_tactic(clean_message, compass_result.tactic, gentle_tone, compass_result)

        # Create trace information
        trace_info = {
            "compass_state": compass_result.state,
            "signals": compass_result.signals,
            "applied_tactic": compass_result.tactic,
            "confidence": round(compass_result.confidence, 2),
            "risk_level": compass_result.risk_level,
            "redaction_enabled": enable_redaction,
            "gentle_mode": gentle_tone
        }

        trace_md = (
            f"**Trace** — state: `{trace_info['compass_state']}` | "
            f"signals: `{', '.join(trace_info['signals']) or 'none'}` | "
            f"tactic: `{trace_info['applied_tactic']}` | "
            f"confidence: `{trace_info['confidence']}` | "
            f"risk: `{trace_info['risk_level']}`"
        )

        full_reply = f"{reply}\n\n---\n{trace_md}"

        # Update history with proper message format
        updated_history = history + [
            {"role": "user", "content": raw_message},
            {"role": "assistant", "content": full_reply}
        ]

        # Audit logging
        audit_event = make_audit_event(clean_message, found_redactions, compass_result, reply)
        audit_dict = asdict(audit_event)
        
        download_path = None
        
        if use_temp_logging:
            try:
                # Use temporary directory that Gradio can access
                temp_audit_file = session_manager.create_temp_audit_file()
                with temp_audit_file.open("a", encoding="utf-8") as f:
                    f.write(json.dumps(audit_dict, ensure_ascii=False) + "\n")
                download_path = str(temp_audit_file)
                logger.info(f"Audit event written to temporary file: {temp_audit_file}")
            except Exception as e:
                logger.error(f"Failed to write temporary audit log: {e}")
        else:
            # Original logging to custom path (without download capability)
            try:
                log_path = Path("~/.outer_void/audit/blux_ca_demo.jsonl").expanduser()
                log_path.parent.mkdir(parents=True, exist_ok=True)
                with log_path.open("a", encoding="utf-8") as f:
                    f.write(json.dumps(audit_dict, ensure_ascii=False) + "\n")
                logger.info(f"Audit event written to persistent log: {log_path}")
            except Exception as e:
                logger.error(f"Failed to write persistent audit log: {e}")

        return updated_history, audit_dict, download_path

    except Exception as e:
        logger.error(f"Error in ca_chat: {e}")
        error_reply = "I apologize, but I encountered an error processing your message. Please try again."
        error_history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": error_reply}
        ]
        return error_history, {"error": str(e)}, None

# --- Enhanced UI with BLUX-cA Styling ----------------------------------------
def build_ui():
    """Build the enhanced Gradio interface."""
    with gr.Blocks(
        title="BLUX-cA Demo", 
        css="""
        footer {visibility: hidden}
        .constitution-box {border: 1px solid #e0e0e0; border-radius: 8px; padding: 16px; margin: 8px 0;}
        .risk-high {color: #d32f2f; font-weight: bold;}
        .risk-medium {color: #f57c00; font-weight: bold;}
        .risk-low {color: #388e3c; font-weight: bold;}
        """
    ) as demo:
        
        gr.Markdown("""
        # 🌌 BLUX-cA Demo
        *Context-Aware Constitutional AI*
        
        This demo showcases the BLUX-cA ideology: a transparent, constitutionally-guided AI 
        that explains its reasoning and adapts to your needs.
        """)

        with gr.Row():
            with gr.Column(scale=2):
                # Enhanced Chatbot with proper type
                chat = gr.Chatbot(
                    height=500, 
                    label="Conversation",
                    type="messages",
                    show_copy_button=True,
                    avatar_images=(
                        "https://em-content.zobj.net/source/microsoft/319/robot_1f916.png", 
                        "https://em-content.zobj.net/source/microsoft/319/brain_1f9e0.png"
                    )
                )
                
                msg = gr.Textbox(
                    placeholder="Share your context, ask for help, or describe a challenge...",
                    autofocus=True,
                    lines=2,
                    max_lines=5
                )
                
                with gr.Row():
                    redact_toggle = gr.Checkbox(
                        True, 
                        label="🔒 Redact personal data (email/phone/IP/SSN)",
                        info="Protects your privacy before analysis"
                    )
                    gentle_tone = gr.Checkbox(
                        True, 
                        label="🌱 Gentle tone mode",
                        info="Softer, more supportive language"
                    )
                    temp_logging = gr.Checkbox(
                        True,
                        label="📥 Enable file downloads",
                        info="Allows downloading audit logs (uses temp files)"
                    )
                
                with gr.Row():
                    send_btn = gr.Button("Send", variant="primary", size="lg")
                    clear_btn = gr.Button("Clear History", variant="secondary")

            with gr.Column(scale=1):
                # Constitution display
                gr.Markdown("### 📜 Constitution")
                with gr.Group(elem_classes="constitution-box"):
                    for i, rule in enumerate(CONSTITUTION, 1):
                        gr.Markdown(f"{i}. {rule}")
                
                gr.Markdown("### 🔍 Live Analysis")
                with gr.Accordion("Current Session Info", open=True):
                    session_info = gr.JSON(
                        value=session_manager.get_session_info(),
                        label="Session"
                    )
                
                gr.Markdown("### 📊 Latest Audit Event")
                audit_json = gr.JSON(
                    label="Event Details",
                    value={}
                )
                
                gr.Markdown("### 💾 Download Audit Log")
                download = gr.File(
                    label="Session Audit Log",
                    visible=False,
                    file_types=[".jsonl"],
                    type="filepath"
                )

        # Event handlers
        def on_send(history, message, redact_enabled, gentle_enabled, temp_logging_enabled):
            if not message or not message.strip():
                gr.Warning("Please enter a message before sending.")
                return history, {}, None, session_manager.get_session_info()
            
            new_history, audit_data, download_path = ca_chat(
                history or [], message, redact_enabled, gentle_enabled, temp_logging_enabled
            )
            
            session_info = session_manager.get_session_info()
            
            # Show download button only if we have a downloadable file
            download_visible = download_path is not None
            download_value = download_path if download_path else None
            
            return new_history, audit_data, gr.update(visible=download_visible, value=download_value), session_info

        def on_clear():
            # Clean up old session and start new one
            session_manager.cleanup()
            session_manager.__init__()  # Reinitialize for new session
            
            return [], {}, gr.update(visible=False, value=None), session_manager.get_session_info()

        # Connect components
        send_btn.click(
            fn=on_send,
            inputs=[chat, msg, redact_toggle, gentle_tone, temp_logging],
            outputs=[chat, audit_json, download, session_info]
        ).then(
            lambda: "",  # Clear message input
            outputs=[msg]
        )

        msg.submit(
            fn=on_send,
            inputs=[chat, msg, redact_toggle, gentle_tone, temp_logging],
            outputs=[chat, audit_json, download, session_info]
        ).then(
            lambda: "",  # Clear message input
            outputs=[msg]
        )

        clear_btn.click(
            fn=on_clear,
            outputs=[chat, audit_json, download, session_info]
        )

        gr.Markdown("""
        ### 🎯 How It Works
        
        **Discernment Compass**: Analyzes your message for patterns (hopelessness, blame, help-seeking, crisis)
        
        **Constitutional Alignment**: Every response follows our core principles
        
        **Transparency**: See the reasoning behind each response in the trace
        
        **Safety First**: Automatic crisis detection with immediate resource guidance
        
        **Privacy**: Optional redaction of personal information before processing
        
        ### 📁 File Downloads
        
        Enable "Download Audit Logs" to get a JSONL file containing all conversation events.
        Files are stored in a temporary directory and cleaned up when you clear the history.
        """)

    return demo

# --- Main Execution ----------------------------------------------------------
if __name__ == "__main__":
    logger.info("Starting BLUX-cA Demo...")
    
    # Validate dependencies
    try:
        import gradio
        logger.info(f"Gradio version: {gradio.__version__}")
    except ImportError:
        logger.error("Gradio not installed. Run: pip install gradio>=4.44.0")
        exit(1)
    
    try:
        # Launch the application with allowed paths for file downloads
        demo = build_ui()
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False,
            show_error=True,
            debug=False,
            ssr_mode=False,
            # Allow temp directory access for file downloads
            allowed_paths=[session_manager.temp_dir]
        )
    except KeyboardInterrupt:
        logger.info("Received interrupt signal - shutting down gracefully...")
    except Exception as e:
        logger.error(f"Failed to launch application: {e}")
    finally:
        # Clean up on exit
        session_manager.cleanup()