"""
InsAIts Live TUI Dashboard
==========================
Runs in a VS Code split terminal panel. Watches the InsAIts audit log
in real time and displays live metrics, anomaly feed, and agent status.

Usage:
    python -m insa_its.dashboard.tui
    python -m insa_its.dashboard.tui --log .insaits_audit_session.jsonl
    python -m insa_its.dashboard.tui --demo

Split terminal in VS Code: Ctrl+Shift+5 (or click the split icon in terminal)

Complete flow:
    Claude Code runs        -> PostToolUse hook triggers insaits_hook.py
    insaits_hook.py         -> inspects output, writes to .insaits_audit_session.jsonl
    insaits_tui.py          -> watches the JSONL file, updates dashboard every 0.8s

Author: Cristi Bogdan -- Steddy Nova SRL / YuyAI
"""

import json
import time
import argparse
import threading
from pathlib import Path
from datetime import datetime
from collections import defaultdict, deque
from typing import Optional

from textual.app import App, ComposeResult
from textual.containers import Container, Horizontal, Vertical, ScrollableContainer
from textual.widgets import (
    Header, Footer, Static, DataTable, Label, ProgressBar, Sparkline
)
from textual.reactive import reactive
from textual.timer import Timer
from textual import events
from rich.text import Text
from rich.panel import Panel
from rich.table import Table
from rich.console import Console
from rich import box


# --- Colours -----------------------------------------------------------------
NAVY    = "#0D1B2A"
BLUE    = "#1D4ED8"
CYAN    = "#06B6D4"
GREEN   = "#10B981"
AMBER   = "#F59E0B"
RED     = "#EF4444"
PURPLE  = "#8B5CF6"
GRAY    = "#6B7280"
WHITE   = "#F9FAFB"

SEVERITY_COLOR = {
    "CRITICAL": RED,
    "HIGH":     AMBER,
    "MEDIUM":   CYAN,
    "LOW":      GREEN,
    "INFO":     GRAY,
}

ANOMALY_ICONS = {
    "blank_response":           "O",
    "waiting_for_description":  "||",
    "incomplete_code":          "!",
    "context_collapse":         "~>",
    "semantic_drift":           "~",
    "repetition_loop":          "<>",
    "truncated_output":         "|>",
    "confidence_collapse":      "v",
    "tool_poisoning":           "X",
    "prompt_injection":         ">>",
    "hallucination_chain":      "oo",
    "jargon_drift":             "?",
    "credential_exposure":      "**",
    "LLM_FINGERPRINT_MISMATCH": "!F",
    "TOOL_DESCRIPTION_DIVERGENCE": "TD",
    "BEHAVIORAL_FINGERPRINT_CHANGE": "BF",
    "CREDENTIAL_EXPOSURE":      "CE",
    "INFORMATION_FLOW_VIOLATION": "IF",
    "TOOL_CALL_ANOMALY":        "TC",
}


# --- Data Layer: watches the audit log file ----------------------------------

class AuditWatcher:
    """
    Tails the InsAIts JSONL audit log and maintains in-memory state.
    Thread-safe -- runs in background thread, TUI reads from it.
    """

    def __init__(self, log_path: str):
        self.log_path    = Path(log_path)
        self._lock       = threading.Lock()
        self._entries: deque   = deque(maxlen=500)
        self._anomalies: deque = deque(maxlen=200)
        self._agents: dict     = {}
        self._stats      = {
            "total_messages":  0,
            "total_anomalies": 0,
            "resolved":        0,
            "critical":        0,
            "high":            0,
            "by_type":         defaultdict(int),
            "anomaly_rate_history": deque(maxlen=60),
        }
        self._last_pos   = 0
        self._running    = True
        self._thread     = threading.Thread(target=self._tail, daemon=True)

    def start(self):
        self._thread.start()

    def stop(self):
        self._running = False

    def get_stats(self) -> dict:
        with self._lock:
            return dict(self._stats)

    def get_recent_anomalies(self, n: int = 20) -> list:
        with self._lock:
            return list(self._anomalies)[-n:]

    def get_agents(self) -> dict:
        with self._lock:
            return dict(self._agents)

    def get_rate_history(self) -> list:
        with self._lock:
            return list(self._stats["anomaly_rate_history"])

    def _tail(self):
        """Background thread: tail the log file and parse new entries."""
        while self._running:
            try:
                if self.log_path.exists():
                    with open(self.log_path, "r", encoding="utf-8") as f:
                        f.seek(self._last_pos)
                        while True:
                            line = f.readline()
                            if not line:
                                break
                            line = line.strip()
                            if line:
                                try:
                                    entry = json.loads(line)
                                    self._process(entry)
                                except json.JSONDecodeError:
                                    pass
                        self._last_pos = f.tell()
            except (IOError, OSError):
                pass
            time.sleep(0.5)

    def _process(self, entry: dict):
        with self._lock:
            self._entries.append(entry)
            self._stats["total_messages"] += 1

            anomalies = entry.get("anomalies", [])
            if isinstance(anomalies, list):
                for a in anomalies:
                    self._stats["total_anomalies"] += 1
                    atype = a if isinstance(a, str) else a.get("type", "unknown")
                    self._stats["by_type"][atype] += 1
                    sev = a.get("severity", "MEDIUM") if isinstance(a, dict) else "MEDIUM"
                    if sev == "CRITICAL":
                        self._stats["critical"] += 1
                    elif sev == "HIGH":
                        self._stats["high"] += 1
                    self._anomalies.append({
                        "ts":       entry.get("ts", datetime.utcnow().isoformat()),
                        "type":     atype,
                        "severity": sev,
                        "model":    entry.get("model", "unknown"),
                        "resolved": entry.get("resolved", False),
                    })

            if entry.get("resolved"):
                self._stats["resolved"] += 1

            # Track agents
            model = entry.get("model", "unknown")
            if model not in self._agents:
                self._agents[model] = {"messages": 0, "anomalies": 0, "status": "HEALTHY"}
            self._agents[model]["messages"] += 1
            self._agents[model]["anomalies"] += len(anomalies) if isinstance(anomalies, list) else 0
            rate = self._agents[model]["anomalies"] / max(self._agents[model]["messages"], 1)
            self._agents[model]["status"] = (
                "CRITICAL" if rate > 0.5 else
                "WARNING"  if rate > 0.2 else
                "HEALTHY"
            )

            # Rolling anomaly rate
            total = self._stats["total_messages"]
            anom  = self._stats["total_anomalies"]
            self._stats["anomaly_rate_history"].append(
                int((anom / max(total, 1)) * 100)
            )


# --- Demo data generator (when no real log exists) ---------------------------

class DemoWatcher(AuditWatcher):
    """Generates synthetic data for demo/development mode."""

    DEMO_AGENTS  = ["OrderBot", "InventoryBot", "ShippingBot", "RefundBot"]
    DEMO_MODELS  = ["gpt-4o", "gpt-4o-mini", "claude-3.5"]
    DEMO_ANOMALY_TYPES = [
        "semantic_drift", "hallucination_chain", "context_collapse",
        "incomplete_code", "LLM_FINGERPRINT_MISMATCH", "prompt_injection",
        "jargon_drift", "blank_response", "TOOL_DESCRIPTION_DIVERGENCE",
        "CREDENTIAL_EXPOSURE", "INFORMATION_FLOW_VIOLATION",
    ]

    def __init__(self):
        super().__init__("demo_nonexistent.jsonl")
        self._demo_tick = 0

    def _tail(self):
        import random
        while self._running:
            self._demo_tick += 1
            # Generate 1-3 messages per tick
            for _ in range(random.randint(1, 3)):
                model   = random.choice(self.DEMO_MODELS)
                anom    = []
                if random.random() < 0.45:
                    atype = random.choice(self.DEMO_ANOMALY_TYPES)
                    sev   = random.choice(["CRITICAL","HIGH","MEDIUM","MEDIUM","LOW"])
                    anom  = [{"type": atype, "severity": sev}]
                entry = {
                    "ts":       datetime.utcnow().isoformat(),
                    "model":    model,
                    "msg_id":   f"msg_{self._demo_tick}",
                    "retries":  1 if anom else 0,
                    "resolved": bool(anom) and random.random() > 0.3,
                    "anomalies": anom,
                }
                self._process(entry)
            time.sleep(1.2)


# --- Widgets -----------------------------------------------------------------

class MetricBox(Static):
    """A single KPI metric box."""

    def __init__(self, label: str, value: str, color: str = WHITE, **kwargs):
        super().__init__(**kwargs)
        self._label = label
        self._value = value
        self._color = color

    def render(self) -> str:
        return (
            f"[bold {self._color}]{self._value}[/bold {self._color}]\n"
            f"[{GRAY}]{self._label}[/{GRAY}]"
        )

    def update_value(self, value: str, color: Optional[str] = None):
        self._value = value
        if color:
            self._color = color
        self.refresh()


class AnomalyFeed(Static):
    """Live scrolling anomaly event feed."""

    MAX_LINES = 18

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._lines: deque = deque(maxlen=self.MAX_LINES)

    def push(self, anomaly: dict):
        ts    = anomaly.get("ts", "")[-8:][:8]  # HH:MM:SS
        atype = anomaly.get("type", "unknown")
        sev   = anomaly.get("severity", "MEDIUM")
        model = anomaly.get("model", "?")[:12]
        icon  = ANOMALY_ICONS.get(atype, "*")
        color = SEVERITY_COLOR.get(sev, GRAY)
        res   = "+" if anomaly.get("resolved") else "x"
        line = (
            f"[{GRAY}]{ts}[/{GRAY}] "
            f"[bold {color}]{icon} {sev[:4]}[/bold {color}] "
            f"[{WHITE}]{atype.replace('_',' ')[:22]:<22}[/{WHITE}] "
            f"[{CYAN}]{model}[/{CYAN}] "
            f"[{GREEN if res == '+' else RED}]{res}[/{GREEN if res == '+' else RED}]"
        )
        self._lines.append(line)
        self.refresh()

    def render(self) -> str:
        if not self._lines:
            return f"[{GRAY}]  Waiting for anomalies...[/{GRAY}]"
        return "\n".join(self._lines)


class AgentTable(Static):
    """Live agent status table."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._agents: dict = {}

    def update_agents(self, agents: dict):
        self._agents = agents
        self.refresh()

    def render(self) -> str:
        if not self._agents:
            return f"[{GRAY}]  No agents registered yet...[/{GRAY}]"

        header = (
            f"[bold {GRAY}]{'AGENT':<16} {'MSGS':>6} {'ANOMS':>6} "
            f"{'RATE':>6}  {'STATUS'}[/bold {GRAY}]"
        )
        lines = [header, f"[{GRAY}]{'---'*17}[/{GRAY}]"]

        for name, data in self._agents.items():
            msgs  = data["messages"]
            anom  = data["anomalies"]
            rate  = anom / max(msgs, 1)
            st    = data["status"]
            sc    = RED if st == "CRITICAL" else AMBER if st == "WARNING" else GREEN
            bar_len = int(rate * 10)
            bar   = "#" * bar_len + "." * (10 - bar_len)
            lines.append(
                f"[{CYAN}]{name[:16]:<16}[/{CYAN}] "
                f"[{WHITE}]{msgs:>6}[/{WHITE}] "
                f"[{AMBER}]{anom:>6}[/{AMBER}] "
                f"[{GRAY}]{rate:>5.0%}[/{GRAY}]  "
                f"[bold {sc}]{st}[/bold {sc}]"
            )
        return "\n".join(lines)


class SparklineWidget(Static):
    """Anomaly rate over time as ASCII sparkline."""

    CHARS = " .:-=+*#@"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._data: list = []

    def update_data(self, data: list):
        self._data = data[-50:] if data else []
        self.refresh()

    def render(self) -> str:
        if not self._data:
            return f"[{GRAY}]  No data yet[/{GRAY}]"
        mx = max(self._data) if self._data else 1
        spark = "".join(
            self.CHARS[min(int(v / max(mx, 1) * 8), 8)]
            for v in self._data
        )
        current = self._data[-1] if self._data else 0
        color = RED if current > 50 else AMBER if current > 25 else GREEN
        return (
            f"[{GRAY}]Anomaly rate (%)[/{GRAY}]  "
            f"[bold {color}]{current:3d}%[/bold {color}]  "
            f"[{CYAN}]{spark}[/{CYAN}]"
        )


class TypeBreakdown(Static):
    """Bar chart of anomaly types."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._data: dict = {}

    def update_data(self, data: dict):
        self._data = data
        self.refresh()

    def render(self) -> str:
        if not self._data:
            return f"[{GRAY}]  No anomalies yet[/{GRAY}]"

        total  = sum(self._data.values())
        sorted_data = sorted(self._data.items(), key=lambda x: -x[1])[:8]
        lines  = []
        for atype, count in sorted_data:
            pct  = count / max(total, 1)
            bar_len = int(pct * 20)
            bar  = "#" * bar_len
            icon = ANOMALY_ICONS.get(atype, "*")
            lines.append(
                f"[{CYAN}]{icon}[/{CYAN}] "
                f"[{WHITE}]{atype.replace('_',' ')[:22]:<22}[/{WHITE}] "
                f"[{AMBER}]{bar:<20}[/{AMBER}] "
                f"[{GRAY}]{count:>4}[/{GRAY}]"
            )
        return "\n".join(lines) if lines else f"[{GRAY}]  No anomalies yet[/{GRAY}]"


# --- Main App ----------------------------------------------------------------

CSS = f"""
Screen {{
    background: {NAVY};
}}

#header-bar {{
    height: 3;
    background: {NAVY};
    border-bottom: solid {BLUE};
    padding: 0 2;
}}

#title {{
    color: {CYAN};
    text-style: bold;
    width: 1fr;
    padding: 1 0;
}}

#subtitle {{
    color: {GRAY};
    width: 1fr;
    text-align: right;
    padding: 1 0;
}}

#metrics-row {{
    height: 5;
    margin: 1 1 0 1;
}}

MetricBox {{
    border: solid {BLUE};
    background: #0a1628;
    width: 1fr;
    height: 5;
    padding: 1 2;
    text-align: center;
    margin: 0 1;
}}

#sparkline-row {{
    height: 3;
    margin: 0 2;
    padding: 1 0;
}}

SparklineWidget {{
    width: 1fr;
    color: {GRAY};
}}

#main-row {{
    height: 1fr;
    margin: 0 1;
}}

#feed-panel {{
    width: 2fr;
    border: solid {BLUE};
    background: #080f1a;
    margin: 0 1;
    padding: 1;
}}

#right-col {{
    width: 1fr;
    margin: 0 1;
}}

#agent-panel {{
    border: solid {BLUE};
    background: #080f1a;
    height: 1fr;
    padding: 1;
    margin-bottom: 1;
}}

#breakdown-panel {{
    border: solid {BLUE};
    background: #080f1a;
    height: 1fr;
    padding: 1;
}}

.panel-label {{
    color: {BLUE};
    text-style: bold;
    margin-bottom: 1;
}}

Footer {{
    background: {NAVY};
    border-top: solid {BLUE};
    color: {GRAY};
}}
"""


class InsAItsDashboard(App):
    """InsAIts Live TUI Dashboard -- runs in VS Code terminal panel."""

    CSS = CSS
    BINDINGS = [
        ("q",     "quit",        "Quit"),
        ("c",     "clear",       "Clear feed"),
        ("r",     "reset",       "Reset stats"),
        ("d",     "toggle_demo", "Toggle demo mode"),
    ]

    # Reactive state
    total_messages  = reactive(0)
    total_anomalies = reactive(0)
    anomaly_rate    = reactive(0.0)
    status          = reactive("HEALTHY")

    def __init__(self, log_path: str = None, demo: bool = False):
        super().__init__()
        self._demo = demo
        self._log_path = log_path or self._find_log()
        self._prev_anomaly_count = 0
        self._watcher: Optional[AuditWatcher] = None
        self._start_time = time.time()

    def _find_log(self) -> str:
        """Auto-detect the most recent InsAIts audit log."""
        candidates = [
            ".insaits_audit_session.jsonl",
            *Path(".").glob(".insaits_audit_*.jsonl"),
            *Path.home().glob(".insaits/*.jsonl"),
        ]
        for c in candidates:
            if Path(c).exists():
                return str(c)
        return ".insaits_audit_session.jsonl"

    def compose(self) -> ComposeResult:
        yield Static(
            f"[bold {CYAN}]InsAIts[/bold {CYAN}]  "
            f"[{GRAY}]Runtime AI Security Monitor[/{GRAY}]  "
            f"[{BLUE}]v3.1[/{BLUE}]",
            id="title"
        )

        with Horizontal(id="metrics-row"):
            yield MetricBox("MESSAGES",  "0",       WHITE,   id="m-messages")
            yield MetricBox("ANOMALIES", "0",        AMBER,  id="m-anomalies")
            yield MetricBox("RATE",      "0%",       GREEN,  id="m-rate")
            yield MetricBox("RESOLVED",  "0",        GREEN,  id="m-resolved")
            yield MetricBox("CRITICAL",  "0",        RED,    id="m-critical")
            yield MetricBox("STATUS",    "HEALTHY",  GREEN,  id="m-status")

        yield SparklineWidget(id="sparkline")

        with Horizontal(id="main-row"):
            with Vertical(id="feed-panel"):
                yield Label(
                    f"[bold {BLUE}]>>  LIVE ANOMALY FEED[/bold {BLUE}]",
                    classes="panel-label"
                )
                yield AnomalyFeed(id="anomaly-feed")

            with Vertical(id="right-col"):
                with Vertical(id="agent-panel"):
                    yield Label(
                        f"[bold {BLUE}]<>  AGENT STATUS[/bold {BLUE}]",
                        classes="panel-label"
                    )
                    yield AgentTable(id="agent-table")

                with Vertical(id="breakdown-panel"):
                    yield Label(
                        f"[bold {BLUE}]##  ANOMALY TYPES[/bold {BLUE}]",
                        classes="panel-label"
                    )
                    yield TypeBreakdown(id="type-breakdown")

        yield Footer()

    def on_mount(self) -> None:
        if self._demo:
            self._watcher = DemoWatcher()
        else:
            self._watcher = AuditWatcher(self._log_path)
        self._watcher.start()
        self.set_interval(0.8, self._refresh_ui)

    def _refresh_ui(self) -> None:
        stats   = self._watcher.get_stats()
        agents  = self._watcher.get_agents()
        history = self._watcher.get_rate_history()

        # Update metrics
        msgs  = stats["total_messages"]
        anoms = stats["total_anomalies"]
        rate  = anoms / max(msgs, 1)
        crit  = stats["critical"]
        res   = stats["resolved"]

        status_color = RED if crit > 0 and rate > 0.3 else AMBER if rate > 0.1 else GREEN
        status_text  = "CRITICAL" if rate > 0.3 and crit > 0 else "WARNING" if rate > 0.1 else "HEALTHY"

        self.query_one("#m-messages",  MetricBox).update_value(str(msgs))
        self.query_one("#m-anomalies", MetricBox).update_value(str(anoms), AMBER if anoms > 0 else GREEN)
        self.query_one("#m-rate",      MetricBox).update_value(f"{rate:.0%}", RED if rate > 0.3 else AMBER if rate > 0.1 else GREEN)
        self.query_one("#m-resolved",  MetricBox).update_value(str(res), GREEN)
        self.query_one("#m-critical",  MetricBox).update_value(str(crit), RED if crit > 0 else GREEN)
        self.query_one("#m-status",    MetricBox).update_value(status_text, status_color)

        # Push new anomalies to feed
        feed  = self.query_one("#anomaly-feed", AnomalyFeed)
        new_a = self._watcher.get_recent_anomalies(50)
        if len(new_a) > self._prev_anomaly_count:
            for anomaly in new_a[self._prev_anomaly_count:]:
                feed.push(anomaly)
            self._prev_anomaly_count = len(new_a)

        # Update sparkline
        self.query_one("#sparkline",      SparklineWidget).update_data(history)
        self.query_one("#agent-table",    AgentTable).update_agents(agents)
        self.query_one("#type-breakdown", TypeBreakdown).update_data(dict(stats["by_type"]))

    def action_clear(self) -> None:
        self.query_one("#anomaly-feed", AnomalyFeed)._lines.clear()
        self.query_one("#anomaly-feed", AnomalyFeed).refresh()

    def action_reset(self) -> None:
        self._prev_anomaly_count = 0
        self.action_clear()

    def action_toggle_demo(self) -> None:
        if self._watcher:
            self._watcher.stop()
        self._demo = not self._demo
        self._prev_anomaly_count = 0
        if self._demo:
            self._watcher = DemoWatcher()
        else:
            self._watcher = AuditWatcher(self._log_path)
        self._watcher.start()
        self.notify(
            f"{'Demo' if self._demo else 'Live'} mode activated",
            severity="information"
        )

    def on_unmount(self) -> None:
        if self._watcher:
            self._watcher.stop()


# --- Entry point -------------------------------------------------------------

def main():
    """CLI entry point for InsAIts TUI dashboard."""
    parser = argparse.ArgumentParser(description="InsAIts Live TUI Dashboard")
    parser.add_argument("--log",  default=None,  help="Path to InsAIts audit JSONL log")
    parser.add_argument("--demo", action="store_true", help="Run in demo mode with synthetic data")
    args = parser.parse_args()

    app = InsAItsDashboard(log_path=args.log, demo=args.demo)
    app.run()


if __name__ == "__main__":
    main()
