#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""D12: Hebbian Swarm Balancer — Model-Task Adaptive Learning.
SWEEP → AGGREGATE → LEARN → EXPOSE. 15min via orchestrator.
"""
from __future__ import annotations
import io, json, os, re, sys, time
from datetime import datetime, timezone
from pathlib import Path

if os.name == "nt" and __name__ == "__main__":
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8")

from pulse.core.paths import get_root_dir
from pulse.core.brain_utils import now_iso
from pulse.core.brain_io import _acquire

ROOT = get_root_dir()
STATE_FILE = ROOT / "state" / "hebbian_matrix.json"
TASK_BOARD = ROOT / "state" / "task_board.json"
COMPLETED_DIR = ROOT / "state" / "completed"
REGISTRY_FILE = ROOT / "state" / "agent_registry.json"

LEARNING_RATE = 0.15        # eta — pull weight toward observed rate
DECAY_RATE = 0.02           # LTD — per-sweep decay for stale pairs
MIN_OBSERVATIONS = 3        # need 3+ outcomes before weight is trusted
COLD_START_WEIGHT = 0.5     # default for unseen pairs
WEIGHT_FLOOR = 0.1          # never drop below (allows recovery)
WEIGHT_CEILING = 0.95       # never go above (always some uncertainty)
STALE_DAYS = 7              # LTD kicks in after 7 days with no events
PROCESSED_CAP = 5000        # rolling cap on processed_tasks

_DEFAULT_STATE = {"matrix": {}, "processed_tasks": [], "last_sweep": None,
                  "runs": 0, "last_elapsed_s": 0}
_WORKER_SUFFIX_RE = re.compile(r"^(.+?)_\d+$")


def _load_state() -> dict:
    if not STATE_FILE.exists():
        return dict(_DEFAULT_STATE)
    try:
        with _acquire(STATE_FILE):
            return json.loads(STATE_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return dict(_DEFAULT_STATE)


def _save_state(state: dict):
    STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
    pt = state.get("processed_tasks", [])
    if len(pt) > PROCESSED_CAP:
        state["processed_tasks"] = pt[-PROCESSED_CAP:]
    with _acquire(STATE_FILE):
        STATE_FILE.write_text(json.dumps(state, indent=2, ensure_ascii=False), encoding="utf-8")


def _load_json(path: Path) -> dict | list:
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {}


def _resolve_model(worker_id: str, registry: dict) -> str | None:
    """Resolve worker_id to model name. Registry preferred, regex fallback."""
    if not worker_id:
        return None
    agent = registry.get(worker_id, {})
    model = (agent.get("model") or "").strip()
    if model:
        return model
    if worker_id.startswith("worker_"):
        return None
    m = _WORKER_SUFFIX_RE.match(worker_id)
    if m:
        candidate = m.group(1).strip()
        if candidate and not candidate.startswith("worker"):
            return candidate
    return worker_id


def _sweep_outcomes(state: dict) -> tuple[list[dict], set[str]]:
    """Scan board + archives for new completed tasks."""
    registry = _load_json(REGISTRY_FILE)
    if isinstance(registry, list):
        registry = {}
    processed = set(state.get("processed_tasks", []))
    events, new_ids = [], set()

    def _process(task: dict):
        tid = task.get("task_id", "")
        if not tid or tid in processed or tid in new_ids:
            return
        status = (task.get("status") or "").lower()
        if status not in ("done", "failed", "escalated"):
            return
        wid = task.get("completed_by", "")
        model = _resolve_model(wid, registry) if wid else None
        if not model:
            return
        domain = (task.get("domain") or "general").lower().strip()
        success = (status == "done") and ("POISON_PILL" not in (task.get("outcome") or ""))
        events.append({"model": model, "domain": domain, "success": success})
        new_ids.add(tid)

    board = _load_json(TASK_BOARD)
    if isinstance(board, dict):
        for d in board.get("dispatches", []):
            for t in d.get("tasks", []):
                _process(t)
    if COMPLETED_DIR.exists():
        for af in COMPLETED_DIR.iterdir():
            if af.suffix != ".json":
                continue
            try:
                arch = json.loads(af.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                continue
            tasks = arch.get("tasks", []) if isinstance(arch, dict) else []
            if isinstance(arch, list):
                for item in arch:
                    tasks.extend(item.get("tasks", []) if isinstance(item, dict) else [])
            for t in tasks:
                _process(t)
    return events, new_ids


def _aggregate(matrix: dict, events: list[dict]):
    """Count successes/failures per model x domain."""
    for ev in events:
        model, domain = ev["model"], ev["domain"]
        entry = matrix.setdefault(model, {}).setdefault(domain, {
            "successes": 0, "failures": 0, "weight": COLD_START_WEIGHT, "last_event": None})
        entry["successes" if ev["success"] else "failures"] += 1
        entry["last_event"] = now_iso()


def _apply_hebbian(matrix: dict):
    """Hebbian update: pull weights toward observed success rate + LTD decay."""
    now = datetime.now(timezone.utc)
    for model in matrix:
        for domain in matrix[model]:
            e = matrix[model][domain]
            total = e["successes"] + e["failures"]
            if total >= MIN_OBSERVATIONS:
                e["weight"] += LEARNING_RATE * (e["successes"] / total - e["weight"])
                e["weight"] = max(WEIGHT_FLOOR, min(WEIGHT_CEILING, e["weight"]))
            le = e.get("last_event")
            if le:
                try:
                    days = (now - datetime.fromisoformat(le.replace("Z", "+00:00"))).total_seconds() / 86400
                    if days > STALE_DAYS:
                        e["weight"] = max(WEIGHT_FLOOR, e["weight"] * (1 - DECAY_RATE))
                except (ValueError, TypeError):
                    pass


def _matrix_stats(matrix: dict) -> tuple[int, int, list[tuple]]:
    """Returns (n_pairs, n_obs, ranked_entries)."""
    n_pairs = sum(len(d) for d in matrix.values())
    n_obs = sum(e["successes"] + e["failures"] for d in matrix.values() for e in d.values())
    entries = [(m, d, e["weight"], e["successes"] + e["failures"])
               for m, doms in matrix.items() for d, e in doms.items()
               if e["successes"] + e["failures"] >= MIN_OBSERVATIONS]
    entries.sort(key=lambda x: x[2], reverse=True)
    return n_pairs, n_obs, entries


def get_model_fitness(model: str, domain: str) -> float:
    """Get Hebbian weight for model x domain. Returns COLD_START_WEIGHT if unseen."""
    if not STATE_FILE.exists():
        return COLD_START_WEIGHT
    try:
        data = json.loads(STATE_FILE.read_text(encoding="utf-8"))
        return data.get("matrix", {}).get(model, {}).get(domain, {}).get("weight", COLD_START_WEIGHT)
    except (json.JSONDecodeError, OSError):
        return COLD_START_WEIGHT


def get_model_ranking(domain: str, tier: str | None = None) -> list[tuple[str, float]]:
    """Rank models by Hebbian fitness for a domain. Returns [(model, weight), ...] desc."""
    if not STATE_FILE.exists():
        return []
    try:
        data = json.loads(STATE_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return []
    rankings = []
    for model, domains in data.get("matrix", {}).items():
        if domain in domains:
            e = domains[domain]
            total = e.get("successes", 0) + e.get("failures", 0)
            if total >= MIN_OBSERVATIONS:
                rankings.append((model, e.get("weight", COLD_START_WEIGHT)))
    rankings.sort(key=lambda x: x[1], reverse=True)
    return rankings


def run_once(verbose: bool = False, dry_run: bool = False) -> dict:
    """Full D12 cycle: SWEEP -> AGGREGATE -> LEARN -> STORE."""
    state, t0 = _load_state(), time.time()
    vp = print if verbose else lambda *a: None
    events, new_ids = _sweep_outcomes(state)
    vp(f"[D12] Swept {len(events)} new outcome(s) from {len(new_ids)} task(s)")
    if not events:
        state["runs"] = state.get("runs", 0) + 1
        state.update(last_sweep=datetime.now(timezone.utc).isoformat(),
                     last_elapsed_s=round(time.time() - t0, 1))
        if not dry_run:
            _save_state(state)
        return {"events": 0, "models": 0, "pairs": 0}
    matrix = state.get("matrix", {})
    _aggregate(matrix, events)
    vp(f"[D12] Aggregated into {len(matrix)} model(s)")
    _apply_hebbian(matrix)
    n_pairs, n_obs, _ = _matrix_stats(matrix)
    vp(f"[D12] {n_pairs} model x domain pairs, {n_obs} total observations")
    if dry_run:
        vp("[D12] Dry run — state not saved")
        return {"events": len(events), "models": len(matrix), "pairs": n_pairs}
    state["matrix"] = matrix
    state["processed_tasks"] = list(set(state.get("processed_tasks", [])) | new_ids)
    state["runs"] = state.get("runs", 0) + 1
    state.update(last_sweep=datetime.now(timezone.utc).isoformat(),
                 last_elapsed_s=round(time.time() - t0, 1))
    _save_state(state)
    vp(f"[D12] Done in {state['last_elapsed_s']}s")
    return {"events": len(events), "models": len(matrix), "pairs": n_pairs}


def show_status():
    """Print D12 Hebbian matrix stats + top model-domain pairs."""
    s = _load_state()
    matrix = s.get("matrix", {})
    last = (s.get("last_sweep") or "never")[:19]
    n_pairs, n_obs, ranked = _matrix_stats(matrix)
    print(f"\n  D12 Hebbian Swarm Balancer\n  {'=' * 40}")
    for k, v in [("Runs", s.get("runs", 0)), ("Last sweep", last),
                 ("Models", f"{len(matrix)} tracked"), ("Pairs", f"{n_pairs} model x domain"),
                 ("Observations", f"{n_obs:,} task outcomes"),
                 ("Processed", f"{len(s.get('processed_tasks', []))} task IDs")]:
        print(f"  {k + ':':<16} {v}")
    if s.get("last_elapsed_s"):
        print(f"  {'Elapsed:':<16} {s['last_elapsed_s']}s")
    if ranked:
        print(f"\n  Top Performers:")
        for m, d, w, t in ranked[:5]:
            print(f"    {m[:24]:<26} {d:<16} {w:.2f}  ({t} obs)")
        weak = [e for e in ranked if e[2] < 0.4]
        if weak:
            print(f"\n  Weak Pairings:")
            for m, d, w, t in weak[:3]:
                print(f"    {m[:24]:<26} {d:<16} {w:.2f}  ({t} obs)")
    print()


def cmd_hebbian(args: list):
    """CLI entry: pulse hebbian [--dry-run] [-v] [status]."""
    if args and args[0] == "status":
        return show_status()
    verbose = "-v" in args or "--verbose" in args
    result = run_once(verbose=verbose, dry_run="--dry-run" in args)
    if not verbose:
        print(f"[D12] {result['events']} events, {result['models']} models, {result['pairs']} pairs")


if __name__ == "__main__":
    cmd_hebbian(sys.argv[1:])
