# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""PULSE Save Script: Write-back agent learnings to the brain."""
import argparse
import io
import json
import logging
import os
import sys
from pathlib import Path
# Windows Unicode fix -- only when running as main script (not when imported)
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')
SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent
# When run directly, ensure pulse package is importable
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))
from pulse.core.brain_utils import (
    load_json_dict, save_json, TASK_BOARD_FILE, AGENT_REGISTRY_FILE,
)
# Re-export sub-module functions for backwards compatibility
from pulse.brain.save_writers import (
    write_learnings, write_failures, write_private_learnings, log_session,
    count_brain_queries, record_skill_usage, record_procedure, _now_iso,
    build_save_card,
)
from pulse.brain.save_knowledge_ops import (
    refine_knowledge, flag_incorrect_knowledge, broadcast_critical_failure,
)
_AGENT_MAP = {
    "claude": "claude_daemon", "gemini": "gemini_daemon",
    "codex": "codex_daemon", "grok": "grok_daemon",
    "haiku": "claude_daemon", "sonnet": "claude_daemon", "opus": "claude_daemon",
}

def mark_task_done(task_id: str, agent_name: str, outcome: str = "") -> bool:
    """Mark a task as done on the board."""
    if not task_id:
        return False
    board = load_json_dict(TASK_BOARD_FILE)
    found = False
    for d in board.get("dispatches", []):
        for t in d.get("tasks", []):
            t_id = t.get("id") or t.get("task_id", "")
            if t_id == task_id:
                t["status"] = "done"
                t["outcome"] = outcome or f"Completed by {agent_name}"
                t["completed_at"] = _now_iso()
                found = True
                break
    if found:
        save_json(TASK_BOARD_FILE, board)
    return found

def update_registry(agent_name: str, task_domain: str = "general",
                    outcome: str = "success") -> None:
    """Update agent registry: increment completed count, refresh heartbeat, track outcome."""
    daemon_id = _AGENT_MAP.get(agent_name.lower(), f"{agent_name.lower()}_daemon")
    registry = load_json_dict(AGENT_REGISTRY_FILE)
    if daemon_id in registry:
        entry = registry[daemon_id]
        entry["last_heartbeat"] = _now_iso()
        entry["status"] = "idle"
        entry["current_task"] = None
        history = entry.get("history", {"tasks_completed": 0, "domains": {}})
        history["tasks_completed"] = history.get("tasks_completed", 0) + 1
        succeeded = outcome != "failure"
        history["tasks_succeeded"] = history.get("tasks_succeeded", 0) + (1 if succeeded else 0)
        history["tasks_failed"] = history.get("tasks_failed", 0) + (0 if succeeded else 1)
        domains = history.get("domains", {})
        ds = domains.get(task_domain, {"success": 0, "failure": 0})
        if isinstance(ds, int):
            ds = {"success": ds, "failure": 0}
        ds["success" if succeeded else "failure"] += 1
        domains[task_domain] = ds
        history["domains"] = domains
        entry["history"] = history
        save_json(AGENT_REGISTRY_FILE, registry)

def save(agent_name: str, task_id: str = "", learnings: list = None,
         failures: list = None, files_touched: list = None,
         outcome: str = "", domain: str = "general",
         brain_queries: int = -1, private: bool = False,
         skills_used: list = None, procedure: str = "") -> dict:
    """Full save sequence. Returns summary."""
    learnings = learnings or []
    failures = failures or []
    files_touched = files_touched or []
    skills_used = skills_used or []
    if brain_queries < 0:
        brain_queries = count_brain_queries(agent_name)
    if not procedure:
        has_bq = brain_queries > 0
        has_fail = bool(failures)
        has_files = len(files_touched) >= 3
        if has_bq and files_touched:
            procedure = "brain_query_before_edit"
        elif has_bq and task_id:
            procedure = "brain_query_before_implement"
        elif has_fail and files_touched:
            procedure = "failure_recovery"
        elif not has_fail and has_files:
            procedure = "test_after_change"
    results = {
        "task_marked_done": False, "learnings_saved": 0, "failures_saved": 0,
        "private_saved": 0, "session_logged": False, "registry_updated": False,
        "brain_queries": brain_queries, "skills_tracked": 0, "procedure_logged": False,
    }
    if task_id:
        results["task_marked_done"] = mark_task_done(task_id, agent_name, outcome)
    update_registry(agent_name, domain, outcome="failure" if failures else "success")
    results["registry_updated"] = True
    if private:
        results["private_saved"] = write_private_learnings(learnings, agent_name)
    else:
        results["learnings_saved"] = write_learnings(learnings, agent_name)
    results["failures_saved"] = write_failures(failures, agent_name)
    try:  # Phase 7: pending queue for instant within-session availability
        from pulse.brain.pending_failures import write_pending_failure
        for f_text in failures:
            if f_text and isinstance(f_text, str): write_pending_failure(f_text.strip(), agent=agent_name, domain=domain)
    except Exception: pass
    for f_text in failures:
        if f_text and broadcast_critical_failure(f_text.strip(), agent_name):
            results["broadcast_sent"] = True
    log_session(agent_name, task_id, learnings, failures, files_touched)
    results["session_logged"] = True

    # Episodic memory: save session-level episode for future recall
    try:
        from pulse.brain.index.episodic_writer import save_episode, build_session_summary
        ep_summary = build_session_summary(
            agent_name, learnings, failures, task_id=task_id, domain=domain)
        key_events = (learnings[:3] if learnings else []) + (failures[:2] if failures else [])
        ep_saved = save_episode(
            agent=agent_name, session_summary=ep_summary,
            key_events=key_events, domain=domain,
            session_id=task_id or "",
        )
        results["episodic_saved"] = ep_saved
    except Exception:
        logging.debug("Episodic save failed in pulse_save", exc_info=True)

    if not skills_used:
        try:
            from pulse.brain.monitoring.compliance_tracker import infer_skills_from_session
            skills_used = infer_skills_from_session(learnings, failures, files_touched)
        except Exception:
            logging.debug("Suppressed exception in pulse_save", exc_info=True)
    results["skills_tracked"] = record_skill_usage(skills_used, bool(failures))

    actions = results["learnings_saved"] + results["failures_saved"] + (1 if task_id else 0)
    try:
        from pulse.brain.monitoring.compliance_tracker import update_compliance
        results["compliance_score"] = update_compliance(agent_name, brain_queries, actions)
    except Exception:
        logging.debug("Suppressed exception in pulse_save", exc_info=True)
    if procedure:
        proc_result = record_procedure(procedure, bool(failures))
        results["procedure_logged"] = bool(proc_result)

    # Wire 1A: Session checkpoint — auto-save on every save() call
    try:
        from pulse.brain.session.session_checkpoint import save_checkpoint
        task_title = task_id or "interactive"
        save_checkpoint(
            agent=agent_name,
            task=task_title,
            progress="completed",
            findings=learnings[:5] if learnings else [],
            files_touched=files_touched,
        )
    except Exception:
        pass

    try:
        from pulse.brain.session.session_lifecycle import end_session, get_active_session
        _sid = get_active_session(agent_name)
        task_title = task_id or "interactive"
        if _sid:
            end_session(_sid, summary=f"Task: {task_title or 'unknown'}")
    except Exception:
        pass
    if task_id:
        try:
            from pulse.daemons.neural.prediction_daemon import record_prediction_outcome
            prediction_updated = record_prediction_outcome(
                task_id, success=not bool(failures), agent_name=agent_name)
            results["prediction_feedback"] = prediction_updated
        except (ImportError, Exception):
            logging.debug("Suppressed exception in pulse_save", exc_info=True)

        try:
            from pulse.core.reward.reward_signal import record_task_outcome
            # Retrieve item IDs from working memory session (items brain served for this task)
            retrieved_ids = []
            try:
                from pulse.brain.memory.working_memory import get_active_contexts
                for ctx in get_active_contexts(agent_name):
                    item_id = ctx.get("id", "")
                    if item_id:
                        retrieved_ids.append(item_id)
            except Exception:
                pass
            # Fallback: resolve from task text if working memory empty
            if not retrieved_ids:
                try:
                    from pulse.brain.metacognition.learned_prompt import resolve_items_from_text
                    retrieved_ids = resolve_items_from_text(outcome or task_id)
                except Exception:
                    pass
            signal = record_task_outcome(
                task_id=task_id or "interactive", agent=agent_name,
                success=not bool(failures), domain=domain,
                retrieved_item_ids=retrieved_ids,
            )
            results["reward_signal"] = signal.rpe_delta if signal else 0.0
        except (ImportError, Exception):
            logging.debug("Reward signal failed in pulse_save", exc_info=True)
    return results

def main():
    parser = argparse.ArgumentParser(description="PULSE Save")
    parser.add_argument("--agent", required=True)
    parser.add_argument("--task-id", default="")
    parser.add_argument("--learnings", nargs="*", default=[])
    parser.add_argument("--failures", nargs="*", default=[])
    parser.add_argument("--files-touched", default="")
    parser.add_argument("--outcome", default="")
    parser.add_argument("--domain", default="general")
    parser.add_argument("--stdin", action="store_true")
    parser.add_argument("--private", action="store_true")
    parser.add_argument("--brain-queries", type=int, default=-1)
    parser.add_argument("--skills-used", default="")
    parser.add_argument("--procedure", default="")
    parser.add_argument("--flag-wrong", default="")
    parser.add_argument("--correction", default="")
    parser.add_argument("--refine", default="")
    parser.add_argument("--updated-text", default="")
    parser.add_argument("--reason", default="")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    if args.flag_wrong:
        result = flag_incorrect_knowledge(
            args.flag_wrong, args.correction or "No correction provided", args.agent)
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            if result["flagged"]:
                print(f"[SAVE] Flagged as [DISPUTED] in {result['file']}")
            else:
                print(f"[SAVE] No matching entry found for: {args.flag_wrong}")
            if result["correction_logged"]:
                print("[SAVE] Correction logged to corrections.json")
        return
    if args.refine:
        if not args.updated_text:
            print("[SAVE] --refine requires --updated-text", file=sys.stderr)
            return
        result = refine_knowledge(
            args.refine, args.updated_text,
            args.reason or "No reason provided", args.agent)
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            if result.get("refined"):
                print(f"[SAVE] Refined in {result['file']}: \"{result['original'][:60]}\"")
                print(f"[SAVE] Revision #{result['revision_count']} logged")
            elif result.get("error"):
                print(f"[SAVE] Refinement blocked: {result['error']}")
            else:
                print(f"[SAVE] No matching entry found for: {args.refine[:60]}")
        return
    learnings = args.learnings
    failures = args.failures
    files_touched = [f.strip() for f in args.files_touched.split(",") if f.strip()]
    if args.stdin:
        try:
            data = json.loads(sys.stdin.read())
            learnings = data.get("learnings", learnings)
            failures = data.get("failures", failures)
            files_touched = data.get("files_touched", files_touched)
        except Exception as e:
            print(f"[SAVE] Error reading stdin: {e}", file=sys.stderr)
    skills_list = [s.strip() for s in args.skills_used.split(",") if s.strip()]
    results = save(
        agent_name=args.agent, task_id=args.task_id,
        learnings=learnings, failures=failures,
        files_touched=files_touched, outcome=args.outcome,
        domain=args.domain, brain_queries=args.brain_queries,
        private=args.private, skills_used=skills_list,
        procedure=args.procedure,
    )
    if args.json:
        print(json.dumps(results, indent=2))
    else:
        print(build_save_card(args.agent, results))
if __name__ == "__main__":
    main()
