bioops/simulation_ui/callbacks.py

"""Gradio callback functions for the BioOps Twin dashboard.

Handles chat interaction (via BioOpsAgent), telemetry DataFrame
generation, state display formatting, health score calculation,
shadow mode toggling, and audit log retrieval.
"""

from __future__ import annotations

import json
import logging
from typing import Any

import pandas as pd

from bioops.llm_engine.client import BioOpsAgent
from bioops.simulation_core.simulator import CentrifugeSimulator
from bioops.simulation_core.state_machine import MachineState
from bioops.security.sanitizer import sanitize_input, sanitize_output
from bioops.security.audit_logger import read_logs
from bioops.industrial_edge.mqtt_client import mqtt_edge

logger = logging.getLogger("bioops_twin.callbacks")

# Module-level references (injected at build time)
_simulator: CentrifugeSimulator | None = None
_agent: BioOpsAgent | None = None


def bind(simulator: CentrifugeSimulator, agent: BioOpsAgent) -> None:
    """Bind the simulator and agent instances for callback use.

    Args:
        simulator: The centrifuge simulator instance.
        agent: The BioOps cognitive agent instance.
    """
    global _simulator, _agent  # noqa: PLW0603
    _simulator = simulator
    _agent = agent


def _sim() -> CentrifugeSimulator:
    assert _simulator is not None, "Callbacks not initialised. Call bind()."
    return _simulator


def _agt() -> BioOpsAgent:
    assert _agent is not None, "Callbacks not initialised. Call bind()."
    return _agent


# ---------------------------------------------------------------------------
# Chat callback
# ---------------------------------------------------------------------------

def chat_respond(
    user_message: str,
    chat_history: list[dict[str, str]],
) -> tuple[list[dict[str, str]], str]:
    """Process operator messages through the BioOps Agent.

    Handles sanitisation (security layer), agent processing (LLM or
    mock), and automatic sensor alert injection when the simulator
    detects a safety violation (Rule 27).

    Uses Gradio 6 ``messages`` format (OpenAI-style dicts).

    Args:
        user_message: Raw text from the operator.
        chat_history: Gradio messages-format chat history.

    Returns:
        Updated chat history and empty string to clear the textbox.
    """
    chat_history = chat_history or []

    # Sanitise input (Veea stub — Rule 10)
    clean_input = sanitize_input(user_message)

    # Process through agent
    raw_response = _agt().process_message(clean_input)

    # Sanitise output
    response = sanitize_output(raw_response)
    chat_history.append({"role": "user", "content": clean_input})
    chat_history.append({"role": "assistant", "content": response})

    # Check for sensor alert generated by the physics tick
    # (feedback loop — Rule 27)
    if _sim().last_alert is not None:
        alert_msg = _agt().format_sensor_alert(_sim().last_alert)
        chat_history.append({"role": "assistant", "content": alert_msg})
        _sim().last_alert = None  # Consume the alert

    return chat_history, ""


# ---------------------------------------------------------------------------
# Shadow Mode toggle
# ---------------------------------------------------------------------------

def toggle_shadow_mode(enabled: bool) -> str:
    """Toggle Shadow Mode for the agent.

    Args:
        enabled: Whether shadow mode should be active.

    Returns:
        A display string showing the current mode.
    """
    _agt().shadow_mode = enabled
    mode = "🛡️ **SHADOW MODE** — Commands require operator confirmation" if enabled else "⚡ **AUTO-RUN** — Commands execute immediately"
    logger.info("Shadow mode set to %s", enabled)
    return mode


# ---------------------------------------------------------------------------
# Telemetry callbacks
# ---------------------------------------------------------------------------

def get_telemetry_dataframe() -> pd.DataFrame:
    """Build a pandas DataFrame from the simulator's telemetry log.

    Advances the simulation by one tick on each call (driven by
    ``gr.Timer``).

    Returns:
        DataFrame with columns ``time_s``, ``RPM``, ``Vibration (g)``,
        ``Z-Score``.
    """
    sim = _sim()
    sim.tick()

    if not sim.telemetry_log:
        return pd.DataFrame(
            {"time_s": [0], "RPM": [0], "Vibration": [0.0], "Z-Score": [0.0]},
        )

    t0: float = sim.telemetry_log[0].timestamp
    rows: list[dict[str, float]] = [
        {
            "time_s": round(s.timestamp - t0, 1),
            "RPM": float(s.rpm),
            "Vibration": s.vibration_rms_g,
            "Z-Score": s.z_score,
        }
        for s in sim.telemetry_log
    ]
    return pd.DataFrame(rows)


def get_dashboard_updates() -> tuple[str, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Retrieve all dashboard telemetry values in a single call.
    
    This bundles the requests to prevent hitting the Hugging Face Space
    Rate Limits (429) that occur when individual components poll via `every=`.
    """
    df = get_telemetry_dataframe()
    return (
        get_state_display(),
        df,
        df,
        df,
        get_audit_log_dataframe(),
    )


# ---------------------------------------------------------------------------
# Health Score
# ---------------------------------------------------------------------------

def get_health_score() -> float:
    """Calculate a 0-100 Health Score based on recent vibration history.

    Uses a weighted assessment of vibration levels relative to the
    critical threshold. 100 = perfect, 0 = imminent failure.

    Returns:
        Health score as a float [0, 100].
    """
    sim = _sim()
    from bioops.simulation_core.simulator import CRITICAL_VIBRATION_G

    if not sim.telemetry_log:
        return 100.0

    recent = sim.telemetry_log[-20:]
    avg_vib = sum(s.vibration_rms_g for s in recent) / len(recent)
    max_vib = max(s.vibration_rms_g for s in recent)
    anomaly_count = sum(1 for s in recent if s.anomaly)

    # Weighted score
    vib_ratio = min(avg_vib / CRITICAL_VIBRATION_G, 1.0)
    peak_ratio = min(max_vib / CRITICAL_VIBRATION_G, 1.0)
    anomaly_penalty = min(anomaly_count * 10, 30)

    score = 100.0 * (1.0 - 0.5 * vib_ratio - 0.3 * peak_ratio) - anomaly_penalty

    if sim.state == MachineState.EMERGENCY_STOP:
        score = min(score, 10.0)
    elif sim.state == MachineState.ERROR:
        score = min(score, 25.0)

    return max(0.0, min(100.0, score))


# ---------------------------------------------------------------------------
# State display (enhanced with Health Score, MQTT status, anomaly)
# ---------------------------------------------------------------------------

def get_state_display() -> str:
    """Return a formatted string showing the current machine state."""
    sim = _sim()
    state_icons: dict[MachineState, str] = {
        MachineState.STANDBY: "🟢",
        MachineState.SPINNING: "🔵",
        MachineState.ERROR: "🔴",
        MachineState.EMERGENCY_STOP: "🛑",
    }
    icon = state_icons.get(sim.state, "⚪")
    mode_tag = "🟢 LIVE" if _agt().is_live else "🟡 MOCK"
    shadow_tag = "🛡️ SHADOW" if _agt().shadow_mode else "⚡ AUTO"

    # Health Score
    health = get_health_score()
    if health >= 80:
        health_icon, health_color = "💚", "green"
    elif health >= 50:
        health_icon, health_color = "💛", "orange"
    else:
        health_icon, health_color = "❤️", "red"

    # MQTT status
    mqtt_status = "🟢 Connected" if mqtt_edge.connected else "🔴 Disconnected"

    # Latest Z-Score
    z_score = sim.telemetry_log[-1].z_score if sim.telemetry_log else 0.0
    anomaly_flag = " ⚠️ **ANOMALY DETECTED**" if (sim.telemetry_log and sim.telemetry_log[-1].anomaly) else ""

    # RCF (Relative Centrifugal Force)
    rcf_val = sim.rcf
    rcf_display = f"{rcf_val:,.0f} ×g" if rcf_val >= 1 else "— ×g"

    # Session uptime
    uptime = sim.uptime_seconds
    mins, secs = divmod(int(uptime), 60)
    hrs, mins = divmod(mins, 60)
    uptime_str = f"{hrs:02d}:{mins:02d}:{secs:02d}"

    return (
        f"{icon} **{sim.state.value}** — Agent: {mode_tag} · {shadow_tag}\n\n"
        f"**RPM:** {sim.current_rpm} / {sim.target_rpm} · **RCF:** {rcf_display}  \n"
        f"**Vibration:** {sim.vibration_rms_g:.4f} g · Z-Score: {z_score:.1f}{anomaly_flag}  \n"
        f"**Device:** {sim.device_id} · ⏱️ Uptime: {uptime_str}  \n"
        f"{health_icon} **Health Score:** {health:.0f}%  \n"
        f"📡 **MQTT Edge:** {mqtt_status}"
    )


# ---------------------------------------------------------------------------
# Audit Log table
# ---------------------------------------------------------------------------

def get_audit_log_dataframe() -> pd.DataFrame:
    """Return the latest audit log entries as a DataFrame for the UI.

    Returns:
        DataFrame with columns: Time, Level, Event, Source, Details.
    """
    logs = read_logs(limit=30)
    if not logs:
        return pd.DataFrame(
            {"Time": [], "Level": [], "Event": [], "Source": [], "Details": []}
        )

    rows = []
    for entry in logs:
        # Format details as readable key=value pairs
        details: dict[str, Any] = entry.get("details", {})
        detail_parts: list[str] = []
        for k, v in details.items():
            if isinstance(v, dict):
                continue  # skip nested dicts for readability
            detail_parts.append(f"{k}={v}")
        detail_str = " · ".join(detail_parts) if detail_parts else "—"
        if len(detail_str) > 80:
            detail_str = detail_str[:77] + "…"

        # Format timestamp to just time (HH:MM:SS)
        ts = entry.get("timestamp_iso", "")
        short_ts = ts[11:19] if len(ts) > 19 else ts

        rows.append({
            "Time": short_ts,
            "Level": entry.get("level", ""),
            "Event": entry.get("event_type", ""),
            "Source": entry.get("source", ""),
            "Details": detail_str,
        })
    return pd.DataFrame(rows)