# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
D21: Self-Correcting Neuroplasticity Engine — Closed-Loop Learning.

Detects agents struggling in specific domains, triggers targeted extraction
from their capture logs, boosts the extracted lessons, and injects them into
brain + hot cache so agents auto-correct on their next boot.

Source: DMN session_20260227T033544 | Usage: pulse neuroplasticity
"""

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

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

from pulse.core.brain_utils import STATE_DIR, LESSONS_JSONL, FAILS_JSONL

METRICS_FILE = STATE_DIR / "health" / "neuroplasticity_metrics.json"
_BOOST_SALIENCE = 6.0         # boosted salience for extracted remedial lessons
_BOOST_CONFIDENCE = 0.75      # baseline confidence for remedial items
_MIN_TASKS_FOR_WEAK = 3       # minimum tasks before judging weakness
_WEAK_THRESHOLD = 0.5         # success rate below this = weak domain
_MAX_EXTRACT_LINES = 500      # max capture log lines to scan per agent
_SCRIBE_PROMPT = """You are a knowledge extractor. From the chat transcript below,
extract lessons and failure patterns specifically about the domain: {domain}.

Return JSON: {{"lessons": ["..."], "failures": ["..."]}}
Only include items about {domain}. Be specific and actionable. Max 5 items each."""


def _detect_weak_agents() -> list[dict]:
    """Find agents with weak domains from registry."""
    try:
        from pulse.brain.metacognition.agent_procedures import get_weak_domains
        reg_file = STATE_DIR / "agent_registry.json"
        if not reg_file.exists():
            return []
        data = json.loads(reg_file.read_text(encoding="utf-8"))
        weak_agents = []
        for agent_id, info in data.items():
            if not isinstance(info, dict):
                continue
            history = info.get("history", {})
            if history.get("tasks_completed", 0) < _MIN_TASKS_FOR_WEAK:
                continue
            weak = get_weak_domains(agent_id, threshold=_WEAK_THRESHOLD)
            if weak:
                weak_agents.append({
                    "agent_id": agent_id,
                    "model": info.get("model", "unknown"),
                    "weak_domains": weak[:3],  # top 3 weakest
                })
        return weak_agents
    except Exception as e:
        _log.warning("Failed to detect weak agents: %s", e)
        return []


def _find_capture_log(agent_id: str) -> Path | None:
    """Find the capture log for an agent."""
    # Workers write to pulse_captured_<worker_id>.log
    log = STATE_DIR / f"pulse_captured_{agent_id}.log"
    if log.exists():
        return log
    # Try agent type prefix (e.g., worker_gemini_123 → pulse_captured_gemini.log)
    for part in agent_id.split("_"):
        if part in ("claude", "gemini", "codex", "grok"):
            alt = STATE_DIR / f"pulse_captured_{part}.log"
            if alt.exists():
                return alt
    return None


def _extract_domain_knowledge(log_path: Path, domain: str) -> dict:
    """Extract domain-specific lessons from capture log via LLM."""
    try:
        with open(log_path, "r", encoding="utf-8", errors="replace") as f:
            lines = f.readlines()
    except OSError:
        return {"lessons": [], "failures": []}

    # Take last N lines (most recent activity)
    tail = lines[-_MAX_EXTRACT_LINES:] if len(lines) > _MAX_EXTRACT_LINES else lines
    text = "".join(tail).strip()
    if len(text) < 100:
        return {"lessons": [], "failures": []}

    # Filter for domain-relevant sections
    domain_lines = []
    for line in tail:
        if domain.lower() in line.lower():
            domain_lines.append(line)
    # If domain-specific lines found, use those; otherwise use full tail
    if len(domain_lines) >= 5:
        text = "".join(domain_lines[-100:])

    try:
        from pulse.core.llm_router import dispatch
        prompt = _SCRIBE_PROMPT.format(domain=domain) + "\n\nTRANSCRIPT:\n" + text[:8000]
        response = dispatch(prompt, task_type="fast")
        # Parse JSON from response
        import re
        m = re.search(r'\{[^{}]*"lessons"[^{}]*\}', response, re.DOTALL)
        if m:
            return json.loads(m.group())
    except Exception as e:
        _log.warning("LLM extraction failed for %s: %s", domain, e)
    return {"lessons": [], "failures": []}


def _boost_and_inject(items: dict, agent_id: str, domain: str) -> int:
    """Write extracted items to brain with boosted salience + hot cache."""
    import hashlib
    count = 0
    now = datetime.now(timezone.utc).isoformat()

    for category, target in [("lessons", LESSONS_JSONL), ("failures", FAILS_JSONL)]:
        for text in items.get(category, []):
            if not isinstance(text, str) or len(text) < 20:
                continue
            entry = {
                "id": f"NP-{hashlib.md5(text.encode()[:120]).hexdigest()[:8]}",
                "fingerprint": hashlib.md5(text.encode()[:120]).hexdigest()[:16],
                "type": "lesson" if category == "lessons" else "failure",
                "timestamp": now,
                "agent": agent_id,
                "content": text[:300],
                "title": text[:120],
                "domain": domain,
                "salience": _BOOST_SALIENCE,
                "confidence": _BOOST_CONFIDENCE,
                "consolidation_count": 1,
                "evidence_count": 1,
                "tags": ["neuroplasticity", "remedial", domain],
                "_provenance": {
                    "source_agent": agent_id,
                    "source_daemon": "d21_neuroplasticity",
                    "domain": domain,
                },
            }
            try:
                from pulse.brain.save_writers import append_knowledge_jsonl
                if append_knowledge_jsonl(target, entry):
                    count += 1
                    # Also push to hot cache for instant retrieval
                    try:
                        from pulse.brain.memory.hot_cache import write_hot
                        write_hot(entry)
                    except Exception:
                        pass
            except Exception:
                continue
    return count


def _load_metrics() -> dict:
    if not METRICS_FILE.exists():
        return {"runs": 0, "agents_treated": 0, "items_injected": 0,
                "domains_treated": {}, "last_run": ""}
    try:
        return json.loads(METRICS_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return {"runs": 0, "agents_treated": 0, "items_injected": 0,
                "domains_treated": {}, "last_run": ""}


def _save_metrics(metrics: 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(metrics, 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 neuroplasticity cycle: detect → extract → boost → inject."""
    metrics = _load_metrics()
    weak_agents = _detect_weak_agents()

    if not weak_agents:
        if verbose:
            _log.info("No agents with weak domains detected")
        metrics["last_run"] = datetime.now(timezone.utc).isoformat()
        metrics["runs"] = metrics.get("runs", 0) + 1
        _save_metrics(metrics)
        return {"agents": 0, "items": 0, "status": "no_weakness"}

    if verbose:
        _log.info("Found %d agent(s) with weak domains", len(weak_agents))

    total_items = 0
    agents_treated = 0

    for agent in weak_agents:
        aid = agent["agent_id"]
        log_path = _find_capture_log(aid)
        if not log_path:
            if verbose:
                _log.info("  [SKIP] %s — no capture log found", aid)
            continue

        for wd in agent["weak_domains"]:
            domain = wd["domain"]
            if verbose:
                _log.info("  [EXTRACT] %s → domain=%s (rate=%.0f%%)",
                          aid, domain, wd["success_rate"] * 100)

            # Extract domain-specific knowledge
            items = _extract_domain_knowledge(log_path, domain)
            lesson_count = len(items.get("lessons", []))
            fail_count = len(items.get("failures", []))
            if lesson_count + fail_count == 0:
                continue

            # Boost and inject
            injected = _boost_and_inject(items, aid, domain)
            total_items += injected
            if verbose:
                _log.info("    Injected %d items (%dL/%dF)", injected, lesson_count, fail_count)

            # Track per-domain treatment
            dom_stats = metrics.get("domains_treated", {})
            dom_stats[domain] = dom_stats.get(domain, 0) + injected
            metrics["domains_treated"] = dom_stats

        agents_treated += 1

    metrics["runs"] = metrics.get("runs", 0) + 1
    metrics["agents_treated"] = metrics.get("agents_treated", 0) + agents_treated
    metrics["items_injected"] = metrics.get("items_injected", 0) + total_items
    metrics["last_run"] = datetime.now(timezone.utc).isoformat()
    _save_metrics(metrics)

    return {"agents": agents_treated, "items": total_items,
            "weak_count": len(weak_agents), "status": "ok"}


def cmd_neuroplasticity(args: list):
    """CLI entry point for `pulse neuroplasticity`."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    if args and args[0] == "status":
        m = _load_metrics()
        print(f"[D21] Neuroplasticity — Runs: {m.get('runs', 0)}, "
              f"Agents treated: {m.get('agents_treated', 0)}, "
              f"Items injected: {m.get('items_injected', 0)}")
        doms = m.get("domains_treated", {})
        if doms:
            print(f"  Domains: {json.dumps(doms)}")
        print(f"  Last run: {m.get('last_run', 'never')}")
        return
    result = run_once(verbose=True)
    print(json.dumps(result, indent=2))


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