#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
PULSE Worker Tools — tool functions, definitions, dispatch, and capture.

Imported by pulse_worker.py and worker_llm.py.
"""
from __future__ import annotations

import json
import os
import re
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

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

from pulse.core.brain_utils import now_iso as _now_iso

_CMD_ALLOWLIST = {
    "git", "python", "pytest", "pip", "node", "npm", "npx",
    "mkdir", "type", "cat", "echo", "ls", "dir", "tsc", "eslint", "prettier",
}


def _guard_path(path_str: str) -> Path:
    """Ensure path stays within project root."""
    p = Path(path_str).resolve()
    if not str(p).startswith(str(ROOT_DIR)):
        raise PermissionError(f"Path outside project root: {path_str}")
    return p


# Tool implementations

def tool_read_file(path: str) -> str:
    p = _guard_path(path)
    if not p.exists():
        return f"ERROR: File not found: {path}"
    return p.read_text(encoding="utf-8", errors="replace")[:50000]


def tool_write_file(path: str, content: str) -> str:
    p = _guard_path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(content, encoding="utf-8")
    return f"OK: Wrote {len(content)} chars to {path}"


_PYTHON_BLOCKED_FLAGS = {"-c", "-m", "--command"}


def tool_run_command(cmd: str) -> str:
    parts = cmd.strip().split()
    if not parts:
        return "ERROR: Empty command"
    base = Path(parts[0]).stem.lower()
    if base not in _CMD_ALLOWLIST:
        return f"ERROR: Command '{base}' not in allowlist: {sorted(_CMD_ALLOWLIST)}"
    # Block python -c / -m to prevent arbitrary code execution
    if base == "python" and len(parts) > 1 and parts[1] in _PYTHON_BLOCKED_FLAGS:
        return f"ERROR: '{parts[1]}' flag blocked for security. Use a .py file instead."
    try:
        result = subprocess.run(parts, capture_output=True, text=True, timeout=120, cwd=str(ROOT_DIR),
                               encoding="utf-8", errors="replace")
        return f"EXIT {result.returncode}\n{(result.stdout + result.stderr)[:10000]}"
    except subprocess.TimeoutExpired:
        return "ERROR: Command timed out (120s)"
    except Exception as e:
        return f"ERROR: {e}"


def tool_brain_query(query: str) -> str:
    try:
        result = subprocess.run(
            [sys.executable, str(TOOLS_DIR / "brain_query.py"), "--query", query],
            capture_output=True, text=True, timeout=30, cwd=str(ROOT_DIR),
            encoding="utf-8", errors="replace",
        )
        return result.stdout[:5000] if result.stdout else "No brain results."
    except Exception as e:
        return f"Brain query error: {e}"


def tool_list_files(pattern: str, path: str = ".") -> str:
    try:
        base = _guard_path(path)
    except PermissionError as e:
        return f"ERROR: {e}"
    if not base.exists():
        return f"ERROR: Path not found: {path}"
    matches = []
    try:
        for p in base.rglob(pattern):
            if p.is_file():
                matches.append(str(p.relative_to(ROOT_DIR)))
                if len(matches) >= 100:
                    break
    except Exception as e:
        return f"ERROR: {e}"
    return f"Found {len(matches)} file(s):\n" + "\n".join(matches) if matches else f"No files matching '{pattern}' in {path}"


def tool_search_code(pattern: str, path: str = ".", context_lines: int = 2) -> str:
    try:
        base = _guard_path(path)
    except PermissionError as e:
        return f"ERROR: {e}"
    context_lines = min(max(int(context_lines), 0), 5)
    # Try git grep first
    try:
        cmd = ["git", "grep", "-n", f"-C{context_lines}", "--max-count=20", "-E", pattern]
        if path != ".":
            cmd.append(str(base.relative_to(ROOT_DIR)))
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=30, cwd=str(ROOT_DIR),
                               encoding="utf-8", errors="replace")
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout[:15000]
    except Exception:
        pass
    # Python fallback
    matches = []
    try:
        compiled = re.compile(pattern)
    except re.error as e:
        return f"ERROR: Invalid regex: {e}"
    _SKIP = (".pyc", ".pyo", ".so", ".dll", ".exe", ".bin", ".png", ".jpg", ".gif")
    try:
        for fp in base.rglob("*"):
            if not fp.is_file() or fp.stat().st_size > 500_000 or fp.suffix in _SKIP:
                continue
            try:
                lines = fp.read_text(encoding="utf-8", errors="replace").splitlines()
            except Exception:
                continue
            for i, line in enumerate(lines):
                if compiled.search(line):
                    rel = str(fp.relative_to(ROOT_DIR))
                    s, e = max(0, i - context_lines), min(len(lines), i + context_lines + 1)
                    snippet = "\n".join(f"  {j+1}: {lines[j]}" for j in range(s, e))
                    matches.append(f"{rel}:{i+1}\n{snippet}")
                    if len(matches) >= 20:
                        return f"Found {len(matches)}+ matches:\n\n" + "\n\n".join(matches)
    except Exception as e:
        return f"ERROR scanning files: {e}"
    return f"Found {len(matches)} match(es):\n\n" + "\n\n".join(matches) if matches else f"No matches for '{pattern}'"


def tool_edit_file(path: str, old_text: str, new_text: str) -> str:
    p = _guard_path(path)
    if not p.exists():
        return f"ERROR: File not found: {path}"
    content = p.read_text(encoding="utf-8", errors="replace")
    count = content.count(old_text)
    if count == 0:
        return f"ERROR: old_text not found in {path}"
    if count > 1:
        idx = content.index(old_text)
        content = content[:idx] + new_text + content[idx + len(old_text):]
        p.write_text(content, encoding="utf-8")
        return f"WARNING: {count} occurrences found, replaced first only in {path}"
    p.write_text(content.replace(old_text, new_text, 1), encoding="utf-8")
    return f"OK: Replaced in {path}"


def tool_submit_task(title: str, domain: str = "general",
                     complexity: str = "medium", description: str = "") -> str:
    try:
        from pulse.tasks.pulse_relay import parse_manual_tasks
        from pulse.tasks.pulse_dispatcher import TaskDispatcher
        task_spec = f"{title}:{domain}:{complexity}"
        tasks = parse_manual_tasks([task_spec])
        if description:
            tasks[0]["description"] = description
        dispatcher = TaskDispatcher(verbose=False)
        if not dispatcher._hmac_key:
            return "ERROR: No HMAC key — cannot dispatch tasks"
        dispatcher.dispatch_tasks(tasks, prompt=title)
        return f"OK: Task '{title}' posted to board (domain={domain}, complexity={complexity})"
    except Exception as e:
        return f"ERROR posting task: {e}"


# Tool definitions (spec → Anthropic format)

def _prop(desc: str, typ: str = "string") -> dict:
    return {"type": typ, "description": desc}

_TOOL_SPECS = [
    ("read_file", "Read a file's contents. Path relative to project root.",
     {"path": _prop("File path relative to project root")}, ["path"]),
    ("write_file", "Write content to a file. Path relative to project root.",
     {"path": _prop("File path relative to project root"), "content": _prop("Content to write")}, ["path", "content"]),
    ("run_command", f"Run a shell command. Allowed: {', '.join(sorted(_CMD_ALLOWLIST))}.",
     {"cmd": _prop("Shell command to run")}, ["cmd"]),
    ("brain_query", "Search the PULSE brain for knowledge, lessons, and anti-patterns.",
     {"query": _prop("Search query")}, ["query"]),
    ("list_files", "List files matching a glob pattern (e.g. '*.py'). Up to 100 results.",
     {"pattern": _prop("Glob pattern"), "path": _prop("Directory (default: '.')")}, ["pattern"]),
    ("search_code", "Search code for a regex pattern with context. Uses git grep with Python fallback.",
     {"pattern": _prop("Regex pattern"), "path": _prop("Directory (default: '.')"),
      "context_lines": _prop("Context lines (0-5, default: 2)", "integer")}, ["pattern"]),
    ("edit_file", "Surgical find-and-replace in a file (first occurrence).",
     {"path": _prop("File path relative to project root"),
      "old_text": _prop("Exact text to find"), "new_text": _prop("Replacement text")}, ["path", "old_text", "new_text"]),
    ("submit_task", "Post a new task to the PULSE task board for other workers.",
     {"title": _prop("Task title"), "domain": _prop("Task domain (default: 'general')"),
      "complexity": _prop("low/medium/high (default: 'medium')"), "description": _prop("Detailed description")}, ["title"]),
    ("complete_task", "Signal task completion with outcome summary.",
     {"outcome": _prop("Summary of outcome")}, ["outcome"]),
    ("extract_knowledge_batch", "Extract engineering knowledge (lessons, failures, patterns) from a batch of chat interactions. Used for [EXTRACT_KNOWLEDGE] memory tasks.",
     {"task_data": _prop("JSON string with keys: batch (list of interactions), agent_name, source_log, batch_num")}, ["task_data"]),
    ("clean_intents_batch", "Distill raw want/dont_want entries into clean actionable preferences. Used for [CLEAN_INTENTS] memory tasks.",
     {"task_data": _prop("JSON string with keys: batch (list of items), key ('want' or 'dont_want'), batch_num")}, ["task_data"]),
    ("check_antipattern", "Check code against known anti-patterns using semantic LLM analysis. Used for [CHECK_ANTIPATTERN] immune system tasks.",
     {"text": _prop("Code or text to check against anti-patterns")}, ["text"]),
]

TOOL_DEFINITIONS_ANTHROPIC = [
    {"name": n, "description": d, "input_schema": {"type": "object", "properties": p, "required": r}}
    for n, d, p, r in _TOOL_SPECS
]

TOOL_DISPATCH = {
    "read_file": lambda a: tool_read_file(a["path"]),
    "write_file": lambda a: tool_write_file(a["path"], a["content"]),
    "run_command": lambda a: tool_run_command(a["cmd"]),
    "brain_query": lambda a: tool_brain_query(a["query"]),
    "list_files": lambda a: tool_list_files(a["pattern"], a.get("path", ".")),
    "search_code": lambda a: tool_search_code(a["pattern"], a.get("path", "."), a.get("context_lines", 2)),
    "edit_file": lambda a: tool_edit_file(a["path"], a["old_text"], a["new_text"]),
    "submit_task": lambda a: tool_submit_task(a["title"], a.get("domain", "general"), a.get("complexity", "medium"), a.get("description", "")),
    "complete_task": lambda a: a.get("outcome", "").strip() or "Task completed without detailed output.",
    "extract_knowledge_batch": lambda a: tool_extract_knowledge_batch(a["task_data"]),
    "clean_intents_batch": lambda a: tool_clean_intents_batch(a["task_data"]),
    "check_antipattern": lambda a: tool_check_antipattern(a["text"]),
}


# Conversation capture

_SALIENCE_SCORES = {
    "reasoning": 1.5, "tool_call": 1.0, "tool_result": 0.8,
    "brain_query": 2.0, "complete": 3.0, "task_start": 1.0,
}


def compute_salience(event_type: str, tool_name: str = "", error: bool = False) -> float:
    """Score event salience 0.1-10.0 for knowledge extraction priority."""
    if error:
        return 4.0
    base = _SALIENCE_SCORES.get(event_type, 0.5)
    if tool_name in ("write_file", "edit_file"):
        base *= 2.0
    elif tool_name == "run_command":
        base *= 1.5
    return min(base, 10.0)


def capture_to_bridge_log(worker_id: str, events: list[dict], task_id: str):
    """Write conversation to bridge-compatible capture log."""
    capture_file = STATE_DIR / f"pulse_captured_{worker_id}.log"
    try:
        with open(capture_file, "a", encoding="utf-8") as f:
            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":
                    f.write(f"[{ts}] {worker_id}: {content}\n")
                elif etype == "tool_call":
                    f.write(f"[{ts}] {worker_id}: [TOOL] {ev.get('tool', '')}({json.dumps(ev.get('args', {}))[:200]})\n")
                elif etype == "tool_result":
                    f.write(f"[{ts}] {worker_id}: [RESULT] {ev.get('tool', '')} -> {content[:300]}\n")
                elif etype == "error":
                    f.write(f"[{ts}] {worker_id}: [ERROR] {content}\n")
                elif etype == "complete":
                    f.write(f"[{ts}] {worker_id}: [DONE] Task {task_id}: {content}\n")
    except Exception:
        pass


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


from pulse.browser.core import browser_goto, browser_get_content, browser_fill, browser_click



# === 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.admin.retroactive.retroactive_intents_cleaner import WANTS_PROMPT, DONT_WANTS_PROMPT
        from pulse.brain.save_writers import append_to_json_list

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

        if not batch:
            return "Empty batch — skipped"

        prompt_template = WANTS_PROMPT if intent_type == "want" else DONT_WANTS_PROMPT
        raw = dispatch(prompt_template + 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]

        written = 0
        key = "want" if intent_type == "want" else "dont_want"
        ts = _now_iso()
        target = Path(file_path) if file_path else (STATE_DIR.parent / "Hippocampus" / "Evidence" / "Intents" / f"{key}s.json")

        for item_text in clean_items:
            entry = {key: item_text, "domain": "general", "visibility": "public",
                     "timestamp": ts, "source": "swarm_cleaned"}
            if append_to_json_list(target, entry, dedup_key=key):
                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}"
