#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""D8 Swarm Reflex — detect struggling agents, push brain knowledge via broadcast.
Covers workers + CLI + API. Scheduled every 5min. CLI: --once --dry-run --status -v"""
from __future__ import annotations

import json
import logging
import os
import signal
import sys
import time
from datetime import datetime, timezone, timedelta
from pathlib import Path

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

ROOT_DIR = Path(__file__).resolve().parents[4]
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

from pulse.core.brain_utils import (
    STATE_DIR, TASK_BOARD_FILE, AGENT_REGISTRY_FILE,
    HIPPOCAMPUS_DIR, load_json_dict, now_iso as _now_iso,
)
from pulse.core.brain_io import _acquire, save_json

log = logging.getLogger("pulse.d8_reflex")

STALL_THRESHOLD = 180       # 3 min = stalled
RETRY_THRESHOLD = 2         # retry_count >= 2 = struggling
GRACE_PERIOD = 60           # skip agents < 60s old
ALERT_TTL_HOURS = 2         # reflex alerts expire after 2h
POLL_INTERVAL = 5 * 60      # 5 minutes
STATE_FILE = STATE_DIR / "swarm_reflex_state.json"
BROADCAST_FILE = HIPPOCAMPUS_DIR / "Evidence" / "Sessions" / "broadcast_alerts.json"
HEARTBEAT_DIR = STATE_DIR / "heartbeats"


def _load_state() -> dict:
    if STATE_FILE.exists():
        return load_json_dict(STATE_FILE)
    return {"total_scans": 0, "total_pushes": 0, "pushed_set": [],
            "last_run": None}


def _save_state(state: dict):
    # Cap pushed_set to last 200 entries
    state["pushed_set"] = state.get("pushed_set", [])[-200:]
    with _acquire(STATE_FILE):
        save_json(STATE_FILE, state)


def _heartbeat_age(agent_id: str, registry_entry: dict) -> float:
    hb_file = HEARTBEAT_DIR / f"{agent_id}.json"
    if hb_file.exists():
        try:
            age = time.time() - hb_file.stat().st_mtime
            return age
        except OSError:
            pass
    # Fallback: registry timestamp
    ts_str = registry_entry.get("last_heartbeat", "")
    if ts_str:
        try:
            ts = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
            return (datetime.now(timezone.utc) - ts).total_seconds()
        except (ValueError, TypeError):
            pass
    return 9999.0  # unknown = very stale


def _build_retry_map(board: dict) -> dict:
    result: dict[str, list] = {}
    for dispatch in board.get("dispatches", []):
        for t in dispatch.get("tasks", []):
            agent = t.get("claimed_by") or t.get("completed_by", "")
            retry = t.get("retry_count", 0)
            if not agent or retry < RETRY_THRESHOLD:
                continue
            if t.get("status") not in ("open", "active", "done"):
                continue
            result.setdefault(agent, []).append({
                "task_id": t.get("task_id", ""),
                "title": t.get("title", "")[:80],
                "domain": t.get("domain", "general"),
                "retry_count": retry,
            })
    return result


def _prune_old_reflex_alerts(alerts: list) -> list:
    cutoff = datetime.now(timezone.utc) - timedelta(hours=ALERT_TTL_HOURS)
    kept = []
    for a in alerts:
        if a.get("source") != "d8_reflex":
            kept.append(a)
            continue
        ts_str = a.get("timestamp", "")
        try:
            ts = datetime.fromisoformat(ts_str.replace("Z", "+00:00"))
            if ts >= cutoff:
                kept.append(a)
        except (ValueError, TypeError):
            kept.append(a)
    return kept


def _write_broadcast_alerts(new_alerts: list):
    BROADCAST_FILE.parent.mkdir(parents=True, exist_ok=True)
    with _acquire(BROADCAST_FILE):
        existing = {}
        if BROADCAST_FILE.exists():
            try:
                existing = json.loads(BROADCAST_FILE.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                existing = {}
        alerts = existing.get("alerts", [])
        alerts = _prune_old_reflex_alerts(alerts)
        alerts.extend(new_alerts)
        existing["alerts"] = alerts
        existing["last_updated"] = _now_iso()
        BROADCAST_FILE.write_text(
            json.dumps(existing, indent=2, ensure_ascii=False), encoding="utf-8")


def _brain_hint_for(query: str) -> str:
    try:
        from pulse.brain.brain_search import brain_search
        results = brain_search(query, max_per_source=2)
        hits = results.get("results", {})
        parts = []
        for src in ("anti_patterns", "fails", "wins", "patterns"):
            for h in hits.get(src, [])[:1]:
                text = (h.get("title") or h.get("text", ""))[:80]
                if text:
                    parts.append(text)
        return ". ".join(parts)[:200] if parts else ""
    except Exception:
        return ""


def detect_struggling(verbose: bool = False) -> list[dict]:
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    board = load_json_dict(TASK_BOARD_FILE)
    retry_map = _build_retry_map(board)
    struggles = []

    for agent_id, info in registry.items():
        if not isinstance(info, dict):
            continue
        age = _heartbeat_age(agent_id, info)
        if age < GRACE_PERIOD:
            continue  # just started

        current_task = info.get("current_task")
        agent_type = info.get("type", "agent")

        # Check 1: Retry struggle — agent has tasks with high retry_count
        if agent_id in retry_map:
            for task_info in retry_map[agent_id]:
                struggles.append({
                    "agent_id": agent_id, "agent_type": agent_type,
                    "struggle_type": "retry",
                    "domain": task_info["domain"],
                    "task_id": task_info["task_id"],
                    "task_title": task_info["title"],
                    "retry_count": task_info["retry_count"],
                    "context": f"retry_count={task_info['retry_count']}",
                })

        # Check 2: Stall struggle — heartbeat stale with active task
        if age > STALL_THRESHOLD and current_task:
            struggles.append({
                "agent_id": agent_id, "agent_type": agent_type,
                "struggle_type": "stall",
                "domain": info.get("history", {}).get("domains", {}) and "general",
                "task_id": current_task,
                "task_title": f"stalled task {current_task}",
                "retry_count": 0,
                "context": f"heartbeat_age={int(age)}s",
            })

    if verbose and struggles:
        log.info("[D8] Detected %d struggling agent(s)", len(struggles))
        for s in struggles[:5]:
            log.info("  %s [%s] %s: %s", s["agent_id"], s["struggle_type"],
                     s["domain"], s["context"])
    return struggles


def push_reflex_alerts(struggles: list[dict], dry_run: bool = False,
                       verbose: bool = False) -> dict:
    state = _load_state()
    pushed_set = set(state.get("pushed_set", []))
    pushed = skipped = 0
    new_alerts = []

    for s in struggles:
        key = f"{s['agent_id']}|{s.get('task_id', '')}"
        if key in pushed_set:
            skipped += 1
            continue

        query = f"{s['domain']} {s['task_title']}"
        hint = _brain_hint_for(query)
        if not hint:
            skipped += 1
            continue

        failure_text = (
            f"[D8 REFLEX] {s['agent_id']} struggling on {s['domain']} "
            f"({s['context']}). Brain: {hint}")

        if dry_run:
            log.info("  [DRY-RUN] Would push: %s", failure_text[:100])
            pushed += 1
            continue

        new_alerts.append({
            "agent": s["agent_id"],
            "failure": failure_text[:200],
            "timestamp": _now_iso(),
            "source": "d8_reflex",
        })
        pushed_set.add(key)
        pushed += 1

    if new_alerts:
        _write_broadcast_alerts(new_alerts)
        if verbose:
            log.info("[D8] Pushed %d reflex alert(s) to broadcast", len(new_alerts))

    if not dry_run:
        state["pushed_set"] = list(pushed_set)
        _save_state(state)

    return {"pushed": pushed, "skipped": skipped}


def get_reflex_stats() -> dict:
    return _load_state()


def run_once(dry_run: bool = False, verbose: bool = False) -> dict:
    struggles = detect_struggling(verbose=verbose)
    push = push_reflex_alerts(struggles, dry_run=dry_run, verbose=verbose)
    if not dry_run:
        state = _load_state()
        state["total_scans"] = state.get("total_scans", 0) + 1
        state["total_pushes"] = state.get("total_pushes", 0) + push["pushed"]
        state["last_run"] = _now_iso()
        _save_state(state)
    return {"struggles": len(struggles), "push": push}


def run_daemon(verbose: bool = False, dry_run: bool = False):
    _running = True

    def _shutdown(*_):
        nonlocal _running
        _running = False

    signal.signal(signal.SIGINT, _shutdown)
    signal.signal(signal.SIGTERM, _shutdown)
    log.info("[D8] Swarm Reflex started (interval=%ds)", POLL_INTERVAL)
    while _running:
        try:
            result = run_once(dry_run=dry_run, verbose=verbose)
            if verbose:
                log.info("[D8] Cycle: %s", result)
        except Exception as e:
            log.error("[D8] Cycle error: %s", e)
        for _ in range(POLL_INTERVAL):
            if not _running:
                break
            time.sleep(1)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [D8] %(message)s", datefmt="%H:%M:%S")
    _v, _dr = "--verbose" in sys.argv or "-v" in sys.argv, "--dry-run" in sys.argv
    if "--once" in sys.argv or _dr:
        r = run_once(dry_run=_dr, verbose=_v or _dr)
        p = r["push"]
        print(f"\n[D8] {r['struggles']} struggling, pushed {p['pushed']}, skipped {p['skipped']}")
    elif "--status" in sys.argv:
        st = get_reflex_stats()
        print(f"  D8 Swarm Reflex — last: {st.get('last_run', 'never')}  "
              f"Scans: {st.get('total_scans', 0)} | Pushes: {st.get('total_pushes', 0)}")
    else:
        print("[D8] Swarm Reflex: --once | --dry-run | --status | -v")
