#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE Worker Tools (Memory) — Brain V4 swarm memory extraction tools.

Split from worker_tools_ext.py for Law 7 compliance.
Contains: tool_extract_knowledge_batch, tool_clean_intents_batch,
          tool_check_antipattern, capture_to_bridge_log, save_session_json.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path

from pulse.core.brain_utils import now_iso as _now_iso, INTENTS_DIR

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent
STATE_DIR = ROOT_DIR / "state"


# === CAPTURE + SESSION PERSISTENCE ===

def capture_to_bridge_log(worker_id: str, events: list[dict], task_id: str,
                          task_title: str = ""):
    """Write signed conversation to bridge-compatible capture log.

    Writes a synthetic USER line (task description) followed by ASSISTANT lines
    so the bridge can pair them as a conversation and flush to knowledge routing.
    """
    capture_file = STATE_DIR / f"pulse_captured_{worker_id}.log"
    # Lazy import: resolve signer from worker_tools sibling at call-time
    import sys as _sys
    _pkg = __package__ or "pulse.workers"
    _sib_mod = _sys.modules.get(_pkg + ".worker_tools")
    if _sib_mod is None:
        import importlib as _il
        _sib_mod = _il.import_module(_pkg + ".worker_tools")
    _sign_fn = getattr(_sib_mod, "_worker_sign", None) or (lambda line, wid: line)
    try:
        with open(capture_file, "a", encoding="utf-8") as f:
            # Synthetic USER line so bridge can pair user+agent for flush
            ts_start = _now_iso()[:19].replace("T", " ")
            user_line = f"[{ts_start}] USER: [TASK {task_id}] {task_title or 'worker task'}"
            f.write(_sign_fn(user_line, worker_id) + "\n")

            for ev in events:
                ts = ev.get("timestamp", _now_iso())[:19].replace("T", " ")
                etype, content = ev.get("type", ""), ev.get("content", "")[:500]
                if etype == "reasoning":
                    line = f"[{ts}] ASSISTANT: {content}"
                elif etype == "tool_call":
                    line = f"[{ts}] ASSISTANT: [TOOL] {ev.get('tool', '')}({json.dumps(ev.get('args', {}))[:200]})"
                elif etype == "tool_result":
                    line = f"[{ts}] ASSISTANT: [RESULT] {ev.get('tool', '')} -> {content[:300]}"
                elif etype == "error":
                    line = f"[{ts}] ASSISTANT: [ERROR] {content}"
                elif etype == "complete":
                    line = f"[{ts}] ASSISTANT: [DONE] Task {task_id}: {content}"
                else:
                    continue
                f.write(_sign_fn(line, worker_id) + "\n")
    except Exception as e:
        logging.getLogger("pulse.workers").warning(
            "Capture log write failed for %s task %s: %s", worker_id, task_id, e)


def save_session_json(worker_id: str, model: str, tier: str, task_id: str,
                      task: dict, events: list[dict], outcome: str,
                      success: bool, tokens_in: int, tokens_out: int, cost: float,
                      skills_used: list[str] | None = None):
    """Save rich structured session for deep analysis."""
    session_dir = STATE_DIR / "worker_sessions"
    session_dir.mkdir(parents=True, exist_ok=True)
    session = {
        "worker_id": worker_id, "task_id": task_id,
        "title": task.get("title", ""), "domain": task.get("domain", "general"),
        "model": model, "tier": tier,
        "started_at": events[0]["timestamp"] if events else _now_iso(),
        "completed_at": _now_iso(), "success": success,
        "outcome": outcome[:1000],
        "tokens_in": tokens_in, "tokens_out": tokens_out, "cost_usd": cost,
        "iterations": len([e for e in events if e["type"] == "reasoning"]),
        "tools_used": list({e["tool"] for e in events if e.get("tool")}),
        "skills_used": skills_used or [],
        "errors": [e["content"] for e in events if e["type"] == "error"],
        "events": events,
    }
    # Layer 3: Sign worker session for integrity tracking
    from pulse.core.pulse_crypto import sign_item
    session = sign_item(session, worker_id)
    try:
        (session_dir / f"{task_id}_{worker_id}.json").write_text(
            json.dumps(session, indent=2, default=str), encoding="utf-8")
    except Exception:
        pass

    # Procedural cache: record habit on success, penalize on failure
    try:
        from pulse.brain.memory.procedural_cache import (
            record_procedure, record_failure, extract_tool_sequence
        )
        _domain = task.get("domain", "general")
        _title = task.get("title", "")
        _desc = task.get("description", _title)
        if success:
            tool_seq = extract_tool_sequence(events)
            record_procedure(_domain, _title, _desc, tool_seq)
        else:
            record_failure(_domain, _title, _desc)
    except Exception:
        pass


# === SWARM MEMORY TOOLS (Brain V4) ===
# Uses the same extraction pipeline as worker_memory.py fast-path.

def tool_extract_knowledge_batch(task_data: str) -> str:
    """Extract engineering knowledge from a batch of chat interactions.
    Called when processing an [EXTRACT_KNOWLEDGE] task."""
    try:
        from pulse.core.llm_router import dispatch, extract_json
        from pulse.admin.retroactive.retroactive_llm_extractor import SCRIBE_PROMPT, write_to_brain

        data = json.loads(task_data)
        batch_text = data.get("batch_text", "")
        agent = data.get("agent_name", "unknown")
        timestamp = data.get("timestamp", _now_iso())
        batch_num = data.get("batch_num", 0)

        if not batch_text:
            return "Empty batch — skipped"

        raw = dispatch(SCRIBE_PROMPT + batch_text, task_type="fast")
        items = extract_json(raw)
        if not isinstance(items, dict):
            return f"LLM returned non-dict: {type(items).__name__}"

        written = write_to_brain(items, source_agent=agent, timestamp=timestamp)
        total = sum(written.values())
        return (f"Batch {batch_num}: {written['lessons']}L/{written['failures']}F/"
                f"{written['patterns']}P ({total} items)")
    except Exception as e:
        return f"ERROR in extract_knowledge_batch: {e}"


def tool_clean_intents_batch(task_data: str) -> str:
    """Distill raw want/dont_want entries into clean actionable preferences.
    Called when processing a [CLEAN_INTENTS] task."""
    try:
        from pulse.core.llm_router import dispatch, extract_json
        from pulse.brain.save_writers import append_knowledge_jsonl

        data = json.loads(task_data)
        batch = data.get("batch", [])
        intent_type = data.get("intent_type", data.get("key", "want"))
        batch_num = data.get("batch_num", 0)

        if not batch:
            return "Empty batch — skipped"

        _purifier_base = (
            "You are the PULSE Intent Purifier. Distill raw {kind} into ONLY actionable, unique {goal}.\n"  
            "REMOVE: duplicates (keep best-worded), fragments <15 chars, code/path fragments, "
            "task-specific requests, contradictions.\n"
            "KEEP: architectural preferences, system behavior rules, tool/workflow prefs, quality standards.\n"
            "OUTPUT: Strict JSON array of cleaned strings. No duplicates. Each a complete, actionable sentence.\n"
            "If NONE are valid, return: []\n\nRaw {kind}:\n"
        )
        _wants_tmpl = _purifier_base.format(kind="wants (user preferences)", goal="preferences")
        _donts_tmpl = _purifier_base.format(kind="dont_wants (prohibitions)", goal="prohibitions")
        raw = dispatch(
            (_wants_tmpl if intent_type == "want" else _donts_tmpl) + json.dumps(batch, indent=None),       
            task_type="fast",
        )
        parsed = extract_json(raw)

        if isinstance(parsed, list):
            clean_items = parsed
        elif isinstance(parsed, dict):
            clean_items = parsed.get("wants", parsed.get("dont_wants", []))
            if not isinstance(clean_items, list):     
                clean_items = []
        else:
            clean_items = []

        clean_items = [s.strip() for s in clean_items if isinstance(s, str) and len(s.strip()) >= 15]       

        from pulse.brain.extraction.wants_classify import _classify_domain

        written = 0
        ts = _now_iso()

        for item_text in clean_items:
            # Use the correct SQLite-first writer
            append_knowledge_jsonl(
                item_type="wants" if intent_type == "want" else "dont_wants",
                description=item_text,
                domain=_classify_domain(item_text),
                timestamp=ts,
                source="swarm_cleaned",
                dedup_key="description",
            )
            written += 1

        return f"Batch {batch_num}: {len(batch)} raw -> {len(clean_items)} clean -> {written} written"      
    except Exception as e:
        return f"ERROR in clean_intents_batch: {e}"   

def tool_check_antipattern(text: str) -> str:
    """Check text/code against known anti-patterns via LLM (V5 Immune System)."""
    try:
        from pulse.admin.patterns.anti_pattern_similarity import check_against_anti_patterns
        violations = check_against_anti_patterns(text)
        if not violations:
            return "CLEAN: No anti-pattern violations detected."
        summary = json.dumps(violations, indent=2)
        return f"VIOLATIONS FOUND ({len(violations)}):\n{summary}"
    except Exception as e:
        return f"ERROR in check_antipattern: {e}"
