#!/usr/bin/env python3
from __future__ import annotations

"""
Minimal analytics script for BloomOS task logs.

This expects a JSONL log file produced by BloomOSClient(log_path=...),
where each line looks roughly like:

    {
      "timestamp": 1732212345.123,
      "timestamp_iso": "2025-11-21T17:38:53",
      "request": {...},
      "result": {
        "task_id": "...",
        "status": "ok",
        "gate_used": "math_bias",
        "anchors_used": ["math"],
        "logits_mean": 0.272,
        "logits_std": 0.05,
        "receipt": {...}
      }
    }

Usage:
    python3 bloomos_dashboard.py path/to/log.jsonl
"""

import json
import sys
from pathlib import Path
from typing import Any, Dict, List

try:
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


def load_log(path: Path) -> List[Dict[str, Any]]:
    entries: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entries.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return entries


def print_summary(entries: List[Dict[str, Any]]) -> None:
    """Print text summary when matplotlib is not available."""
    by_gate: Dict[str, Dict[str, List[float]]] = {}

    for e in entries:
        res = e.get("result") or {}
        gate = res.get("gate_used") or "unknown"
        lm = res.get("logits_mean")
        ls = res.get("logits_std")
        if lm is None:
            continue
        if gate not in by_gate:
            by_gate[gate] = {"logits_mean": [], "logits_std": []}
        by_gate[gate]["logits_mean"].append(float(lm))
        if ls is not None:
            by_gate[gate]["logits_std"].append(float(ls))

    print("\n=== BloomOS Task Log Summary ===\n")
    print(f"Total entries: {len(entries)}")
    print()

    for gate, series in sorted(by_gate.items()):
        means = series["logits_mean"]
        stds = series["logits_std"]
        if not means:
            continue
        avg_mean = sum(means) / len(means)
        min_mean = min(means)
        max_mean = max(means)
        avg_std = sum(stds) / len(stds) if stds else 0.0

        print(f"Gate: {gate}")
        print(f"  Tasks: {len(means)}")
        print(f"  logits_mean: avg={avg_mean:.4f}, min={min_mean:.4f}, max={max_mean:.4f}")
        if stds:
            print(f"  logits_std:  avg={avg_std:.4f}")
        print()


def plot_dashboard(entries: List[Dict[str, Any]]) -> None:
    """Plot logits_mean distributions per gate."""
    by_gate: Dict[str, Dict[str, List[float]]] = {}

    for e in entries:
        res = e.get("result") or {}
        gate = res.get("gate_used") or "unknown"
        lm = res.get("logits_mean")
        ls = res.get("logits_std")
        if lm is None:
            continue
        if gate not in by_gate:
            by_gate[gate] = {"logits_mean": [], "logits_std": []}
        by_gate[gate]["logits_mean"].append(float(lm))
        if ls is not None:
            by_gate[gate]["logits_std"].append(float(ls))

    # Plot logits_mean distributions per gate (simple line plot)
    plt.figure(figsize=(10, 6))
    for gate, series in by_gate.items():
        means = series["logits_mean"]
        if not means:
            continue
        plt.plot(means, label=gate, marker='o', markersize=4)

    plt.xlabel("Task index (per gate, in log order)")
    plt.ylabel("logits_mean")
    plt.title("BloomOS logits_mean by gate_used")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def main() -> None:
    if len(sys.argv) < 2:
        print("Usage: python3 bloomos_dashboard.py path/to/log.jsonl")
        raise SystemExit(1)

    log_path = Path(sys.argv[1])

    if not log_path.exists():
        print(f"Error: Log file not found: {log_path}")
        raise SystemExit(1)

    entries = load_log(log_path)

    if not entries:
        print("No entries loaded from", log_path)
        raise SystemExit(1)

    print(f"Loaded {len(entries)} entries from {log_path}")

    if HAS_MATPLOTLIB:
        plot_dashboard(entries)
    else:
        print("(matplotlib not available, showing text summary)")
        print_summary(entries)


if __name__ == "__main__":
    main()
