"""NEUROVAULT CLI."""

from __future__ import annotations

import csv
import json
import sys
import time
from pathlib import Path
from io import StringIO

import click
from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from neurovault.daemon import NeurovaultDaemon, stop_running_daemon
from neurovault.engine import NeurovaultEngine

console = Console()


def _get_vault(name: str = "default") -> NeurovaultEngine:
    try:
        return NeurovaultEngine.open(name)
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)


def _fmt_ago(ts_ms: int) -> str:
    delta = time.time() - (ts_ms / 1000)
    if delta < 60:
        return f"{int(delta)}s ago"
    if delta < 3600:
        return f"{int(delta / 60)}m ago"
    if delta < 86400:
        return f"{int(delta / 3600)}h ago"
    return f"{int(delta / 86400)}d ago"


@click.group()
@click.version_option(version="0.3.1")
def main():
    """NEUROVAULT - The Human Brain's Digital Twin."""


@main.command()
@click.argument("name", default="default")
@click.option("--description", "description", default="")
@click.option("--no-embedding", is_flag=True)
def init(name: str, description: str, no_embedding: bool):
    """Initialize a NEUROVAULT brain."""
    try:
        NeurovaultEngine.init(name, description, no_embedding)
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)

    console.print(
        Panel(
            f"[green]Initialized[/green] vault [bold]{name}[/bold]\n"
            f"Embedding: {'disabled' if no_embedding else 'auto'}\n"
            "Sensors: clipboard, file_watcher, browser, screenshot, audio, calendar, git",
            title="NEUROVAULT",
            border_style="magenta",
        )
    )


@main.command()
@click.argument("content")
@click.option("--vault", "vault", default="default")
@click.option("--source", "source", default="manual")
@click.option("--importance", "importance", type=float, default=0.5)
def ingest(content: str, vault: str, source: str, importance: float):
    """Manually ingest content into memory."""
    engine = _get_vault(vault)
    memory = engine.ingest(content, source=source, importance=importance)
    if memory is None:
        console.print("[yellow]Blocked by privacy filter.[/yellow]")
        return
    console.print(f"[green]Ingested[/green] memory [{memory.short_id}] from [cyan]{source}[/cyan]")


@main.command()
@click.argument("query")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=10)
@click.option("--mode", "mode", type=click.Choice(["keyword", "semantic", "graph", "multi"]), default="multi")
@click.option("--format", "fmt", type=click.Choice(["text", "json", "csv"]), default="text")
@click.option("--with-entities", "with_entities", is_flag=True, help="Show extracted entities")
@click.option("--with-graph", "with_graph", is_flag=True, help="Show graph depth where available")
def recall(query: str, vault: str, limit: int, mode: str, fmt: str, with_entities: bool, with_graph: bool):
    """Search memories."""
    engine = _get_vault(vault)
    results = engine.recall(query, limit=limit, mode=mode)

    if fmt == "json":
        data = [
            {
                "id": r.memory.short_id,
                "content": r.memory.content,
                "score": r.score,
                "type": getattr(r, "match_type", mode),
                "graph_depth": getattr(r, "graph_depth", None),
            }
            for r in results
        ]
        console.print_json(json.dumps(data, indent=2))
        return
    if fmt == "csv":
        buffer = StringIO()
        writer = csv.DictWriter(
            buffer,
            fieldnames=["id", "content", "score", "type", "graph_depth"],
        )
        writer.writeheader()
        for r in results:
            writer.writerow(
                {
                    "id": r.memory.short_id,
                    "content": r.memory.content.replace("\n", " "),
                    "score": f"{r.score:.6f}",
                    "type": getattr(r, "match_type", mode),
                    "graph_depth": getattr(r, "graph_depth", ""),
                }
            )
        console.print(buffer.getvalue())
        return

    if not results:
        console.print("[yellow]No memories found.[/yellow]")
        return

    table = Table(
        title=f'Recall: "{query}" ({len(results)} results, mode={mode})',
        box=box.ROUNDED,
        show_lines=True,
    )
    table.add_column("#", width=3)
    table.add_column("ID", width=10, style="cyan")
    table.add_column("Content", style="white")
    table.add_column("Score", width=8, style="green")
    table.add_column("Age", width=10, style="dim")

    for i, result in enumerate(results, start=1):
        entity_suffix = ""
        if with_entities and getattr(result, "entities", None):
            entity_suffix = "\nEntities: " + ", ".join(e.name for e in result.entities[:5])
        graph_suffix = ""
        if with_graph and getattr(result, "graph_depth", None) is not None:
            graph_suffix = f"\nGraph depth: {result.graph_depth}"
        table.add_row(
            str(i),
            result.memory.short_id,
            (result.memory.content[:100] + entity_suffix + graph_suffix)[:260],
            f"{result.score:.3f}",
            _fmt_ago(result.memory.created_at),
        )

    console.print(table)


@main.command("predict")
@click.option("--vault", "vault", default="default")
def predict_cmd(vault: str):
    """Show proactive memory predictions based on current context."""
    engine = _get_vault(vault)
    predictions = engine.predict()
    if not predictions:
        console.print("[dim]No predictions right now. Update context and try again.[/dim]")
        return

    table = Table(title="Predictions", box=box.ROUNDED, show_lines=True)
    table.add_column("#", width=3)
    table.add_column("ID", width=10, style="cyan")
    table.add_column("Memory", style="white")
    table.add_column("Score", width=8, style="green")

    for i, prediction in enumerate(predictions, start=1):
        table.add_row(
            str(i),
            prediction.memory.short_id,
            prediction.memory.content[:120],
            f"{prediction.score:.3f}",
        )
    console.print(table)


@main.command("context")
@click.option("--vault", "vault", default="default")
@click.option("--type", "context_type", required=True, help="Context key, e.g. active_file")
@click.option("--value", "value", required=True, help="Context value to track")
def context_cmd(vault: str, context_type: str, value: str):
    """Update active context for prediction engine."""
    engine = _get_vault(vault)
    engine.update_context(context_type, value)
    console.print(f"[green]Context updated[/green] {context_type}={value}")


@main.command("consolidate")
@click.option("--vault", "vault", default="default")
@click.option("--force", "force", is_flag=True)
@click.option("--phases", "phases", default=None, help="Comma-separated phases")
@click.option("--verbose", "verbose", is_flag=True)
def consolidate_cmd(vault: str, force: bool, phases: str | None, verbose: bool):
    """Run a consolidation cycle."""
    engine = _get_vault(vault)
    phase_list = [p.strip() for p in phases.split(",")] if phases else None
    if verbose:
        console.print(f"[dim]Running consolidation phases: {phase_list or 'all'}[/dim]")
    report = engine.consolidate(force=force, phases=phase_list)
    console.print(
        Panel(
            f"Entities extracted:  [green]{report.entities_extracted}[/green]\n"
            f"Relationships found: [green]{report.relationships_found}[/green]\n"
            f"Clusters formed:     [green]{report.clusters_formed}[/green]\n"
            f"Insights generated:  [green]{report.insights_generated}[/green]\n"
            f"Cross-domain links:  [green]{int(getattr(report, 'cross_domain_insights', 0))}[/green]\n"
            f"Memories promoted:   [cyan]{report.memories_promoted}[/cyan]\n"
            f"Memories decayed:    [yellow]{report.memories_decayed}[/yellow]\n"
            f"Memories archived:   [red]{report.memories_archived}[/red]\n"
            f"Duration:            {report.duration_seconds:.2f}s",
            title="Consolidation Report",
            border_style="magenta",
        )
    )


@main.command("insights")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=10)
def insights_cmd(vault: str, limit: int):
    """List generated insight memories."""
    engine = _get_vault(vault)
    insights = engine.get_insights(limit=limit)
    if not insights:
        console.print("[yellow]No insights available.[/yellow]")
        return
    for i, memory in enumerate(insights, start=1):
        console.print(f"{i}. [{memory.short_id}] {memory.content[:160]}")


@main.command("entities")
@click.option("--vault", "vault", default="default")
@click.option("--type", "entity_type", default=None)
@click.option("--limit", "limit", type=int, default=20)
@click.option(
    "--sort",
    "sort",
    type=click.Choice(["mention_count", "importance", "last_seen"]),
    default="mention_count",
)
def entities_cmd(vault: str, entity_type: str | None, limit: int, sort: str):
    """List entities in the knowledge graph."""
    engine = _get_vault(vault)
    entities = engine.get_entities(entity_type=entity_type, limit=limit, sort=sort)

    if not entities:
        console.print("[yellow]No entities found.[/yellow]")
        return

    table = Table(title="Knowledge Graph Entities", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Type", style="magenta")
    table.add_column("Mentions", style="green", width=8)
    table.add_column("Importance", style="yellow", width=10)
    table.add_column("Last Seen", style="dim", width=10)

    for entity in entities:
        table.add_row(
            entity.name,
            entity.entity_type,
            str(entity.mention_count),
            f"{entity.importance:.2f}",
            _fmt_ago(entity.last_seen),
        )

    console.print(table)


@main.command("relations")
@click.argument("entity_name")
@click.option("--vault", "vault", default="default")
def relations_cmd(entity_name: str, vault: str):
    """Show relationships for an entity."""
    engine = _get_vault(vault)
    relationships = engine.get_relationships(entity_name)
    if not relationships:
        console.print(f"[yellow]No relationships found for '{entity_name}'.[/yellow]")
        return

    for rel in relationships:
        try:
            source = engine._nv.get_entity(rel.source_entity_id)
            target = engine._nv.get_entity(rel.target_entity_id)
        except Exception:
            continue
        strength_bar = "█" * max(1, int(rel.strength * 10))
        console.print(
            f"[cyan]{source.name}[/cyan] --[magenta]{rel.relation_type}[/magenta]--> "
            f"[cyan]{target.name}[/cyan] [{strength_bar}] {rel.strength:.2f}"
        )


@main.command("graph")
@click.argument("entity_name")
@click.option("--vault", "vault", default="default")
@click.option("--depth", "depth", type=int, default=2)
@click.option("--format", "fmt", type=click.Choice(["text", "json", "mermaid"]), default="text")
def graph_cmd(entity_name: str, vault: str, depth: int, fmt: str):
    """Show entity graph neighborhood."""
    engine = _get_vault(vault)
    graph = engine.get_entity_graph(entity_name, depth=depth)
    if graph is None:
        console.print(f"[yellow]Entity not found: {entity_name}[/yellow]")
        return

    if fmt == "json":
        payload = {
            "root": graph.root_entity.name,
            "depth": graph.depth,
            "nodes": [n.name for n in graph.nodes],
            "relationships": [
                {
                    "type": r.relation_type,
                    "source": r.source_entity_id.hex()[:8],
                    "target": r.target_entity_id.hex()[:8],
                    "strength": r.strength,
                }
                for r in graph.relationships
            ],
        }
        console.print_json(json.dumps(payload, indent=2))
        return
    if fmt == "mermaid":
        seen = set()
        lines = ["graph LR"]
        for rel in graph.relationships:
            try:
                src = engine._nv.get_entity(rel.source_entity_id)
                tgt = engine._nv.get_entity(rel.target_entity_id)
            except Exception:
                continue
            edge = (src.name, rel.relation_type, tgt.name)
            if edge in seen:
                continue
            seen.add(edge)
            src_id = src.name.replace(" ", "_")
            tgt_id = tgt.name.replace(" ", "_")
            label = rel.relation_type.replace("\"", "'")
            lines.append(f'  {src_id}["{src.name}"] -->|"{label}"| {tgt_id}["{tgt.name}"]')
        console.print("\n".join(lines))
        return

    console.print(f"[bold]Graph root:[/bold] {graph.root_entity.name} (depth={depth})")
    for rel in graph.relationships[:100]:
        try:
            src = engine._nv.get_entity(rel.source_entity_id)
            tgt = engine._nv.get_entity(rel.target_entity_id)
        except Exception:
            continue
        console.print(f"{src.name} --{rel.relation_type}--> {tgt.name} ({rel.strength:.2f})")


@main.command("review")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=20)
def review_cmd(vault: str, limit: int):
    """Show memories due for spaced review."""
    engine = _get_vault(vault)
    items = engine.review_due(limit=limit)
    if not items:
        console.print("[dim]No memories due for review.[/dim]")
        return
    for i, memory in enumerate(items, start=1):
        console.print(f"{i}. [{memory.short_id}] {memory.content[:140]}")


@main.command("decay")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=20)
@click.option("--run", "run_cycle", is_flag=True, help="Run CHRONOS decay cycle before listing")
def decay_cmd(vault: str, limit: int, run_cycle: bool):
    """Show lowest-importance active memories (decay candidates)."""
    engine = _get_vault(vault)
    if run_cycle:
        stats = engine.run_decay()
        console.print(
            f"[green]Decay cycle complete[/green] "
            f"updated={stats.get('updated', 0)} archived={stats.get('archived', 0)}"
        )

    records = engine._bridge.get_active_memories_with_metadata()
    records.sort(key=lambda r: float(r.get("importance", 0.0)))
    records = records[: max(1, limit)]
    if not records:
        console.print("[dim]No active memories.[/dim]")
        return

    table = Table(title="Decay Candidates", box=box.ROUNDED)
    table.add_column("ID", style="cyan", width=10)
    table.add_column("Importance", style="yellow", width=10)
    table.add_column("Tier", style="magenta", width=12)
    table.add_column("Content", style="white")

    for rec in records:
        try:
            mem = engine._bridge.get_memory(rec["id"])
        except Exception:
            continue
        meta = rec.get("metadata", {}) or {}
        table.add_row(
            mem.short_id,
            f"{float(rec.get('importance', 0.0)):.3f}",
            str(meta.get("tier", "unknown")),
            mem.content[:120],
        )
    console.print(table)


@main.command("reinforce")
@click.argument("memory_id")
@click.option("--vault", "vault", default="default")
def reinforce_cmd(memory_id: str, vault: str):
    """Reinforce a memory by ID."""
    engine = _get_vault(vault)
    engine.reinforce(memory_id)
    console.print(f"[green]Reinforced[/green] {memory_id}")


@main.command("forget")
@click.argument("memory_id")
@click.option("--vault", "vault", default="default")
@click.option("--hard", "hard", is_flag=True)
def forget_cmd(memory_id: str, vault: str, hard: bool):
    """Forget a memory by ID."""
    engine = _get_vault(vault)
    engine.forget(memory_id, hard=hard)
    console.print(f"[green]Forgot memory[/green] {memory_id} (hard={hard})")


@main.command("audit")
@click.option("--vault", "vault", default="default")
@click.option("--since", "since", type=int, default=None, help="Unix ms lower bound")
@click.option("--format", "fmt", type=click.Choice(["object", "dict", "json"]), default="json")
def audit_cmd(vault: str, since: int | None, fmt: str):
    """Generate audit report."""
    engine = _get_vault(vault)
    report = engine.audit(since=since, format=fmt)
    if fmt == "json":
        console.print(report)
    elif fmt == "dict":
        console.print_json(json.dumps(report, indent=2))
    else:
        console.print(report)


@main.command("export")
@click.option("--vault", "vault", default="default")
@click.option("--path", "path", required=True)
@click.option("--format", "fmt", type=click.Choice(["json", "csv"]), default="json")
@click.option("--since", "since", type=int, default=None)
def export_cmd(vault: str, path: str, fmt: str, since: int | None):
    """Export memories to file."""
    engine = _get_vault(vault)
    engine.export(path=path, format=fmt, since=since)
    console.print(f"[green]Exported[/green] memories to {path}")


@main.command("import")
@click.option("--vault", "vault", default="default")
@click.option("--path", "path", required=True)
def import_cmd(vault: str, path: str):
    """Import memories from JSON export."""
    engine = _get_vault(vault)
    result = engine.import_memories(path)
    console.print(
        f"[green]Import complete[/green] imported={result.imported} skipped={result.skipped} "
        f"errors={len(result.errors)}"
    )


@main.command("status")
@click.option("--vault", "vault", default="default")
def status_cmd(vault: str):
    """Show brain status."""
    engine = _get_vault(vault)
    status_obj = engine.status
    size_mb = status_obj.db_size_bytes / (1024 * 1024)

    tiers = "\n".join(
        f"  - {name.replace('_', ' ').title()}: {count}"
        for name, count in status_obj.tier_counts.items()
    )
    sensors = "\n".join(
        f"  - {name}: {'on' if meta.get('enabled') else 'off'}"
        for name, meta in status_obj.sensor_status.items()
    )

    console.print(
        Panel(
            f"Daemon: {'RUNNING' if status_obj.daemon_running else 'STOPPED'}"
            + (f" (PID {status_obj.daemon_pid})" if status_obj.daemon_pid else "")
            + f"\n\nMEMORY TIERS\n{tiers}"
            + "\n\nKNOWLEDGE GRAPH\n"
            + f"  - Entities: {status_obj.entity_count}\n"
            + f"  - Relationships: {status_obj.relationship_count}\n"
            + f"  - Insights: {status_obj.insight_count}\n"
            + (
                f"\nLast Consolidation: {_fmt_ago(status_obj.last_consolidation)}"
                if status_obj.last_consolidation
                else "\nLast Consolidation: never"
            )
            + f"\n\nSENSORS\n{sensors}"
            + f"\nDB Size: {size_mb:.2f} MB",
            title="NEUROVAULT Status",
            border_style="magenta",
        )
    )


@main.command("stats")
@click.option("--vault", "vault", default="default")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON")
def stats_cmd(vault: str, as_json: bool):
    """Show compact runtime stats snapshot."""
    engine = _get_vault(vault)
    payload = engine.stats()
    if as_json:
        console.print_json(json.dumps(payload, indent=2))
        return

    table = Table(title="NEUROVAULT Stats", box=box.ROUNDED)
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="white")
    for key in (
        "memory_count",
        "entity_count",
        "relationship_count",
        "concept_count",
        "cluster_count",
        "insight_count",
        "plugin_count",
        "daemon_running",
        "daemon_pid",
    ):
        table.add_row(key, str(payload.get(key)))
    console.print(table)


@main.group("hnsw")
def hnsw_group():
    """Inspect and manage the v0.3 HNSW vector index."""


@hnsw_group.command("status")
@click.option("--vault", "vault", default="default")
def hnsw_status_cmd(vault: str):
    """Show HNSW index runtime status."""
    engine = _get_vault(vault)
    console.print_json(json.dumps(engine.hnsw_status(), indent=2))


@hnsw_group.command("rebuild")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=5000)
def hnsw_rebuild_cmd(vault: str, limit: int):
    """Rebuild HNSW index from stored memories."""
    engine = _get_vault(vault)
    payload = engine.rebuild_hnsw_index(limit=limit)
    console.print_json(json.dumps(payload, indent=2))


@main.group("federation")
def federation_group():
    """Operate v0.3 P2P sync and federation surfaces."""


@federation_group.command("status")
@click.option("--vault", "vault", default="default")
def federation_status_cmd(vault: str):
    """Show federation and sync status."""
    engine = _get_vault(vault)
    payload = {
        "sync": engine.sync_status(),
        "hnsw": engine.hnsw_status(),
    }
    console.print_json(json.dumps(payload, indent=2))


@federation_group.command("bundle")
@click.option("--vault", "vault", default="default")
@click.option("--since", "since_checkpoint", default=None)
@click.option("--out", "out_path", default=None, help="Write bundle JSON to file")
def federation_bundle_cmd(vault: str, since_checkpoint: str | None, out_path: str | None):
    """Create local delta bundle for peer transfer."""
    engine = _get_vault(vault)
    payload = engine.create_delta_bundle(since_checkpoint=since_checkpoint)
    if out_path:
        Path(out_path).expanduser().write_text(json.dumps(payload, indent=2))
        console.print(f"[green]Wrote bundle[/green] {out_path}")
        return
    console.print_json(json.dumps(payload, indent=2))


@federation_group.command("apply")
@click.option("--vault", "vault", default="default")
@click.option("--bundle", "bundle_path", required=True, help="Path to bundle JSON file")
def federation_apply_cmd(vault: str, bundle_path: str):
    """Apply an incoming delta bundle from disk."""
    engine = _get_vault(vault)
    try:
        payload = json.loads(Path(bundle_path).expanduser().read_text())
    except Exception as exc:
        console.print(f"[red]Failed to read bundle:[/red] {exc}")
        sys.exit(1)
    result = engine.apply_delta_bundle(payload)
    console.print_json(json.dumps(result, indent=2))


@federation_group.command("sync")
@click.option("--vault", "vault", default="default")
@click.option("--host", "host", required=True)
@click.option("--port", "port", type=int, default=9473)
@click.option("--since", "since_checkpoint", default=None)
def federation_sync_cmd(vault: str, host: str, port: int, since_checkpoint: str | None):
    """Perform bidirectional sync with a peer endpoint."""
    engine = _get_vault(vault)
    result = engine.sync_with_peer(host=host, port=port, since_checkpoint=since_checkpoint)
    console.print_json(json.dumps(result, indent=2))


@main.group("access")
def access_group():
    """Manage v0.3 users and device tokens."""


@access_group.command("create-user")
@click.argument("username")
@click.option("--vault", "vault", default="default")
@click.option("--role", "role", default="reader")
@click.option("--created-by", "created_by", default="system")
def access_create_user_cmd(username: str, vault: str, role: str, created_by: str):
    """Create a new user and output its token."""
    engine = _get_vault(vault)
    try:
        payload = engine.create_user(username=username, role=role, created_by=created_by)
    except Exception as exc:
        console.print(f"[red]Create user failed:[/red] {exc}")
        sys.exit(1)
    console.print_json(json.dumps(payload, indent=2))


@access_group.command("list-users")
@click.option("--vault", "vault", default="default")
def access_list_users_cmd(vault: str):
    """List configured users."""
    engine = _get_vault(vault)
    console.print_json(json.dumps(engine.list_users(), indent=2))


@access_group.command("revoke-user")
@click.argument("username")
@click.option("--vault", "vault", default="default")
def access_revoke_user_cmd(username: str, vault: str):
    """Deactivate a user account."""
    engine = _get_vault(vault)
    console.print_json(json.dumps({"ok": engine.revoke_user(username)}, indent=2))


@access_group.command("create-device-token")
@click.argument("device_name")
@click.option("--vault", "vault", default="default")
@click.option("--permissions", "permissions", default="read", help="Comma-separated permissions")
def access_create_device_token_cmd(device_name: str, vault: str, permissions: str):
    """Create a device token for mobile companion auth."""
    engine = _get_vault(vault)
    perms = {p.strip() for p in permissions.split(",") if p.strip()}
    token = engine.create_device_token(device_name=device_name, permissions=perms)
    console.print_json(json.dumps({"device_name": device_name, "token": token, "permissions": sorted(perms)}, indent=2))


@main.command("audit-log")
@click.option("--vault", "vault", default="default")
@click.option("--user", "user", default=None)
@click.option("--action", "action", default=None)
@click.option("--since", "since", type=int, default=None)
@click.option("--limit", "limit", type=int, default=50)
def audit_log_cmd(vault: str, user: str | None, action: str | None, since: int | None, limit: int):
    """Query security audit events."""
    engine = _get_vault(vault)
    rows = engine.query_audit(user=user, action=action, since=since, limit=limit)
    console.print_json(json.dumps(rows, indent=2))


@main.command("config")
@click.option("--vault", "vault", default="default")
def config_cmd(vault: str):
    """Show effective configuration."""
    engine = _get_vault(vault)
    cfg = {
        "sensors": engine._config.sensors,
        "privacy": engine._config.privacy,
        "consolidation": engine._config.consolidation,
        "plugin_configs": engine._config.plugin_configs,
        "encryption": engine._config.encryption,
        "llm": engine._config.llm,
        "embedding": engine._config.embedding,
        "dashboard": engine._config.dashboard,
        "pipeline": engine._config.pipeline,
        "observability": engine._config.observability,
        "network": engine._config.network,
        "search": engine._config.search,
        "api": engine._config.api,
        "monitoring": engine._config.monitoring,
        "disabled_sensors": engine._config.disabled_sensors,
    }
    console.print_json(json.dumps(cfg, indent=2))


@main.command("sensors")
@click.option("--vault", "vault", default="default")
@click.option("--enable", "enable_sensor", default=None, help="Enable one sensor by name")
@click.option("--disable", "disable_sensor", default=None, help="Disable one sensor by name")
def sensors_cmd(vault: str, enable_sensor: str | None, disable_sensor: str | None):
    """Show or update sensor enablement status."""
    engine = _get_vault(vault)
    if enable_sensor:
        engine._config.set_sensor_enabled(enable_sensor, True)
        engine._config.save()
    if disable_sensor:
        engine._config.set_sensor_enabled(disable_sensor, False)
        engine._config.save()

    rows = []
    disabled = set(engine._config.disabled_sensors)
    for name, cfg in engine._config.sensors.items():
        enabled = name not in disabled and bool(cfg.get("enabled", True))
        rows.append((name, enabled, cfg.get("poll_interval")))

    table = Table(title="Sensors", box=box.ROUNDED)
    table.add_column("Sensor", style="cyan")
    table.add_column("Enabled", style="green")
    table.add_column("Poll Interval", style="dim")
    for name, enabled, interval in rows:
        table.add_row(name, "yes" if enabled else "no", str(interval))
    console.print(table)


@main.group("plugin")
def plugin_group():
    """Manage plugin integrations (local + CHRONOS)."""


@plugin_group.command("list")
@click.option("--vault", "vault", default="default")
@click.option("--scope", "scope", type=click.Choice(["all", "local", "chronos"]), default="all")
def plugin_list_cmd(vault: str, scope: str):
    """List discovered plugins."""
    engine = _get_vault(vault)
    local_plugins = engine.list_plugins() if scope in ("all", "local") else []
    chronos_plugins = engine.list_chronos_plugins() if scope in ("all", "chronos") else []

    if not local_plugins and not chronos_plugins:
        console.print("[dim]No plugins discovered.[/dim]")
        return

    table = Table(title="Plugins", box=box.ROUNDED)
    table.add_column("Scope", style="dim")
    table.add_column("ID", style="cyan")
    table.add_column("Name", style="white")
    table.add_column("Version", style="green")
    table.add_column("Enabled", style="magenta")
    table.add_column("Description", style="dim")
    for item in local_plugins:
        table.add_row(
            "local",
            str(item.get("id", "")),
            str(item.get("name", "")),
            str(item.get("version", "")),
            "yes" if item.get("enabled") else "no",
            str(item.get("description", "")),
        )
    for item in chronos_plugins:
        table.add_row(
            "chronos",
            str(item.get("id", "")),
            str(item.get("name", "")),
            str(item.get("version", "")),
            "yes" if item.get("active") else "no",
            str(item.get("description", "")),
        )
    console.print(table)


@plugin_group.command("status")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_status_cmd(plugin_id: str, vault: str):
    """Show runtime status for a CHRONOS plugin."""
    engine = _get_vault(vault)
    payload = engine.get_chronos_plugin_status(plugin_id)
    console.print_json(json.dumps(payload, indent=2))


@plugin_group.command("activate")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
@click.option("--config-json", "config_json", default="{}", help="JSON object for plugin config")
def plugin_activate_cmd(plugin_id: str, vault: str, config_json: str):
    """Activate a CHRONOS plugin."""
    try:
        config = json.loads(config_json)
    except Exception as exc:
        console.print(f"[red]Invalid --config-json:[/red] {exc}")
        sys.exit(1)
    if not isinstance(config, dict):
        console.print("[red]--config-json must decode to a JSON object[/red]")
        sys.exit(1)

    engine = _get_vault(vault)
    ok = engine.activate_chronos_plugin(plugin_id, config=config)
    if not ok:
        console.print(f"[red]Failed to activate plugin:[/red] {plugin_id}")
        sys.exit(1)
    console.print(f"[green]Activated plugin[/green] {plugin_id}")


@plugin_group.command("deactivate")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_deactivate_cmd(plugin_id: str, vault: str):
    """Deactivate a CHRONOS plugin."""
    engine = _get_vault(vault)
    ok = engine.deactivate_chronos_plugin(plugin_id)
    if not ok:
        console.print(f"[red]Failed to deactivate plugin:[/red] {plugin_id}")
        sys.exit(1)
    console.print(f"[green]Deactivated plugin[/green] {plugin_id}")


@plugin_group.command("install")
@click.argument("source")
@click.option("--vault", "vault", default="default")
def plugin_install_cmd(source: str, vault: str):
    """Install a local plugin directory."""
    engine = _get_vault(vault)
    result = engine.install_plugin(source)
    if not result.get("ok"):
        console.print(f"[red]Install failed:[/red] {result.get('error', 'unknown error')}")
        sys.exit(1)
    console.print(f"[green]Installed plugin[/green] {result.get('plugin_id')}")


@plugin_group.command("remove")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_remove_cmd(plugin_id: str, vault: str):
    """Remove an installed local plugin."""
    engine = _get_vault(vault)
    if not engine.remove_plugin(plugin_id):
        console.print(f"[yellow]Plugin not found on disk:[/yellow] {plugin_id}")
        return
    console.print(f"[green]Removed plugin[/green] {plugin_id}")


@plugin_group.command("enable")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_enable_cmd(plugin_id: str, vault: str):
    """Enable a local plugin."""
    engine = _get_vault(vault)
    if not engine.enable_plugin(plugin_id):
        console.print(f"[red]Failed to enable plugin:[/red] {plugin_id}")
        sys.exit(1)
    console.print(f"[green]Enabled plugin[/green] {plugin_id}")


@plugin_group.command("disable")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_disable_cmd(plugin_id: str, vault: str):
    """Disable a local plugin."""
    engine = _get_vault(vault)
    if not engine.disable_plugin(plugin_id):
        console.print(f"[red]Failed to disable plugin:[/red] {plugin_id}")
        sys.exit(1)
    console.print(f"[green]Disabled plugin[/green] {plugin_id}")


@plugin_group.command("health")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
def plugin_health_cmd(plugin_id: str, vault: str):
    """Show local plugin health state."""
    engine = _get_vault(vault)
    payload = engine.plugin_health(plugin_id)
    console.print_json(json.dumps(payload, indent=2))


@plugin_group.command("config")
@click.argument("plugin_id")
@click.option("--vault", "vault", default="default")
@click.option("--set", "set_items", multiple=True, help="Set config key-value pairs (key=value)")
def plugin_config_cmd(plugin_id: str, vault: str, set_items: tuple[str, ...]):
    """Configure a local plugin."""
    payload: dict[str, Any] = {}
    for item in set_items:
        if "=" not in item:
            console.print(f"[red]Invalid --set item:[/red] {item} (expected key=value)")
            sys.exit(1)
        key, raw_value = item.split("=", 1)
        key = key.strip()
        raw_value = raw_value.strip()
        if not key:
            console.print(f"[red]Invalid empty key in --set:[/red] {item}")
            sys.exit(1)
        try:
            value = json.loads(raw_value)
        except Exception:
            value = raw_value
        payload[key] = value

    engine = _get_vault(vault)
    if not engine.set_plugin_config(plugin_id, payload):
        console.print(f"[red]Failed to update config:[/red] {plugin_id}")
        sys.exit(1)
    console.print(f"[green]Updated config[/green] {plugin_id}")


@main.group("event")
def event_group():
    """Inspect CHRONOS event stream."""


@event_group.command("list")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=20)
@click.option("--type", "event_type", default=None, help="Filter by event type")
@click.option("--format", "fmt", type=click.Choice(["table", "json"]), default="table")
def event_list_cmd(vault: str, limit: int, event_type: str | None, fmt: str):
    """List recent CHRONOS events."""
    engine = _get_vault(vault)
    events = engine.list_chronos_events(limit=limit, event_type=event_type)

    if fmt == "json":
        console.print_json(json.dumps(events, indent=2))
        return

    if not events:
        console.print("[dim]No events found.[/dim]")
        return

    table = Table(title="CHRONOS Events", box=box.ROUNDED)
    table.add_column("When", style="dim", width=10)
    table.add_column("Type", style="cyan")
    table.add_column("Payload", style="white")
    for item in events:
        ts = int(item.get("created_at", 0))
        payload = item.get("payload", {})
        text = json.dumps(payload) if isinstance(payload, dict) else str(payload)
        table.add_row(_fmt_ago(ts), str(item.get("event_type", "")), text[:180])
    console.print(table)


@main.group("chronos-graph")
def chronos_graph_group():
    """Operate CHRONOS native graph engine hooks."""


@chronos_graph_group.command("rebuild")
@click.option("--vault", "vault", default="default")
def chronos_graph_rebuild_cmd(vault: str):
    """Rebuild CHRONOS native graph index (if available)."""
    engine = _get_vault(vault)
    ok = engine.rebuild_chronos_graph()
    if not ok:
        console.print("[yellow]CHRONOS graph rebuild unavailable in current runtime.[/yellow]")
        return
    console.print("[green]CHRONOS graph rebuild complete.[/green]")


@chronos_graph_group.command("neighbors")
@click.argument("entity_name")
@click.option("--vault", "vault", default="default")
@click.option("--format", "fmt", type=click.Choice(["table", "json"]), default="table")
def chronos_graph_neighbors_cmd(entity_name: str, vault: str, fmt: str):
    """List CHRONOS-native graph neighbors for an entity."""
    engine = _get_vault(vault)
    rows = engine.get_chronos_graph_neighbors(entity_name)

    if fmt == "json":
        console.print_json(json.dumps(rows, indent=2))
        return

    if not rows:
        console.print(f"[dim]No CHRONOS graph neighbors found for '{entity_name}'.[/dim]")
        return

    table = Table(title=f"CHRONOS Graph Neighbors: {entity_name}", box=box.ROUNDED)
    table.add_column("Entity", style="cyan")
    table.add_column("Relationship", style="magenta")
    table.add_column("Weight", style="green", width=8)
    for row in rows:
        table.add_row(
            str(row.get("entity", "")),
            str(row.get("relationship", "")),
            f"{float(row.get('weight', 0.0)):.2f}",
        )
    console.print(table)


@chronos_graph_group.command("path")
@click.argument("source")
@click.argument("target")
@click.option("--vault", "vault", default="default")
@click.option("--format", "fmt", type=click.Choice(["table", "json"]), default="table")
def chronos_graph_path_cmd(source: str, target: str, vault: str, fmt: str):
    """Resolve CHRONOS-native graph path between two entities."""
    engine = _get_vault(vault)
    rows = engine.get_chronos_graph_path(source, target)

    if fmt == "json":
        console.print_json(json.dumps(rows, indent=2))
        return

    if not rows:
        console.print(f"[dim]No CHRONOS graph path found from '{source}' to '{target}'.[/dim]")
        return

    table = Table(title=f"CHRONOS Graph Path: {source} -> {target}", box=box.ROUNDED)
    table.add_column("#", style="dim", width=3)
    table.add_column("Source", style="cyan")
    table.add_column("Relationship", style="magenta")
    table.add_column("Target", style="cyan")
    for index, row in enumerate(rows, start=1):
        table.add_row(
            str(index),
            str(row.get("source", "")),
            str(row.get("relationship", "")),
            str(row.get("target", "")),
        )
    console.print(table)


@main.group("persona")
def persona_group():
    """Manage CHRONOS personas (branches)."""


@persona_group.command("list")
@click.option("--vault", "vault", default="default")
def persona_list_cmd(vault: str):
    """List personas and active branch."""
    engine = _get_vault(vault)
    personas = engine.list_personas()
    if not personas:
        console.print("[dim]No personas found.[/dim]")
        return

    table = Table(title="Personas", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Active", style="green", width=8)
    table.add_column("ID", style="dim", width=16)
    table.add_column("Description", style="white")
    for item in personas:
        table.add_row(
            str(item.get("name", "")),
            "yes" if item.get("active") else "no",
            str(item.get("id", ""))[:14],
            str(item.get("description", "")),
        )
    console.print(table)


@persona_group.command("create")
@click.argument("name")
@click.option("--vault", "vault", default="default")
@click.option("--description", "description", default="")
@click.option("--fork-from", "fork_from", default=None)
@click.option("--use", "make_active", is_flag=True, help="Switch to this persona after creation")
def persona_create_cmd(name: str, vault: str, description: str, fork_from: str | None, make_active: bool):
    """Create a persona, optionally forking from another."""
    engine = _get_vault(vault)
    try:
        persona = engine.create_persona(
            name=name,
            description=description,
            fork_from=fork_from,
            make_active=make_active,
        )
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)

    console.print(
        f"[green]Persona created[/green] {persona.get('name')} "
        f"(active={engine.get_active_persona()})"
    )


@persona_group.command("use")
@click.argument("name")
@click.option("--vault", "vault", default="default")
def persona_use_cmd(name: str, vault: str):
    """Switch active persona."""
    engine = _get_vault(vault)
    if not engine.use_persona(name):
        console.print(f"[red]Failed to switch persona:[/red] {name}")
        sys.exit(1)
    console.print(f"[green]Active persona[/green] {engine.get_active_persona()}")


@persona_group.command("delete")
@click.argument("name")
@click.option("--vault", "vault", default="default")
@click.option("--force", "force", is_flag=True)
def persona_delete_cmd(name: str, vault: str, force: bool):
    """Delete a persona."""
    engine = _get_vault(vault)
    if not engine.delete_persona(name, force=force):
        console.print(f"[red]Failed to delete persona:[/red] {name}")
        sys.exit(1)
    console.print(f"[green]Deleted persona[/green] {name}")


@main.command("merge")
@click.argument("source_persona")
@click.option("--vault", "vault", default="default")
@click.option("--strategy", "strategy", type=click.Choice(["semantic", "additive"]), default="semantic")
@click.option("--threshold", "threshold", type=float, default=0.95)
@click.option("--dry-run", "dry_run", is_flag=True)
def merge_cmd(source_persona: str, vault: str, strategy: str, threshold: float, dry_run: bool):
    """Merge memories from source persona into active persona."""
    engine = _get_vault(vault)
    result = engine.merge_persona(
        source_persona=source_persona,
        strategy=strategy,
        threshold=threshold,
        dry_run=dry_run,
    )
    if result.get("error"):
        console.print(f"[red]Merge failed:[/red] {result['error']}")
        sys.exit(1)

    console.print(
        f"[green]Merge {('dry-run ' if dry_run else '')}complete[/green] "
        f"strategy={result.get('strategy')} added={result.get('added', 0)} "
        f"skipped={result.get('skipped', 0)} conflicts={result.get('conflicts', 0)}"
    )


@main.group("milestone")
def milestone_group():
    """Manage named checkpoint pointers."""


@milestone_group.command("create")
@click.argument("name")
@click.option("--vault", "vault", default="default")
@click.option("--ref", "ref", default="HEAD")
@click.option("--message", "message", default="")
def milestone_create_cmd(name: str, vault: str, ref: str, message: str):
    """Create a milestone for a checkpoint ref."""
    engine = _get_vault(vault)
    try:
        result = engine.create_milestone(name=name, ref=ref, message=message)
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)
    console.print(
        f"[green]Milestone created[/green] {result.get('name')} -> {str(result.get('checkpoint_id', ''))[:8]}"
    )


@milestone_group.command("list")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=20)
def milestone_list_cmd(vault: str, limit: int):
    """List milestones."""
    engine = _get_vault(vault)
    rows = engine.list_milestones(limit=limit)
    if not rows:
        console.print("[dim]No milestones found.[/dim]")
        return

    table = Table(title="Milestones", box=box.ROUNDED)
    table.add_column("Name", style="cyan")
    table.add_column("Checkpoint", style="green", width=12)
    table.add_column("Created", style="dim", width=10)
    table.add_column("Message", style="white")
    for item in rows:
        created = int(item.get("created_at", 0))
        table.add_row(
            str(item.get("name", "")),
            str(item.get("checkpoint_id", ""))[:10],
            _fmt_ago(created) if created else "-",
            str(item.get("message", "")),
        )
    console.print(table)


@main.command("revert")
@click.argument("ref")
@click.option("--vault", "vault", default="default")
@click.option("--hard", "hard", is_flag=True)
def revert_cmd(ref: str, vault: str, hard: bool):
    """Revert active persona to a checkpoint ref."""
    engine = _get_vault(vault)
    result = engine.revert(ref=ref, hard=hard)
    if result.get("error"):
        console.print(f"[red]Revert failed:[/red] {result['error']}")
        sys.exit(1)
    console.print(
        f"[green]Reverted[/green] to {result.get('short_id', '')} "
        f"({result.get('message', '')}) hard={hard}"
    )


@main.group("sync")
def sync_group():
    """Sync bundles across CHRONOS peers."""


@sync_group.command("push")
@click.argument("peer_name")
@click.option("--vault", "vault", default="default")
@click.option("--out", "output_path", required=True, help="Output bundle path")
def sync_push_cmd(peer_name: str, vault: str, output_path: str):
    """Export a sync bundle for a peer."""
    engine = _get_vault(vault)
    result = engine.sync_push(peer_name=peer_name, output_path=output_path)
    console.print_json(json.dumps(result, indent=2))


@sync_group.command("pull")
@click.option("--vault", "vault", default="default")
@click.option("--bundle", "bundle_path", required=True, help="Bundle file path")
@click.option("--strategy", "strategy", default="newer_wins")
def sync_pull_cmd(vault: str, bundle_path: str, strategy: str):
    """Import and merge a sync bundle."""
    engine = _get_vault(vault)
    result = engine.sync_pull(bundle_path=bundle_path, strategy=strategy)
    console.print_json(json.dumps(result, indent=2))


@main.group("keys")
def keys_group():
    """Manage signing keys."""


@keys_group.command("generate")
@click.argument("name", default="default")
@click.option("--vault", "vault", default="default")
def key_generate_cmd(name: str, vault: str):
    """Generate a signing keypair if crypto is available."""
    engine = _get_vault(vault)
    result = engine.generate_key(name=name)
    console.print_json(json.dumps(result, indent=2))


@main.command("verify")
@click.option("--vault", "vault", default="default")
def verify_cmd(vault: str):
    """Run deep integrity verification."""
    engine = _get_vault(vault)
    result = engine.verify_integrity()
    console.print_json(json.dumps(result, indent=2))


@main.command("start")
@click.option("--vault", "vault", default="default")
@click.option("--sensors", "sensors", default="", help="Comma-separated sensor list")
@click.option("--no-consolidation", "no_consolidation", is_flag=True, help="Disable auto-consolidation")
@click.option("--verbose", "verbose", is_flag=True, help="Show ingestion events")
def start_cmd(vault: str, sensors: str, no_consolidation: bool, verbose: bool):
    """Start the background daemon in the current process."""
    sensor_list = [s.strip() for s in sensors.split(",") if s.strip()] if sensors else None
    daemon = NeurovaultDaemon(
        vault_name=vault,
        sensors=sensor_list,
        enable_consolidation=not no_consolidation,
        verbose=verbose,
    )
    console.print("[dim]Starting daemon. Press Ctrl+C to stop.[/dim]")
    try:
        daemon.start()
    except KeyboardInterrupt:
        daemon.stop()


@main.command("stop")
def stop_cmd():
    """Stop a running daemon by PID file."""
    if stop_running_daemon():
        console.print("[green]Stop signal sent to daemon.[/green]")
    else:
        console.print("[yellow]No running daemon found.[/yellow]")


@main.command("dashboard")
@click.option("--vault", "vault", default="default")
@click.option("--port", "port", default=8501, type=int)
def dashboard_cmd(vault: str, port: int):
    """Launch the Streamlit dashboard."""
    try:
        import streamlit.web.cli as stcli  # type: ignore
    except Exception as exc:
        console.print(f"[red]Error:[/red] {exc}")
        console.print("Install dashboard extras: pip install neurovault[dashboard]")
        sys.exit(1)

    app_path = Path(__file__).parent / "dashboard" / "app.py"
    sys.argv = [
        "streamlit",
        "run",
        str(app_path),
        "--server.port",
        str(port),
        "--server.headless",
        "true",
        "--",
        vault,
    ]
    stcli.main()


@main.group("encrypt")
def encrypt_group():
    """Manage local encryption settings."""


@encrypt_group.command("set-key")
@click.option("--vault", "vault", default="default")
@click.option("--source", "source", type=click.Choice(["keychain", "file"]), default="keychain")
@click.option("--passphrase", "passphrase", prompt=True, hide_input=True, confirmation_prompt=True)
def encrypt_set_key_cmd(vault: str, source: str, passphrase: str):
    """Derive and store encryption key from a passphrase."""
    from neurovault.security.encryption import EncryptionManager

    engine = _get_vault(vault)
    manager = EncryptionManager(vault_dir=engine._nv.db_path.parent, vault_name=vault)
    key = manager.derive_key(passphrase)
    if not manager.store_key(key, source=source):
        console.print("[red]Failed to persist encryption key.[/red]")
        sys.exit(1)

    engine._config.encryption["enabled"] = True
    engine._config.encryption["key_source"] = source
    engine._config.save()
    console.print(f"[green]Encryption key stored[/green] source={source}")


@encrypt_group.command("status")
@click.option("--vault", "vault", default="default")
def encrypt_status_cmd(vault: str):
    """Show encryption config state."""
    from neurovault.security.encryption import EncryptionManager

    engine = _get_vault(vault)
    cfg = dict(engine._config.encryption)
    manager = EncryptionManager(vault_dir=engine._nv.db_path.parent, vault_name=vault)
    key_present = manager.get_key(cfg) is not None
    payload = {
        "enabled": bool(cfg.get("enabled", False)),
        "key_source": str(cfg.get("key_source", "keychain")),
        "key_present": key_present,
    }
    console.print_json(json.dumps(payload, indent=2))


@main.command("serve")
@click.option("--vault", "vault", default="default")
@click.option("--host", "host", default="127.0.0.1")
@click.option("--port", "port", default=8787, type=int)
def serve_cmd(vault: str, host: str, port: int):
    """Serve REST API (requires fastapi+uvicorn extras)."""
    try:
        import uvicorn  # type: ignore
    except Exception as exc:  # pragma: no cover
        console.print(f"[red]Error:[/red] {exc}")
        console.print("Install server extras: pip install neurovault[server]")
        return

    from neurovault.api import create_app

    app = create_app(vault_name=vault)
    uvicorn.run(app, host=host, port=port)


@main.command("checkpoint")
@click.option("--vault", "vault", default="default")
@click.option("--message", "message", required=True)
def checkpoint_cmd(vault: str, message: str):
    """Create a memory checkpoint."""
    engine = _get_vault(vault)
    checkpoint = engine.checkpoint(message)
    console.print(f"[green]Checkpoint[/green] [{checkpoint.short_id}] {message}")


@main.command("log")
@click.option("--vault", "vault", default="default")
@click.option("--limit", "limit", type=int, default=20)
def log_cmd(vault: str, limit: int):
    """Show checkpoint history."""
    engine = _get_vault(vault)
    checkpoints = engine._bridge.list_checkpoints(limit=limit)
    if not checkpoints:
        console.print("[dim]No checkpoints yet.[/dim]")
        return

    table = Table(title="Checkpoint Log", box=box.ROUNDED)
    table.add_column("ID", style="cyan")
    table.add_column("Message", style="white")
    table.add_column("Created", style="dim")

    for cp in checkpoints:
        table.add_row(cp.short_id, cp.message, _fmt_ago(cp.created_at))
    console.print(table)


@main.command("diff")
@click.argument("checkpoint_a")
@click.argument("checkpoint_b")
@click.option("--vault", "vault", default="default")
def diff_cmd(checkpoint_a: str, checkpoint_b: str, vault: str):
    """Compare two checkpoints."""
    engine = _get_vault(vault)
    payload = engine.checkpoint_diff(checkpoint_a, checkpoint_b, limit=20)
    added = payload.get("added", [])
    removed = payload.get("removed", [])

    lines = [
        f"Ref A: {payload.get('ref1', checkpoint_a)}",
        f"Ref B: {payload.get('ref2', checkpoint_b)}",
        f"Added: {len(added)}",
        f"Removed: {len(removed)}",
    ]
    if added:
        lines.append("\nAdded samples:")
        for row in added[:5]:
            lines.append(f"  + {str(row.get('id', ''))[:8]} {str(row.get('content', ''))[:80]}")
    if removed:
        lines.append("\nRemoved samples:")
        for row in removed[:5]:
            lines.append(f"  - {str(row.get('id', ''))[:8]} {str(row.get('content', ''))[:80]}")

    console.print(
        Panel(
            "\n".join(lines),
            title="Checkpoint Diff",
            border_style="magenta",
        )
    )


if __name__ == "__main__":
    main()
