# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
D26: JIT Synaptic Inoculation — Preemptive Weakness Remediation.

Monitors task claims in near-real-time. When an agent claims a task in
a domain where it's historically weak, preemptively searches the brain
and injects remedial knowledge into the hot cache BEFORE work begins.

Source: DMN session_20260311T163020 | Usage: pulse inoculate
"""

import json
import logging
from datetime import datetime, timezone
from pathlib import Path

_log = logging.getLogger("pulse.d26_jit_inoculation")

try:
    from pulse.core.brain_utils import STATE_DIR, TASK_BOARD_FILE
except Exception:
    STATE_DIR = Path("state")
    TASK_BOARD_FILE = STATE_DIR / "task_board.json"

METRICS_FILE = STATE_DIR / "health" / "jit_inoculation.json"
_BOOST_SALIENCE = 6.0      # boosted salience for inoculation items
_MAX_INJECT = 5            # max hot cache items per inoculation
_MAX_EVENTS = 100          # cap recent events list


def _load_task_board() -> list[dict]:
    """Load active tasks from task board."""
    if not TASK_BOARD_FILE.exists():
        return []
    try:
        data = json.loads(TASK_BOARD_FILE.read_text(encoding="utf-8"))
        tasks = []
        for d in data.get("dispatches", []):
            tasks.extend(d.get("tasks", []))
        return tasks
    except (json.JSONDecodeError, OSError):
        return []


def _scan_new_claims(tasks: list[dict], seen: dict) -> list[dict]:
    """Find tasks recently claimed (status=active) not yet inoculated."""
    claims = []
    for t in tasks:
        tid = t.get("task_id", "")
        status = t.get("status", "")
        agent = t.get("claimed_by", "")
        if status != "active" or not agent or not tid:
            continue
        if tid in seen:
            continue
        claims.append({
            "task_id": tid,
            "agent": agent,
            "domain": t.get("domain", "general"),
            "title": t.get("title", "")[:120],
        })
    return claims


def _check_weakness(agent_id: str, domain: str) -> bool:
    """Check if agent has a weak domain."""
    try:
        from pulse.brain.metacognition.agent_procedures import get_weak_domains
        weak = get_weak_domains(agent_id)
        return any(w["domain"] == domain for w in weak)
    except Exception:
        return False


def _search_remedial(domain: str, title: str) -> list[dict]:
    """Search brain for remedial knowledge in the weak domain."""
    try:
        from pulse.brain.brain_search import brain_search
        query = f"{domain} best practices common failures {title}"
        results = brain_search(query, max_per_source=3)
        # Flatten all sources into a single ranked list
        items = []
        for source_items in results.get("results", {}).values():
            if isinstance(source_items, list):
                items.extend(source_items)
        # Sort by relevance * salience, take top N
        items.sort(key=lambda x: x.get("relevance", 0) * x.get("salience", 1),
                   reverse=True)
        return items[:_MAX_INJECT]
    except Exception as e:
        _log.warning("Brain search failed for domain %s: %s", domain, e)
        return []


def _inject_to_hot_cache(items: list[dict], agent_id: str,
                          domain: str) -> int:
    """Inject remedial items into hot cache with boosted salience."""
    try:
        from pulse.brain.memory.hot_cache import write_hot
    except ImportError:
        return 0
    injected = 0
    now = datetime.now(timezone.utc).isoformat()
    for item in items:
        entry = {
            "content": item.get("text", item.get("content", ""))[:300],
            "title": item.get("title", "")[:80],
            "type": item.get("type", "lesson"),
            "salience": _BOOST_SALIENCE,
            "confidence": item.get("confidence", 0.7),
            "source": f"jit_inoculation_{agent_id}",
            "domain": domain,
            "timestamp": now,
        }
        if write_hot(entry):
            injected += 1
    return injected


def _load_state() -> dict:
    if not METRICS_FILE.exists():
        return {"runs": 0, "inoculations": 0, "items_injected": 0,
                "seen_claims": {}, "recent_events": []}
    try:
        return json.loads(METRICS_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {"runs": 0, "inoculations": 0, "items_injected": 0,
                "seen_claims": {}, "recent_events": []}


def _save_state(state: dict):
    import os
    METRICS_FILE.parent.mkdir(parents=True, exist_ok=True)
    tmp = METRICS_FILE.with_suffix(".tmp")
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2, ensure_ascii=False)
        os.replace(str(tmp), str(METRICS_FILE))
    except OSError:
        if tmp.exists():
            tmp.unlink()


def run_once(verbose: bool = False) -> dict:
    """Run one inoculation cycle: scan claims → check weakness → inject."""
    state = _load_state()
    seen = state.get("seen_claims", {})

    tasks = _load_task_board()
    claims = _scan_new_claims(tasks, seen)
    if not claims:
        if verbose:
            _log.info("No new task claims to check")
        state["runs"] = state.get("runs", 0) + 1
        state["last_run"] = datetime.now(timezone.utc).isoformat()
        _save_state(state)
        return {"claims": 0, "inoculated": 0, "status": "no_claims"}

    if verbose:
        _log.info("Checking %d new claim(s) for weakness", len(claims))

    inoculated = 0
    total_injected = 0
    done_pairs = set()  # (agent, domain) dedup per cycle

    for claim in claims:
        agent = claim["agent"]
        domain = claim["domain"]
        pair = (agent, domain)
        seen[claim["task_id"]] = datetime.now(timezone.utc).isoformat()

        if pair in done_pairs:
            continue

        # Check if agent is weak in this domain
        # Extract base agent ID (strip _NN suffix)
        import re
        base_agent = re.sub(r'_\d+$', '', agent)
        if not _check_weakness(base_agent, domain):
            continue

        if verbose:
            _log.info("  [WEAK] %s in domain=%s — searching remedial...",
                      agent, domain)

        items = _search_remedial(domain, claim["title"])
        if not items:
            continue

        injected = _inject_to_hot_cache(items, agent, domain)
        if injected > 0:
            inoculated += 1
            total_injected += injected
            done_pairs.add(pair)
            if verbose:
                _log.info("    Injected %d items into hot cache", injected)

            # Record event
            events = state.get("recent_events", [])
            events.append({
                "agent": agent, "domain": domain,
                "items": injected,
                "ts": datetime.now(timezone.utc).isoformat(),
            })
            state["recent_events"] = events[-_MAX_EVENTS:]

    # Cap seen_claims to prevent unbounded growth
    if len(seen) > 2000:
        sorted_keys = sorted(seen.keys(), key=lambda k: seen[k])
        seen = dict(list(zip(sorted_keys, [seen[k] for k in sorted_keys]))[-1000:])

    state["seen_claims"] = seen
    state["inoculations"] = state.get("inoculations", 0) + inoculated
    state["items_injected"] = state.get("items_injected", 0) + total_injected
    state["runs"] = state.get("runs", 0) + 1
    state["last_run"] = datetime.now(timezone.utc).isoformat()
    _save_state(state)

    return {"claims": len(claims), "inoculated": inoculated,
            "injected": total_injected, "status": "ok"}


def cmd_inoculate(args: list):
    """CLI entry point for `pulse inoculate`."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    if args and args[0] == "status":
        s = _load_state()
        print(f"[D26] JIT Inoculation — Runs: {s.get('runs', 0)}, "
              f"Inoculations: {s.get('inoculations', 0)}, "
              f"Items injected: {s.get('items_injected', 0)}")
        events = s.get("recent_events", [])
        if events:
            print(f"  Recent: {len(events)} events")
            for e in events[-3:]:
                print(f"    {e['agent']} → {e['domain']} ({e['items']} items)")
        print(f"  Last run: {s.get('last_run', 'never')}")
        return
    result = run_once(verbose=True)
    print(json.dumps(result, indent=2))


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