OBLITERATUS/obliteratus/ui_watchtower.py

"""Watchtower + One-Click Obliterate UI tabs for the OBLITERATUS Gradio app.

Exports a single function `build_watchtower_tabs()` that creates two tabs
inside the existing `gr.Tabs()` block:

    ⚡ One-Click  — Simple one-click obliteration with auto-populated dropdown
    🔭 Watchtower — Model scanner management and discovery dashboard

Usage in app.py:
    with gr.Tabs():
        # ... existing tabs ...
        from obliteratus.ui_watchtower import build_watchtower_tabs
        build_watchtower_tabs()
"""

from __future__ import annotations

import os
import threading
import time
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr


# ── Lazy imports to avoid circular deps and slow startup ──────────────

def _get_watchtower():
    from obliteratus.watchtower import get_watchtower
    return get_watchtower()


def _get_trending_choices() -> list[str]:
    """Get trending model IDs for dropdown, with graceful fallback."""
    try:
        wt = _get_watchtower()
        choices = wt.get_model_choices()
        if choices:
            return choices
    except Exception:
        pass
    # Fallback: common popular models
    return [
        "meta-llama/Llama-3.1-8B-Instruct",
        "Qwen/Qwen3-4B",
        "mistralai/Mistral-7B-Instruct-v0.3",
        "google/gemma-2-9b-it",
        "microsoft/Phi-3-mini-4k-instruct",
        "deepseek-ai/DeepSeek-V2-Lite-Chat",
        "tiiuae/falcon-7b-instruct",
        "NousResearch/Hermes-3-Llama-3.1-8B",
    ]


# ── Scheduler state ──────────────────────────────────────────────────

_scheduler_intervals = {
    "15 minutes": 15 * 60,
    "30 minutes": 30 * 60,
    "1 hour": 60 * 60,
    "6 hours": 6 * 60 * 60,
    "24 hours": 24 * 60 * 60,
}


# ── One-Click Obliterate handler ─────────────────────────────────────

def _run_one_click(
    model_id: str,
    max_iterations: int,
    target_refusal_pct: float,
    progress=gr.Progress(),
):
    """Generator that runs auto-obliteration, yielding (status, log, metrics).

    Designed as a Gradio generator function for streaming output.
    """
    if not model_id or not model_id.strip():
        yield (
            "⚠️ Please enter or select a model ID.",
            "Error: No model selected.",
            "",
            gr.update(interactive=False),
        )
        return

    model_id = model_id.strip()
    target_rate = target_refusal_pct / 100.0  # Convert % to fraction

    try:
        from obliteratus.auto_obliterate import AutoObliterator
    except ImportError as e:
        yield (
            f"⚠️ Import error: {e}",
            f"Failed to import AutoObliterator: {e}",
            "",
            gr.update(interactive=False),
        )
        return

    obliterator = AutoObliterator(
        model_id=model_id,
        max_iterations=int(max_iterations),
        target_refusal_rate=target_rate,
    )

    # Run the generator
    gen = obliterator.run()
    final_result = None
    try:
        while True:
            try:
                status, log, metrics = next(gen)
                # Show download button only when we have output
                has_output = (
                    obliterator._result.final_output_dir
                    and Path(obliterator._result.final_output_dir).exists()
                )
                yield (
                    status,
                    log,
                    metrics,
                    gr.update(
                        interactive=has_output,
                        visible=has_output,
                    ),
                )
            except StopIteration as e:
                final_result = e.value
                break
    except Exception as e:
        yield (
            f"⚠️ Error: {e}",
            f"Fatal error during auto-obliteration:\n{e}",
            "",
            gr.update(interactive=False),
        )
        return

    # Final yield with download button
    if final_result and final_result.final_output_dir:
        output_path = final_result.final_output_dir
        yield (
            "✅ Complete!" if final_result.success else "⚠️ Complete (target not fully met)",
            f"Auto-obliteration finished.\n"
            f"Output saved to: {output_path}\n\n"
            f"Final refusal rate: {final_result.final_refusal_rate or 'unknown'}\n"
            f"Total time: {final_result.total_time_seconds}s",
            obliterator._format_metrics(),
            gr.update(interactive=True, visible=True),
        )


def _download_result(model_id: str):
    """Return the path to the obliterated model for download."""
    if not model_id:
        return None
    safe_name = model_id.strip().replace("/", "_").replace("\\", "_")
    base = Path.home() / ".obliteratus" / "auto_obliterate" / safe_name

    # Find the latest iteration output
    if base.exists():
        iter_dirs = sorted(base.glob("iter_*"), reverse=True)
        for d in iter_dirs:
            if (d / "config.json").exists():
                return str(d)
    return None


# ── Watchtower handlers ──────────────────────────────────────────────

def _scan_now():
    """Run a watchtower scan and return updated UI components."""
    log_lines = []

    def on_log(msg):
        log_lines.append(msg)

    try:
        wt = _get_watchtower()
        new_models = wt.scan(on_log=on_log)
        stats = wt.get_stats()
        table_data = wt.format_table()

        status_html = _format_status_html(stats, len(new_models))
        log_text = "\n".join(log_lines)

        return status_html, table_data, log_text
    except Exception as e:
        return (
            f"<div style='color: #ff003c;'>Error: {e}</div>",
            [],
            f"Scan failed: {e}",
        )


def _get_watchtower_status():
    """Get current watchtower status for display."""
    try:
        wt = _get_watchtower()
        stats = wt.get_stats()
        return _format_status_html(stats)
    except Exception:
        return "<div style='color: #4a5568;'>Watchtower not initialized. Click 'Scan Now' to start.</div>"


def _get_watchtower_table():
    """Get current watchtower table data."""
    try:
        wt = _get_watchtower()
        return wt.format_table()
    except Exception:
        return []


def _get_obliteration_history():
    """Get history of past obliterations from watchtower."""
    try:
        wt = _get_watchtower()
        obliterated = wt.get_obliterated()
        if not obliterated:
            return "*No models obliterated yet via Watchtower.*"

        lines = [
            "| Model | Method | Refusal | Perplexity | Date |",
            "|-------|--------|---------|------------|------|",
        ]
        for m in obliterated:
            metrics = m.obliteration_metrics or {}
            ref = f"{metrics.get('refusal_rate', '—')}"
            ppl = f"{metrics.get('perplexity', '—')}"
            method = metrics.get("method", "—")
            date = m.last_updated[:16] if m.last_updated else "—"
            lines.append(f"| {m.model_id} | {method} | {ref} | {ppl} | {date} |")

        return "\n".join(lines)
    except Exception:
        return "*Error loading history.*"


def _set_scan_frequency(freq_label: str):
    """Update the scan scheduler frequency."""
    try:
        wt = _get_watchtower()
        interval = _scheduler_intervals.get(freq_label, 3600)
        wt.start_scheduler(interval=interval)
        return f"Scheduler set to scan every {freq_label}."
    except Exception as e:
        return f"Error: {e}"


def _toggle_auto_obliterate(enabled: bool):
    """Toggle auto-obliteration of newly discovered models."""
    try:
        wt = _get_watchtower()
        if enabled:
            def _auto_queue(model):
                wt.set_status(model.model_id, "queued")
                # In a full implementation, this would trigger auto-obliteration
            wt.on_new_model(_auto_queue)
            return "✅ Auto-obliterate enabled. New models will be queued automatically."
        else:
            wt._on_new_model_callbacks.clear()
            return "Auto-obliterate disabled."
    except Exception as e:
        return f"Error: {e}"


def _format_status_html(stats: dict, new_count: int = 0) -> str:
    """Format watchtower stats as styled HTML."""
    total = stats.get("total_tracked", 0)
    by_status = stats.get("by_status", {})
    last_scan = stats.get("last_scan", "Never")
    scan_count = stats.get("scan_count", 0)

    if last_scan and last_scan != "Never":
        try:
            dt = datetime.fromisoformat(last_scan)
            last_scan = dt.strftime("%Y-%m-%d %H:%M UTC")
        except Exception:
            pass

    obliterated = by_status.get("obliterated", 0)
    queued = by_status.get("queued", 0)
    new = by_status.get("new", 0)

    new_badge = ""
    if new_count > 0:
        new_badge = (
            f'<span style="background: #ff6600; color: #000; padding: 2px 8px; '
            f'border-radius: 4px; font-size: 0.85em; margin-left: 8px;">'
            f'+{new_count} NEW</span>'
        )

    return f"""
    <div style="
        background: linear-gradient(135deg, #0d0d14 0%, #1a1f2e 100%);
        border: 1px solid #00ff41;
        border-radius: 8px;
        padding: 16px 20px;
        font-family: 'Fira Code', monospace;
    ">
        <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 12px;">
            <span style="color: #00ff41; font-size: 1.1em; font-weight: bold;">
                🔭 WATCHTOWER STATUS {new_badge}
            </span>
            <span style="color: #4a5568; font-size: 0.85em;">
                Scan #{scan_count}
            </span>
        </div>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(140px, 1fr)); gap: 12px;">
            <div style="text-align: center;">
                <div style="color: #00ff41; font-size: 1.8em; font-weight: bold;">{total}</div>
                <div style="color: #4a5568; font-size: 0.8em;">TRACKED</div>
            </div>
            <div style="text-align: center;">
                <div style="color: #ff6600; font-size: 1.8em; font-weight: bold;">{new}</div>
                <div style="color: #4a5568; font-size: 0.8em;">NEW</div>
            </div>
            <div style="text-align: center;">
                <div style="color: #ffcc00; font-size: 1.8em; font-weight: bold;">{queued}</div>
                <div style="color: #4a5568; font-size: 0.8em;">QUEUED</div>
            </div>
            <div style="text-align: center;">
                <div style="color: #00cc33; font-size: 1.8em; font-weight: bold;">{obliterated}</div>
                <div style="color: #4a5568; font-size: 0.8em;">OBLITERATED</div>
            </div>
        </div>
        <div style="color: #4a5568; font-size: 0.8em; margin-top: 12px; text-align: right;">
            Last scan: {last_scan}
        </div>
    </div>
    """


def _refresh_one_click_dropdown():
    """Refresh the one-click dropdown with latest trending models."""
    choices = _get_trending_choices()
    return gr.update(choices=choices, value=choices[0] if choices else "")


# ── Build tabs ────────────────────────────────────────────────────────

def build_watchtower_tabs():
    """Build the ⚡ One-Click and 🔭 Watchtower tabs.

    Call this inside a `with gr.Tabs():` block in app.py.
    Creates and wires up all UI components.
    """

    # ══════════════════════════════════════════════════════════════════
    # ⚡ ONE-CLICK TAB
    # ══════════════════════════════════════════════════════════════════

    with gr.Tab("⚡ One-Click", id="one_click"):

        gr.Markdown(
            "### One-Click Obliterate\n"
            "Select a model, press the button, and let OBLITERATUS handle everything. "
            "Iteratively escalates methods until refusal is eliminated."
        )

        with gr.Row():
            with gr.Column(scale=3):
                oc_model_dd = gr.Dropdown(
                    choices=_get_trending_choices(),
                    value="",
                    label="Target Model",
                    info="Select from trending models or type any HuggingFace model ID",
                    allow_custom_value=True,
                    interactive=True,
                )
            with gr.Column(scale=1):
                oc_refresh_btn = gr.Button(
                    "🔄 Refresh",
                    variant="secondary",
                    size="sm",
                )

        with gr.Row():
            with gr.Column(scale=1):
                oc_max_iter = gr.Slider(
                    minimum=1,
                    maximum=5,
                    value=3,
                    step=1,
                    label="Max Iterations",
                    info="Number of escalating obliteration attempts",
                )
            with gr.Column(scale=1):
                oc_target_refusal = gr.Slider(
                    minimum=0,
                    maximum=10,
                    value=5,
                    step=1,
                    label="Target Refusal Rate (%)",
                    info="Stop when refusal rate drops below this",
                )

        oc_obliterate_btn = gr.Button(
            "⚡ O B L I T E R A T E ⚡",
            variant="primary",
            size="lg",
            elem_classes=["one-click-btn"],
        )

        oc_status = gr.Markdown(
            value="*Ready. Select a model and click OBLITERATE.*",
            label="Status",
        )

        with gr.Row():
            with gr.Column(scale=2):
                oc_log = gr.Textbox(
                    label="Live Log",
                    lines=20,
                    max_lines=40,
                    interactive=False,
                    elem_classes=["log-box"],
                )
            with gr.Column(scale=1):
                oc_metrics = gr.Markdown(
                    value="*Metrics will appear here after each iteration.*",
                    label="Iteration Metrics",
                )

        oc_download_btn = gr.Button(
            "📦 Download Obliterated Model",
            variant="secondary",
            visible=False,
            interactive=False,
        )

        # ── Wire events ──────────────────────────────────────────────

        oc_refresh_btn.click(
            fn=_refresh_one_click_dropdown,
            outputs=[oc_model_dd],
        )

        oc_obliterate_btn.click(
            fn=_run_one_click,
            inputs=[oc_model_dd, oc_max_iter, oc_target_refusal],
            outputs=[oc_status, oc_log, oc_metrics, oc_download_btn],
        )

    # ══════════════════════════════════════════════════════════════════
    # 🔭 WATCHTOWER TAB
    # ══════════════════════════════════════════════════════════════════

    with gr.Tab("🔭 Watchtower", id="watchtower"):

        gr.Markdown(
            "### Model Watchtower\n"
            "Continuously scans HuggingFace for new popular open-weight "
            "instruction-tuned models. Tracks discoveries and obliteration status."
        )

        # Status display
        wt_status_html = gr.HTML(
            value=_get_watchtower_status(),
            label="Status",
        )

        with gr.Row():
            with gr.Column(scale=1):
                wt_scan_btn = gr.Button(
                    "🔍 Scan Now",
                    variant="primary",
                    size="lg",
                )
            with gr.Column(scale=1):
                wt_freq_dd = gr.Dropdown(
                    choices=list(_scheduler_intervals.keys()),
                    value="1 hour",
                    label="Auto-Scan Frequency",
                    info="How often to check for new models",
                )
            with gr.Column(scale=1):
                wt_auto_toggle = gr.Checkbox(
                    value=False,
                    label="Auto-Obliterate New Models",
                    info="Automatically queue newly discovered models",
                )

        wt_freq_status = gr.Textbox(
            value="",
            label="Scheduler Status",
            interactive=False,
            visible=True,
            lines=1,
        )

        # Model discovery table
        gr.Markdown("#### Discovered Models")
        wt_table = gr.Dataframe(
            value=_get_watchtower_table(),
            headers=[
                "Model ID", "Org", "Size", "Downloads (7d)",
                "Likes", "License", "Discovered", "Status",
            ],
            datatype=["str"] * 8,
            interactive=False,
            wrap=True,
            row_count=(10, "dynamic"),
            col_count=(8, "fixed"),
        )

        # Scan log
        wt_scan_log = gr.Textbox(
            label="Scan Log",
            lines=8,
            max_lines=20,
            interactive=False,
            elem_classes=["log-box"],
        )

        # Obliteration history
        gr.Markdown("#### Obliteration History")
        wt_history = gr.Markdown(
            value=_get_obliteration_history(),
        )

        # ── Wire events ──────────────────────────────────────────────

        wt_scan_btn.click(
            fn=_scan_now,
            outputs=[wt_status_html, wt_table, wt_scan_log],
        )

        wt_freq_dd.change(
            fn=_set_scan_frequency,
            inputs=[wt_freq_dd],
            outputs=[wt_freq_status],
        )

        wt_auto_toggle.change(
            fn=_toggle_auto_obliterate,
            inputs=[wt_auto_toggle],
            outputs=[wt_freq_status],
        )


# ── Extra CSS for the one-click button ────────────────────────────────
# This can be appended to the main CSS variable in app.py.

ONE_CLICK_CSS = """
/* ── One-Click Obliterate button ── */
.one-click-btn button {
    background: linear-gradient(135deg, #ff4400 0%, #ff6600 30%, #ff0044 70%, #cc0033 100%) !important;
    color: #fff !important;
    font-size: 1.4em !important;
    font-weight: bold !important;
    letter-spacing: 0.15em !important;
    border: 2px solid #ff4400 !important;
    padding: 16px 32px !important;
    text-shadow: 0 0 10px rgba(255,68,0,0.5) !important;
    box-shadow: 0 0 20px rgba(255,68,0,0.3), inset 0 0 20px rgba(255,255,255,0.05) !important;
    transition: all 0.3s ease !important;
    min-height: 60px !important;
}
.one-click-btn button:hover {
    box-shadow: 0 0 40px rgba(255,68,0,0.6), inset 0 0 30px rgba(255,255,255,0.1) !important;
    transform: scale(1.02) !important;
    border-color: #ff6600 !important;
}
.one-click-btn button:active {
    transform: scale(0.98) !important;
}
"""