# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""CLI entry point for retroactive_llm_extractor (SoC split)."""
from __future__ import annotations

import importlib, time
from datetime import datetime, timezone
from pulse.admin.retroactive.retroactive_llm_extractor import (
    batch_interactions, write_to_brain, _archive_brain,
    CAPTURE_LOGS, SCRIBE_PROMPT, BATCH_SIZE, STATE_DIR,
    _ALLOWED_REBUILD_MODULES, dispatch, extract_json,
)


def main():
    import argparse
    ap = argparse.ArgumentParser(description="Brain V2 Retroactive LLM Extraction")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--no-archive", action="store_true")
    ap.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    args = ap.parse_args()
    sep = "=" * 60
    print(sep + "\nBRAIN V2 — Retroactive LLM Extraction\n" + sep)
    print("Batch size: " + str(args.batch_size) + (" | DRY RUN" if args.dry_run else ""))
    if not args.no_archive and not args.dry_run:
        print("\n[1/4] Archiving brain files..."); _archive_brain()
    print("\n[2/4] Processing capture logs...")
    total_batches, total_errors = 0, 0
    total_items = {"lessons": 0, "failures": 0, "patterns": 0, "facts": 0}
    start = time.time()

    for log_path in CAPTURE_LOGS:
        if not log_path.exists():
            continue
        agent = log_path.stem.replace("pulse_captured_", "")
        print("\n  [" + agent + "] " + str(round(log_path.stat().st_size / 1048576, 1)) + "MB...")
        batch_num, log_items = 0, {"lessons": 0, "failures": 0, "patterns": 0, "facts": 0}
        est = int(log_path.stat().st_size / (args.batch_size * 800)) or 1

        for batch_text, count, first_ts in batch_interactions(log_path, args.batch_size):
            batch_num += 1
            total_batches += 1
            if args.dry_run:
                print("    Batch " + str(batch_num) + ": " + str(count) + " interactions, ~" + str(len(batch_text)//4) + " tokens")
                continue
            try:
                raw = dispatch(SCRIBE_PROMPT + batch_text, task_type="fast", confidence=0.5)
                written = write_to_brain(
                    extract_json(raw), agent,
                    first_ts or datetime.now(timezone.utc).isoformat())
                for k in log_items:
                    log_items[k] += written[k]
                    total_items[k] += written[k]
                grand = sum(total_items.values())
                rate = total_batches / (time.time() - start) if time.time() > start else 0
                if batch_num % 10 == 0:
                    print("    [" + agent + "] " + str(batch_num) + "/~" + str(est) + " (" + str(grand) + " total) [" + str(round(rate, 1)) + " b/s]")
            except Exception as e:
                total_errors += 1
                print("\n    [" + agent + "] Batch " + str(batch_num) + " ERROR: " + str(e))
        print("  [" + agent + "] " + str(batch_num) + " batches -> " + str(sum(log_items.values())) + " items")

    if not args.dry_run:
        print("\n[3/4] Updating bridge state...")
        try:
            from pulse.core.brain_utils import atomic_update_json
            def _reset(data):
                data = data if isinstance(data, dict) else {}
                for lp in CAPTURE_LOGS:
                    if lp.exists():
                        data[str(lp.name)] = lp.stat().st_size
                return data
            atomic_update_json(STATE_DIR / "pulse_bridge_state.json", _reset, default={})
        except Exception as e:
            print("  Bridge state: " + str(e))
        print("\n[4/4] Rebuilding index + graph...")
        for label, mod, fn in [("Index", "pulse.brain.index.search_index", "build_index"),
                                ("Graph", "pulse.graph.knowledge_graph_build", "build_graph")]:
            try:
                if mod not in _ALLOWED_REBUILD_MODULES:
                    raise ValueError("Module not in allowlist: " + repr(mod))
                getattr(importlib.import_module(mod), fn)(verbose=True)
            except Exception as e:
                print("  " + label + ": " + str(e))

    elapsed, grand = time.time() - start, sum(total_items.values())
    ti = total_items
    print(
        "\n" + "="*60 + "\nDONE: " + str(total_batches) + " batches -> " + str(grand) + " items "
        "(L:" + str(ti["lessons"]) + " F:" + str(ti["failures"]) + " P:" + str(ti["patterns"]) + " Facts:" + str(ti["facts"]) + ") "
        "| " + str(total_errors) + " errors | " + str(int(elapsed)) + "s\n" + "="*60)


if __name__ == "__main__":
    main()
