app.py

import gradio as gr
import json
import logging
import logging.handlers
import numpy as np
from datetime import datetime
from typing import Dict, Any, List, Optional, Tuple
import threading
import time
import os
import sqlite3
import contextlib
import signal
import sys
import functools
import csv
import io
import tempfile
from collections import deque
from scipy.stats import beta
import plotly.graph_objects as go

# ----------------------------------------------------------------------
# Prometheus metrics (optional)
# ----------------------------------------------------------------------
try:
    from prometheus_client import Counter, Histogram, Gauge, generate_latest, CONTENT_TYPE_LATEST
    PROMETHEUS_AVAILABLE = True
except ImportError:
    PROMETHEUS_AVAILABLE = False
    logging.warning("prometheus-client not installed, metrics endpoint disabled")

# ----------------------------------------------------------------------
# Configuration from environment variables with validation
# ----------------------------------------------------------------------
LOW_THRESHOLD = float(os.getenv("ARF_LOW_THRESHOLD", "0.2"))
HIGH_THRESHOLD = float(os.getenv("ARF_HIGH_THRESHOLD", "0.8"))
ALPHA_PRIOR = float(os.getenv("ARF_ALPHA_PRIOR", "1.0"))
BETA_PRIOR = float(os.getenv("ARF_BETA_PRIOR", "1.0"))
DB_PATH = os.getenv("ARF_DB_PATH", "/data/arf_decisions.db")
LOG_LEVEL = os.getenv("ARF_LOG_LEVEL", "INFO").upper()
VERSION = "4.2.0+oss-enhanced"

# Validate thresholds
if not (0 <= LOW_THRESHOLD < HIGH_THRESHOLD <= 1):
    logging.warning(f"Invalid thresholds: low={LOW_THRESHOLD}, high={HIGH_THRESHOLD}. Using defaults.")
    LOW_THRESHOLD = 0.2
    HIGH_THRESHOLD = 0.8

# Validate priors
if ALPHA_PRIOR <= 0 or BETA_PRIOR <= 0:
    logging.warning(f"Invalid priors: alpha={ALPHA_PRIOR}, beta={BETA_PRIOR}. Using defaults.")
    ALPHA_PRIOR = 1.0
    BETA_PRIOR = 1.0

# ----------------------------------------------------------------------
# Logging setup (file rotation + console)
# ----------------------------------------------------------------------
os.makedirs("/var/log/arf", exist_ok=True)
logger = logging.getLogger(__name__)
logger.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))

file_handler = logging.handlers.RotatingFileHandler(
    "/var/log/arf/app.log", maxBytes=10_485_760, backupCount=5
)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))

logger.addHandler(file_handler)
logger.addHandler(console_handler)
logger.propagate = False

# ----------------------------------------------------------------------
# SQLite persistence with secure permissions
# ----------------------------------------------------------------------
def init_db():
    db_dir = os.path.dirname(DB_PATH)
    if db_dir and not os.path.exists(db_dir):
        os.makedirs(db_dir, exist_ok=True)
    with contextlib.closing(sqlite3.connect(DB_PATH)) as conn:
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS decisions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                decision_json TEXT NOT NULL,
                risk REAL NOT NULL
            )
        ''')
        conn.commit()
    try:
        os.chmod(DB_PATH, 0o600)
    except Exception as e:
        logger.warning(f"Could not set secure permissions on DB: {e}")
    logger.info(f"Database initialized at {DB_PATH}")

def save_decision_to_db(decision: dict, risk: float):
    try:
        with contextlib.closing(sqlite3.connect(DB_PATH)) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO decisions (timestamp, decision_json, risk) VALUES (?, ?, ?)",
                (decision["timestamp"], json.dumps(decision), risk)
            )
            conn.commit()
    except Exception as e:
        logger.error(f"Failed to save decision to DB: {e}")

def load_recent_decisions(limit: int = 100) -> List[Tuple[str, dict, float]]:
    decisions = []
    try:
        with contextlib.closing(sqlite3.connect(DB_PATH)) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT timestamp, decision_json, risk FROM decisions ORDER BY timestamp DESC LIMIT ?",
                (limit,)
            )
            rows = cursor.fetchall()
            for ts, json_str, risk in rows:
                decisions.append((ts, json.loads(json_str), risk))
        decisions.reverse()
    except Exception as e:
        logger.error(f"Failed to load decisions from DB: {e}")
    return decisions

def vacuum_db():
    try:
        with contextlib.closing(sqlite3.connect(DB_PATH)) as conn:
            conn.execute("VACUUM")
        logger.info("Database vacuumed")
    except Exception as e:
        logger.error(f"Vacuum failed: {e}")

# ----------------------------------------------------------------------
# Prometheus metrics (renamed to avoid conflict with Gradio components)
# ----------------------------------------------------------------------
if PROMETHEUS_AVAILABLE:
    prom_decisions_total = Counter('arf_decisions_total', 'Total decisions made', ['action'])
    prom_risk_gauge = Gauge('arf_current_risk', 'Current risk score')
    prom_decision_latency = Histogram('arf_decision_latency_seconds', 'Time to evaluate intent')
    prom_mcmc_runs = Counter('arf_mcmc_runs_total', 'Total MCMC runs')
else:
    prom_decisions_total = None
    prom_risk_gauge = None
    prom_decision_latency = None
    prom_mcmc_runs = None

# ----------------------------------------------------------------------
# Thread‑safe history
# ----------------------------------------------------------------------
decision_history = []
risk_history = []
history_lock = threading.Lock()
shutdown_event = threading.Event()

def update_dashboard_data(decision: dict, risk: float):
    with history_lock:
        decision_history.append((datetime.utcnow().isoformat(), decision, risk))
        risk_history.append((datetime.utcnow().isoformat(), risk))
        if len(decision_history) > 100:
            decision_history.pop(0)
        if len(risk_history) > 100:
            risk_history.pop(0)
    save_decision_to_db(decision, risk)
    if PROMETHEUS_AVAILABLE:
        prom_decisions_total.labels(action=decision.get("risk_level", "unknown")).inc()
        prom_risk_gauge.set(risk)

def refresh_history_from_db():
    global decision_history, risk_history
    decisions = load_recent_decisions(100)
    with history_lock:
        decision_history.clear()
        risk_history.clear()
        for ts, dec, risk in decisions:
            decision_history.append((ts, dec, risk))
            risk_history.append((ts, risk))
    if PROMETHEUS_AVAILABLE and risk_history:
        prom_risk_gauge.set(risk_history[-1][1])

# ----------------------------------------------------------------------
# Memory monitoring
# ----------------------------------------------------------------------
def get_memory_usage():
    try:
        import resource
        rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        if rss < 1e9:
            return rss / 1024.0
        else:
            return rss / (1024.0 * 1024.0)
    except ImportError:
        try:
            with open("/proc/self/status") as f:
                for line in f:
                    if line.startswith("VmRSS:"):
                        parts = line.split()
                        if len(parts) >= 2:
                            return int(parts[1]) / 1024.0
        except Exception:
            pass
    return None

def memory_monitor_loop():
    while not shutdown_event.is_set():
        try:
            mem_mb = get_memory_usage()
            if mem_mb is not None:
                logger.info(f"Process memory: {mem_mb:.1f} MB")
            else:
                logger.info("Process memory: unknown")
        except Exception as e:
            logger.error(f"Memory logging error: {e}")
        for _ in range(60):
            if shutdown_event.is_set():
                break
            time.sleep(1)

# ----------------------------------------------------------------------
# Bayesian Risk Engine
# ----------------------------------------------------------------------
class BayesianRiskEngine:
    def __init__(self, alpha=ALPHA_PRIOR, beta=BETA_PRIOR, maxlen=None):
        self.alpha = alpha
        self.beta = beta
        self.maxlen = maxlen
        self.events = deque(maxlen=maxlen)
        self.total_failures = 0
        self.total_successes = 0

    def update(self, failures, successes):
        self.events.append((failures, successes))
        self.total_failures += failures
        self.total_successes += successes
        if self.maxlen is not None and len(self.events) == self.maxlen:
            self.total_failures = sum(f for f, _ in self.events)
            self.total_successes = sum(s for _, s in self.events)
        self.alpha = ALPHA_PRIOR + self.total_failures
        self.beta = BETA_PRIOR + self.total_successes

    def risk(self):
        return self.alpha / (self.alpha + self.beta)

    def risk_interval(self, prob=0.95):
        lo = beta.ppf((1 - prob) / 2, self.alpha, self.beta)
        hi = beta.ppf((1 + prob) / 2, self.alpha, self.beta)
        return lo, hi

# ----------------------------------------------------------------------
# Policy Engine
# ----------------------------------------------------------------------
class PolicyEngine:
    def __init__(self):
        self.thresholds = {"low": LOW_THRESHOLD, "high": HIGH_THRESHOLD}

    def evaluate(self, risk):
        if risk < self.thresholds["low"]:
            return "approve", "Risk within safe limits"
        elif risk > self.thresholds["high"]:
            return "deny", f"Risk exceeds high threshold ({self.thresholds['high']})"
        else:
            return "escalate", f"Risk in escalation zone ({self.thresholds['low']}-{self.thresholds['high']})"

# ----------------------------------------------------------------------
# Infrastructure analysis
# ----------------------------------------------------------------------
def handle_infra_with_governance(fault_type: str, context_window: int, session_state: dict):
    start_time = time.time()
    try:
        fault_type = fault_type.strip()
        if fault_type not in ["none", "switch_down", "server_overload", "cascade"]:
            fault_type = "none"
        context_window = max(0, min(1000, int(context_window)))

        fault_map = {
            "none": (1, 99),
            "switch_down": (20, 80),
            "server_overload": (35, 65),
            "cascade": (60, 40)
        }
        failures, successes = fault_map.get(fault_type, (1, 99))

        maxlen = context_window if context_window > 0 else None
        risk_engine = BayesianRiskEngine(maxlen=maxlen)
        risk_engine.update(failures, successes)
        risk = risk_engine.risk()
        ci_low, ci_high = risk_engine.risk_interval(0.95)

        policy_engine = PolicyEngine()
        action, reason = policy_engine.evaluate(risk)
        control_decision = autonomous_control_decision(risk, risk_engine, policy_engine)

        analysis_result = {
            "risk": risk,
            "risk_ci": [ci_low, ci_high],
            "decision": action,
            "justification": reason,
            "healing_actions": ["restart"] if action == "deny" else ["monitor"],
            "posterior_parameters": {
                "alpha": risk_engine.alpha,
                "beta": risk_engine.beta
            },
            "context_window": context_window
        }
        output = {
            **analysis_result,
            "governance": {
                "policy_evaluation": {
                    "action": action,
                    "reason": reason,
                    "thresholds": policy_engine.thresholds
                },
                "control_plane_decision": control_decision
            }
        }
        if PROMETHEUS_AVAILABLE:
            prom_decision_latency.observe(time.time() - start_time)
        return output, session_state
    except Exception as e:
        logger.exception("Error in handle_infra_with_governance")
        return {"error": str(e)}, session_state

def autonomous_control_decision(risk, risk_engine, policy_engine):
    action, reason = policy_engine.evaluate(risk)
    risk_level = "low" if risk < LOW_THRESHOLD else "medium" if risk < HIGH_THRESHOLD else "high"
    decision = {
        "timestamp": datetime.utcnow().isoformat(),
        "approved": action == "approve",
        "actions": ["escalate_human"] if action == "escalate" else [],
        "reason": reason,
        "risk_level": risk_level
    }
    update_dashboard_data(decision, risk)
    return decision

# ----------------------------------------------------------------------
# MCMC
# ----------------------------------------------------------------------
class MHMCMC:
    def __init__(self, log_target, proposal_sd=0.1):
        self.log_target = log_target
        self.proposal_sd = proposal_sd

    def sample(self, n_samples, initial_state, burn_in=0):
        samples = np.zeros((n_samples, len(initial_state)))
        current = np.array(initial_state)
        current_log = self.log_target(current)
        accepted = 0
        for i in range(n_samples + burn_in):
            proposal = current + np.random.normal(0, self.proposal_sd, size=len(current))
            proposal_log = self.log_target(proposal)
            accept_prob = min(1, np.exp(proposal_log - current_log))
            if np.random.rand() < accept_prob:
                current = proposal
                current_log = proposal_log
                accepted += 1
            if i >= burn_in:
                samples[i - burn_in] = current
        acceptance_rate = accepted / (n_samples + burn_in)
        return samples, acceptance_rate

def run_hmc_mcmc(samples: int, warmup: int):
    try:
        samples = max(500, min(10000, int(samples)))
        warmup = max(100, min(2000, int(warmup)))
        if PROMETHEUS_AVAILABLE:
            prom_mcmc_runs.inc()

        np.random.seed(42)
        data = np.random.normal(0.5, 0.2, 10)

        def log_prior(mu):
            return -0.5 * (mu ** 2)

        def log_likelihood(mu):
            return -0.5 * np.sum(((data - mu) / 0.2) ** 2)

        def log_posterior(mu):
            return log_prior(mu) + log_likelihood(mu)

        sampler = MHMCMC(log_posterior, proposal_sd=0.05)
        mu_samples, acceptance = sampler.sample(samples, initial_state=[0.0], burn_in=warmup)
        mu_samples = mu_samples.flatten()

        mean = np.mean(mu_samples)
        median = np.median(mu_samples)
        credible_interval = np.percentile(mu_samples, [2.5, 97.5])

        fig_trace = go.Figure()
        fig_trace.add_trace(go.Scatter(y=mu_samples, mode='lines', name='μ', line=dict(width=1)))
        fig_trace.update_layout(title="Trace of μ (Metropolis-Hastings)", xaxis_title="Iteration", yaxis_title="μ")

        fig_hist = go.Figure()
        fig_hist.add_trace(go.Histogram(x=mu_samples, nbinsx=50, name='Posterior'))
        fig_hist.update_layout(title="Posterior Distribution of μ", xaxis_title="μ", yaxis_title="Density")

        summary = {
            "mean": mean,
            "median": median,
            "credible_interval_95": f"[{credible_interval[0]:.3f}, {credible_interval[1]:.3f}]",
            "acceptance_rate": f"{acceptance:.2%}"
        }
        return summary, fig_trace, fig_hist
    except Exception as e:
        logger.exception("MCMC computation failed")
        return {"error": str(e)}, go.Figure(), go.Figure()

# ----------------------------------------------------------------------
# Dashboard plots
# ----------------------------------------------------------------------
class TTLCache:
    def __init__(self, ttl_seconds=5):
        self.ttl = ttl_seconds
        self.cache = {}
        self.lock = threading.Lock()

    def __call__(self, func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            key = (func.__name__, args, frozenset(kwargs.items()))
            now = time.time()
            with self.lock:
                if key in self.cache:
                    result, timestamp = self.cache[key]
                    if now - timestamp < self.ttl:
                        return result
            result = func(*args, **kwargs)
            with self.lock:
                self.cache[key] = (result, now)
            return result
        return wrapper

dashboard_cache = TTLCache(ttl_seconds=2)

@dashboard_cache
def generate_risk_gauge():
    with history_lock:
        if not risk_history:
            return go.Figure()
        latest_risk = risk_history[-1][1]
    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=latest_risk,
        title={'text': "Current Risk"},
        gauge={
            'axis': {'range': [0, 1]},
            'bar': {'color': "darkblue"},
            'steps': [
                {'range': [0, LOW_THRESHOLD], 'color': "lightgreen"},
                {'range': [LOW_THRESHOLD, HIGH_THRESHOLD], 'color': "yellow"},
                {'range': [HIGH_THRESHOLD, 1], 'color': "red"}
            ]
        }))
    return fig

@dashboard_cache
def generate_decision_pie():
    with history_lock:
        if not decision_history:
            return go.Figure()
        approved = sum(1 for _, d, _ in decision_history if d.get("approved", False))
        blocked = len(decision_history) - approved
    fig = go.Figure(data=[go.Pie(labels=["Approved", "Blocked"], values=[approved, blocked])])
    fig.update_layout(title="Policy Decisions")
    return fig

@dashboard_cache
def generate_action_timeline():
    with history_lock:
        if not decision_history:
            return go.Figure()
        times = [d["timestamp"] for _, d, _ in decision_history]
        approvals = [1 if d.get("approved", False) else 0 for _, d, _ in decision_history]
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=times, y=approvals, mode='markers+lines', name='Approvals'))
    fig.update_layout(title="Autonomous Actions Timeline", xaxis_title="Time", yaxis_title="Approved (1) / Blocked (0)")
    return fig

@dashboard_cache
def generate_risk_trend():
    with history_lock:
        if not risk_history:
            return go.Figure()
        times = [ts for ts, _ in risk_history]
        risks = [r for _, r in risk_history]
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=times, y=risks, mode='lines+markers', name='Risk', line=dict(color='red', width=2)))
    fig.add_hline(y=LOW_THRESHOLD, line_dash="dash", line_color="green", annotation_text=f"Low ({LOW_THRESHOLD})")
    fig.add_hline(y=HIGH_THRESHOLD, line_dash="dash", line_color="orange", annotation_text=f"High ({HIGH_THRESHOLD})")
    fig.update_layout(title="Risk Trend", xaxis_title="Time", yaxis_title="Risk Score", yaxis_range=[0, 1])
    return fig

def refresh_dashboard():
    with history_lock:
        total = len(decision_history)
        approved = sum(1 for _, d, _ in decision_history if d.get("approved", False))
        blocked = total - approved
        avg_risk = np.mean([r for _, r in risk_history]) if risk_history else 0.5
        control_stats = {
            "total_decisions": total,
            "approved_actions": approved,
            "blocked_actions": blocked,
            "average_risk": float(avg_risk)
        }
    return (
        control_stats,
        generate_risk_gauge(),
        generate_decision_pie(),
        generate_action_timeline(),
        generate_risk_trend()
    )

# ----------------------------------------------------------------------
# Batch simulation
# ----------------------------------------------------------------------
def run_batch_simulation(context_window: int):
    fault_types = ["none", "switch_down", "server_overload", "cascade"]
    results = []
    for fault in fault_types:
        output, _ = handle_infra_with_governance(fault, context_window, {})
        if "error" in output:
            results.append([
                fault,
                "Error",
                output["error"],
                "N/A",
                "N/A"
            ])
        else:
            results.append([
                fault,
                f"{output['risk']:.4f}",
                output["decision"],
                output["governance"]["control_plane_decision"]["risk_level"],
                f"[{output['risk_ci'][0]:.3f}, {output['risk_ci'][1]:.3f}]"
            ])
    # Return as list of lists for gr.Dataframe
    return results

# ----------------------------------------------------------------------
# Data export
# ----------------------------------------------------------------------
def export_history_to_csv():
    try:
        with contextlib.closing(sqlite3.connect(DB_PATH)) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT timestamp, decision_json, risk FROM decisions ORDER BY timestamp")
            rows = cursor.fetchall()
        if not rows:
            return None
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow(["Timestamp", "Decision", "Risk", "Approved", "Risk Level", "Reason"])
        for ts, json_str, risk in rows:
            dec = json.loads(json_str)
            writer.writerow([
                ts,
                json_str,
                risk,
                dec.get("approved", False),
                dec.get("risk_level", ""),
                dec.get("reason", "")
            ])
        output.seek(0)
        # Write to a temporary file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
            f.write(output.getvalue())
            return f.name
    except Exception as e:
        logger.error(f"Export failed: {e}")
        return None

# ----------------------------------------------------------------------
# Update thresholds
# ----------------------------------------------------------------------
def update_thresholds(low: float, high: float):
    global LOW_THRESHOLD, HIGH_THRESHOLD
    if 0 <= low < high <= 1:
        LOW_THRESHOLD = low
        HIGH_THRESHOLD = high
        logger.info(f"Updated thresholds: low={low}, high={high}")
        return f"Thresholds updated: approve < {low}, escalate {low}-{high}, deny > {high}"
    else:
        return f"Invalid thresholds: low={low}, high={high}. Must satisfy 0 ≤ low < high ≤ 1."

# ----------------------------------------------------------------------
# OSS capabilities
# ----------------------------------------------------------------------
oss_caps = {
    "edition": "OSS (Demo)",
    "version": VERSION,
    "license": "Apache 2.0",
    "execution": {"modes": ["advisory"], "max_incidents": 100},
    "memory": {"type": "in-memory", "faiss_index_type": "flat", "max_incident_nodes": 100},
    "enterprise_features": [
        "Real-time HMC (using PyMC)",
        "Hyperpriors",
        "Decision Engine",
        "Full audit trails & compliance reporting",
        "Blast radius limits & automatic rollback",
        "Multi-cloud & hybrid deployment support"
    ]
}

# ----------------------------------------------------------------------
# Graceful shutdown
# ----------------------------------------------------------------------
def shutdown_handler(signum, frame):
    logger.info("Received shutdown signal, cleaning up...")
    shutdown_event.set()
    time.sleep(2)
    logger.info("Shutdown complete")
    sys.exit(0)

signal.signal(signal.SIGTERM, shutdown_handler)
signal.signal(signal.SIGINT, shutdown_handler)

# ----------------------------------------------------------------------
# Startup
# ----------------------------------------------------------------------
init_db()
refresh_history_from_db()

mem_thread = threading.Thread(target=memory_monitor_loop, daemon=True)
mem_thread.start()

def vacuum_scheduler():
    while not shutdown_event.is_set():
        time.sleep(86400)
        if not shutdown_event.is_set():
            vacuum_db()
vacuum_thread = threading.Thread(target=vacuum_scheduler, daemon=True)
vacuum_thread.start()

# ----------------------------------------------------------------------
# Gradio UI
# ----------------------------------------------------------------------
with gr.Blocks(title=f"ARF v{VERSION} – Bayesian Risk Scoring Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"""
    # 🧠 ARF v{VERSION} – Bayesian Risk Scoring for AI Reliability (Demo)
    **Mathematically rigorous risk estimation using conjugate priors and MCMC**
    This demo showcases:
    - **Bayesian conjugate prior (Beta-Binomial)** – online risk update from observed failures/successes.
    - **Policy thresholds** – approve (<{LOW_THRESHOLD}), escalate ({LOW_THRESHOLD}‑{HIGH_THRESHOLD}), deny (>{HIGH_THRESHOLD}).
    - **Metropolis-Hastings MCMC** – sampling from a posterior distribution (simulating HMC concepts).
    - **Autonomous control decisions** – based on the current risk estimate.
    All components are implemented with only `numpy`, `scipy`, and standard libraries.
    """)

    with gr.Tabs():
        # Control Plane Dashboard
        with gr.TabItem("Control Plane Dashboard"):
            gr.Markdown("### 🎮 Control Plane")
            with gr.Row():
                with gr.Column():
                    system_status = gr.JSON(label="System Status", value={
                        "edition": oss_caps["edition"],
                        "version": oss_caps["version"],
                        "governance_mode": "advisory",
                        "policies_loaded": 2,
                        "risk_threshold_low": LOW_THRESHOLD,
                        "risk_threshold_high": HIGH_THRESHOLD
                    })
                with gr.Column():
                    control_stats = gr.JSON(label="Control Statistics", value={
                        "total_decisions": 0,
                        "approved_actions": 0,
                        "blocked_actions": 0,
                        "average_risk": 0.5
                    })
            with gr.Row():
                risk_gauge = gr.Plot(label="Current Risk Gauge")
                decision_pie = gr.Plot(label="Policy Decisions")
            with gr.Row():
                action_timeline = gr.Plot(label="Autonomous Actions Timeline")
                risk_trend = gr.Plot(label="Risk Trend")
            with gr.Row():
                auto_refresh = gr.Checkbox(label="Auto-refresh (3s)", value=False)
                refresh_btn = gr.Button("Refresh Now")
            timer = gr.Timer(value=3, active=False)
            def refresh_if_enabled(auto):
                if auto:
                    return refresh_dashboard()
                else:
                    return [gr.update() for _ in range(5)]
            timer.tick(refresh_if_enabled, inputs=[auto_refresh], outputs=[control_stats, risk_gauge, decision_pie, action_timeline, risk_trend])
            refresh_btn.click(
                fn=refresh_dashboard,
                outputs=[control_stats, risk_gauge, decision_pie, action_timeline, risk_trend]
            )
            auto_refresh.change(lambda v: gr.Timer(active=v), inputs=[auto_refresh], outputs=[timer])

        # Infrastructure Reliability
        with gr.TabItem("Infrastructure Reliability"):
            gr.Markdown("### 🏗️ Infrastructure Intent Evaluation with Bayesian Risk")
            infra_state = gr.State(value={})
            with gr.Row():
                with gr.Column():
                    infra_fault = gr.Dropdown(
                        ["none", "switch_down", "server_overload", "cascade"],
                        value="none",
                        label="Inject Fault",
                        info="Select a fault type to simulate infrastructure issues."
                    )
                    context_window_input = gr.Number(
                        value=50,
                        label="Context Window (number of recent events)",
                        minimum=0,
                        maximum=1000,
                        step=1,
                        info="How many past incidents to consider for risk calculation (0 = unlimited)"
                    )
                    with gr.Row():
                        infra_btn = gr.Button("Evaluate Intent")
                        batch_btn = gr.Button("Run Batch Simulation", variant="secondary")
                with gr.Column():
                    infra_output = gr.JSON(label="Analysis Result")
            batch_results = gr.Dataframe(
                headers=["Fault Type", "Risk", "Decision", "Risk Level", "Confidence Interval"],
                label="Batch Simulation Results",
                datatype=["str", "str", "str", "str", "str"]
            )
            infra_btn.click(
                fn=handle_infra_with_governance,
                inputs=[infra_fault, context_window_input, infra_state],
                outputs=[infra_output, infra_state]
            )
            batch_btn.click(
                fn=run_batch_simulation,
                inputs=[context_window_input],
                outputs=[batch_results]
            )

        # Deep Analysis (MCMC)
        with gr.TabItem("Deep Analysis (MCMC)"):
            gr.Markdown("### Markov Chain Monte Carlo (Metropolis‑Hastings)")
            with gr.Row():
                with gr.Column():
                    hmc_samples = gr.Slider(500, 10000, value=5000, step=500, label="Number of Samples")
                    hmc_warmup = gr.Slider(100, 2000, value=1000, step=100, label="Burn‑in Steps")
                    hmc_run_btn = gr.Button("Run MCMC")
                with gr.Column():
                    hmc_summary = gr.JSON(label="Posterior Summary")
            with gr.Row():
                hmc_trace_plot = gr.Plot(label="Trace Plot")
                hmc_pair_plot = gr.Plot(label="Posterior Histogram")
            hmc_run_btn.click(
                fn=run_hmc_mcmc,
                inputs=[hmc_samples, hmc_warmup],
                outputs=[hmc_summary, hmc_trace_plot, hmc_pair_plot]
            )

        # Policy Management
        with gr.TabItem("Policy Management"):
            gr.Markdown("### 📋 Execution Policies")
            with gr.Row():
                low_slider = gr.Slider(0, 1, value=LOW_THRESHOLD, step=0.01, label="Low Threshold (Approve <)")
                high_slider = gr.Slider(0, 1, value=HIGH_THRESHOLD, step=0.01, label="High Threshold (Deny >)")
                update_thresh_btn = gr.Button("Update Thresholds")
            thresh_status = gr.Markdown(f"Current: approve < {LOW_THRESHOLD}, escalate {LOW_THRESHOLD}-{HIGH_THRESHOLD}, deny > {HIGH_THRESHOLD}")
            policies_json = [
                {"name": "Low Risk Policy", "conditions": [f"risk < {LOW_THRESHOLD}"], "action": "approve", "priority": 1},
                {"name": "Medium Risk Policy", "conditions": [f"{LOW_THRESHOLD} ≤ risk ≤ {HIGH_THRESHOLD}"], "action": "escalate", "priority": 2},
                {"name": "High Risk Policy", "conditions": [f"risk > {HIGH_THRESHOLD}"], "action": "deny", "priority": 3}
            ]
            policy_display = gr.JSON(label="Active Policies", value=policies_json)
            update_thresh_btn.click(
                fn=update_thresholds,
                inputs=[low_slider, high_slider],
                outputs=[thresh_status]
            ).then(
                fn=lambda: [
                    [
                        {"name": "Low Risk Policy", "conditions": [f"risk < {LOW_THRESHOLD}"], "action": "approve", "priority": 1},
                        {"name": "Medium Risk Policy", "conditions": [f"{LOW_THRESHOLD} ≤ risk ≤ {HIGH_THRESHOLD}"], "action": "escalate", "priority": 2},
                        {"name": "High Risk Policy", "conditions": [f"risk > {HIGH_THRESHOLD}"], "action": "deny", "priority": 3}
                    ]
                ],
                outputs=[policy_display]
            )

        # Enterprise / OSS
        with gr.TabItem("Enterprise / OSS"):
            gr.Markdown(f"""
            <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 2rem; border-radius: 12px; margin-bottom: 2rem; text-align: center; color: white;">
                <h1 style="margin: 0; font-size: 2.5rem;">🚀 ARF {oss_caps['edition'].upper()} Edition</h1>
                <p style="font-size: 1.2rem; opacity: 0.9;">Version {oss_caps['version']} · Apache 2.0 License</p>
            </div>

            <div style="display: flex; gap: 1.5rem; margin-bottom: 2rem;">
                <div style="flex: 1; background: #f8f9fa; padding: 1.5rem; border-radius: 8px;">
                    <h3>📦 OSS Capabilities (Demo)</h3>
                    <ul>
                        <li>✅ Bayesian conjugate prior – Beta-Binomial risk scoring</li>
                        <li>✅ Policy thresholds – configurable approve/escalate/deny</li>
                        <li>✅ MCMC sampling – Metropolis-Hastings (simulates HMC concepts)</li>
                        <li>✅ In-memory storage – no persistence</li>
                        <li>✅ Full open-source transparency</li>
                    </ul>
                </div>
                <div style="flex: 1; background: #f8f9fa; padding: 1.5rem; border-radius: 8px;">
                    <h3>🏢 Enterprise Features</h3>
                    <ul>
                        <li>🔒 Real-time HMC (using PyMC) – Bayesian deep learning for risk</li>
                        <li>🔒 Hyperpriors – hierarchical models for better generalization</li>
                        <li>🔒 Decision Engine with full audit trails</li>
                        <li>🔒 Blast radius limits & automatic rollback</li>
                        <li>🔒 Multi-cloud & hybrid deployment support</li>
                        <li>🔒 Compliance reporting (SOC2, ISO 27001)</li>
                        <li>🔒 24/7 enterprise support & SLAs</li>
                    </ul>
                </div>
            </div>

            <div style="background: #e9ecef; padding: 1.5rem; border-radius: 8px; text-align: center;">
                <h3 style="margin-top: 0;">✨ Why Upgrade to Enterprise?</h3>
                <p>ARF Enterprise delivers the same mathematically rigorous foundation but with <strong>production‑grade reliability</strong> and <strong>governance controls</strong> that meet the strictest compliance requirements.</p>
                <ul style="display: inline-block; text-align: left; margin: 1rem auto;">
                    <li>📊 **Persistent storage** – every decision logged and queryable</li>
                    <li>⚙️ **Advanced risk fusion** – combine conjugate, hyperprior, and HMC estimates</li>
                    <li>🛡️ **Semantic memory** – FAISS vector search for context‑aware policies</li>
                    <li>📈 **Real‑time dashboards** with Grafana & Prometheus integration</li>
                </ul>
            </div>

            <div style="text-align: center; margin-top: 2rem;">
                <a href="https://calendly.com/petter2025us/30min" target="_blank" style="background: #764ba2; color: white; padding: 12px 24px; text-decoration: none; border-radius: 8px; font-weight: bold; margin-right: 1rem;">📅 Book a Demo</a>
                <a href="mailto:petter2025us@outlook.com" style="background: #667eea; color: white; padding: 12px 24px; text-decoration: none; border-radius: 8px; font-weight: bold;">📧 Contact Sales</a>
            </div>
            """)
            gr.Markdown("### 📥 Export Decision History")
            export_btn = gr.DownloadButton("Download CSV", variant="primary")
            export_btn.click(
                fn=export_history_to_csv,
                outputs=[gr.File(label="decision_history.csv")]
            )

# ----------------------------------------------------------------------
# Launch
# ----------------------------------------------------------------------
if __name__ == "__main__":
    demo.queue()
    if PROMETHEUS_AVAILABLE and hasattr(demo, 'app') and demo.app:
        demo.app.add_api_route("/metrics", lambda: (generate_latest(), 200, {"Content-Type": CONTENT_TYPE_LATEST}), methods=["GET"])
    demo.launch(theme="soft", server_name="0.0.0.0", server_port=7860)