#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Memory Dispatcher — Brain V4 Phase 2. Chunks logs/intents into swarm tasks."""
from __future__ import annotations

import io, json, os, sys, time
from datetime import datetime, timezone
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent.parent

if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

if os.name == "nt" and __name__ == "__main__":
    if hasattr(sys.stdout, 'reconfigure'):
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    if hasattr(sys.stderr, 'reconfigure'):
        sys.stderr.reconfigure(encoding='utf-8', errors='replace')

from pulse.core.brain_utils import HIPPOCAMPUS_DIR, STATE_DIR
from pulse.tasks.pulse_dispatcher import TaskDispatcher
from pulse.admin.retroactive.retroactive_llm_extractor import (
    batch_interactions, SCRIBE_PROMPT, CAPTURE_LOGS, BRAIN_FILES,
)

INTENTS_DIR = HIPPOCAMPUS_DIR / "Intents"
WANTS_FILE = INTENTS_DIR / "wants.json"
DONT_WANTS_FILE = INTENTS_DIR / "dont_wants.json"
BRIDGE_STATE = STATE_DIR / "pulse_bridge_state.json"
EXTRACTION_BATCH_SIZE = 15
INTENT_BATCH_SIZE = 80


def _load_bridge_positions() -> dict[str, int]:
    """Load dispatch positions from bridge state (handles nested positions dict)."""
    if not BRIDGE_STATE.exists():
        return {}
    try:
        data = json.loads(BRIDGE_STATE.read_text(encoding="utf-8"))
        # P47 FIX: bridge state is {"positions": {agent: byte_offset}} not flat
        if isinstance(data, dict) and "positions" in data:
            data = data["positions"]
        return {k: v for k, v in data.items() if isinstance(v, int)}
    except Exception:
        return {}


def _save_dispatch_position(log_name: str, position: int):
    from pulse.core.brain_utils import atomic_update_json
    def _update(data):
        if not isinstance(data, dict):
            data = {}
        # Write into the "positions" sub-dict so _load_bridge_positions() can find it
        if "positions" not in data or not isinstance(data["positions"], dict):
            data["positions"] = {}
        data["positions"][f"dispatch_{log_name}"] = position
        return data
    atomic_update_json(BRIDGE_STATE, _update, default={})


def _make_task(task_id: str, title: str, desc: str, memory_data: dict) -> dict:
    return {"task_id": task_id, "title": title, "description": desc,
            "domain": "brain", "complexity": "low", "recommended_tier": "fast",
            "required_capabilities": [], "depends_on": [], "status": "open",
            "claimed_by": None, "claimed_at": None, "files_touched": [],
            "outcome": None, "memory_data": memory_data}


def dispatch_extraction_tasks(log_path: Path, batch_size: int = EXTRACTION_BATCH_SIZE,
                              dry_run: bool = False) -> int:
    """Chunk a capture log and post EXTRACT_KNOWLEDGE tasks to the board."""
    if not log_path.exists():
        return 0
    agent = log_path.stem.replace("pulse_captured_", "")
    positions = _load_bridge_positions()
    last_pos = positions.get(f"dispatch_{log_path.name}", 0)
    current_size = log_path.stat().st_size
    if current_size <= last_pos:
        print(f"  [{agent}] No new data (pos {last_pos}/{current_size})")
        return 0

    MAX_EXTRACTION_TASKS = 10  # Cap per dispatch to prevent terminal flood
    tasks, ts_prefix = [], int(time.time())
    for batch_text, count, first_ts in batch_interactions(log_path, batch_size, start_offset=last_pos):
        if len(tasks) >= MAX_EXTRACTION_TASKS:
            break
        n = len(tasks) + 1
        tasks.append(_make_task(
            f"t_{ts_prefix}_{n:03d}",
            f"[EXTRACT_KNOWLEDGE] {agent} batch {n} ({count} interactions)",
            f"Extract knowledge from {count} interactions. Source: {log_path.name}",
            {"type": "EXTRACT_KNOWLEDGE", "source_log": str(log_path), "agent_name": agent,
             "batch_text": batch_text, "batch_num": n, "interaction_count": count,
             "timestamp": first_ts or datetime.now(timezone.utc).isoformat()}))
    if not tasks:
        return 0
    if dry_run:
        print(f"  [{agent}] Would post {len(tasks)} EXTRACT_KNOWLEDGE tasks")
        return len(tasks)
    dispatcher = TaskDispatcher(verbose=False)
    dispatcher.dispatch(
        prompt=f"[MEMORY_DISPATCHER] Extract knowledge from {agent} ({len(tasks)} batches)",
        tasks=tasks, founder_override=True)
    _save_dispatch_position(log_path.name, current_size)
    print(f"  [{agent}] Posted {len(tasks)} tasks")
    return len(tasks)


MAX_INTENT_TASKS_PER_DISPATCH = 10  # Cap to prevent terminal flood (was 7,384)


def dispatch_intent_cleaning(file_path: Path, key: str,
                             batch_size: int = INTENT_BATCH_SIZE, dry_run: bool = False) -> int:
    """Chunk intents from SQLite and post CLEAN_INTENTS tasks to the board."""
    # JSON Kill List: read from SQLite instead of JSON file
    try:
        if key == "want":
            from pulse.core.brain_db import get_wants
            data = get_wants(limit=5000)
            texts = [e.get("want", "") for e in data if e.get("want")]
        elif key == "dont_want":
            from pulse.core.brain_db import get_dont_wants
            data = get_dont_wants(limit=5000)
            texts = [e.get("dont_want", "") for e in data if e.get("dont_want")]
        else:
            texts = []
    except Exception:
        return 0
    if not texts:
        return 0

    label = file_path.stem if hasattr(file_path, 'stem') else key
    tasks, ts_prefix = [], int(time.time())
    for i in range(0, len(texts), batch_size):
        if len(tasks) >= MAX_INTENT_TASKS_PER_DISPATCH:
            break
        batch = texts[i : i + batch_size]
        n = i // batch_size + 1
        tasks.append(_make_task(
            f"t_{ts_prefix}_{n:03d}",
            f"[CLEAN_INTENTS] {label} batch {n} ({len(batch)} entries)",
            f"Distill {len(batch)} raw {key} entries into clean preferences.",
            {"type": "CLEAN_INTENTS", "file": str(file_path), "key": key,
             "batch": batch, "batch_num": n,
             "intent_type": "want" if key == "want" else "dont_want"}))
    if dry_run:
        print(f"  [{label}] Would post {len(tasks)} CLEAN_INTENTS tasks")
        return len(tasks)
    dispatcher = TaskDispatcher(verbose=False)
    dispatcher.dispatch(
        prompt=f"[MEMORY_DISPATCHER] Clean {label} ({len(tasks)} batches)",
        tasks=tasks, founder_override=True)
    print(f"  [{label}] Posted {len(tasks)} tasks")
    return len(tasks)


def memory_status() -> dict:
    """Show pending memory tasks on the board."""
    from pulse.core.brain_utils import TASK_BOARD_FILE
    if not TASK_BOARD_FILE.exists():
        return {}
    board = json.loads(TASK_BOARD_FILE.read_text(encoding="utf-8"))
    counts: dict[str, int] = {}
    for d in board.get("dispatches", []):
        for t in d.get("tasks", []):
            title, status = t.get("title", ""), t.get("status", "open")
            if "[EXTRACT_KNOWLEDGE]" in title:
                counts[f"extract_{status}"] = counts.get(f"extract_{status}", 0) + 1
            elif "[CLEAN_INTENTS]" in title:
                counts[f"clean_{status}"] = counts.get(f"clean_{status}", 0) + 1
    return counts


def dispatch_all(dry_run: bool = False, extraction: bool = True, intents: bool = True) -> dict:
    """Dispatch all pending memory work."""
    results = {"extraction_tasks": 0, "intent_tasks": 0}
    if extraction:
        print("[1] Dispatching extraction tasks...")
        for lp in CAPTURE_LOGS:
            results["extraction_tasks"] += dispatch_extraction_tasks(lp, dry_run=dry_run)
    if intents:
        print("[2] Dispatching intent cleaning tasks...")
        results["intent_tasks"] += dispatch_intent_cleaning(WANTS_FILE, "want", dry_run=dry_run)
        results["intent_tasks"] += dispatch_intent_cleaning(DONT_WANTS_FILE, "dont_want", dry_run=dry_run)
    return results


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Brain V4 Memory Dispatcher")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--extraction", action="store_true")
    parser.add_argument("--intents", action="store_true")
    parser.add_argument("--status", action="store_true")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose output")
    args = parser.parse_args()

    if args.status:
        for k, v in memory_status().items():
            print(f"  {k}: {v}")
        return

    do_ext = args.extraction or (not args.extraction and not args.intents)
    do_int = args.intents or (not args.extraction and not args.intents)
    print(f"{'='*60}\nBRAIN V4 — Memory Dispatcher\n{'='*60}")
    if args.dry_run:
        print("MODE: DRY RUN\n")
    results = dispatch_all(dry_run=args.dry_run, extraction=do_ext, intents=do_int)
    print(f"\n{'='*60}")
    print(f"DISPATCHED: {results['extraction_tasks']} extraction + {results['intent_tasks']} intent")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()
