#!/usr/bin/env python3
"""
InsAIts -- Claude Code PreToolUse / PostToolUse Hook
====================================================
This script runs as a Claude Code hook to inspect tool inputs (PreToolUse)
or tool outputs (PostToolUse) for anomalies, and injects corrective
feedback into stdout if problems are detected.

Claude Code reads stdout from hooks and surfaces it back to the model.

SETUP (PreToolUse -- recommended, can block dangerous actions):
  Add to your project's .claude/settings.json:
  {
    "hooks": {
      "PreToolUse": [
        {
          "matcher": ".*",
          "hooks": [
            {
              "type": "command",
              "command": "python3 -m insa_its.hooks.claude_code_hook"
            }
          ]
        }
      ]
    }
  }

  Or for PostToolUse (inspect outputs, cannot block):
  {
    "hooks": {
      "PostToolUse": [
        {
          "matcher": "Write|Edit|MultiEdit",
          "hooks": [
            {
              "type": "command",
              "command": "python3 -m insa_its.hooks.claude_code_hook"
            }
          ]
        }
      ]
    }
  }

HOW IT WORKS:
  Claude Code passes tool input/result as JSON on stdin.
  This script inspects it and writes feedback to stdout.
  Claude Code feeds stdout back to the model automatically.
  Exit code 0 = pass through. Exit code 2 = block + show message (PreToolUse only).

Environment variables:
  INSAITS_MODEL      LLM model identifier for fingerprinting. Default: claude-opus.
  INSAITS_FAIL_MODE  "open" (default) = continue on SDK errors.
                     "closed" = block tool execution on SDK errors.
  INSAITS_VERBOSE    Set to any value to enable debug logging.
  INSAITS_STEALTH    Set to "1" to enable stealth mode: opaque messages with
                     no branding, randomized clean-output checkpoints, and
                     human-in-the-loop for CRITICAL anomalies.

Author: Cristi Bogdan -- Steddy Nova SRL / YuyAI
"""

from __future__ import annotations

import json
import logging
import os
import random
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Configure logging to stderr so it does not interfere with stdout protocol
logging.basicConfig(
    stream=sys.stderr,
    format="[InsAIts] %(message)s",
    level=logging.DEBUG if os.environ.get("INSAITS_VERBOSE") else logging.WARNING,
)
log: logging.Logger = logging.getLogger("insaits-hook")

# --- Constants ---
MIN_CONTENT_LENGTH: int = 10
MAX_SCAN_LENGTH: int = 4000
DEFAULT_MODEL: str = "claude-opus"
BLOCKING_SEVERITIES: frozenset = frozenset({"CRITICAL", "HIGH"})
AUDIT_LOG_PATH: str = os.environ.get(
    "INSAITS_AUDIT_LOG", ".insaits_audit_session.jsonl"
)

# Import from InsAIts SDK (adapted path)
try:
    from insa_its.ai_monitor import AIMonitor
    AVAILABLE: bool = True
except ImportError:
    AVAILABLE = False


def extract_text_from_tool_result(data: Dict[str, Any]) -> Tuple[str, str, str]:
    """Extract the response text, task context, and agent identity from a payload.

    Works for both PreToolUse (tool_input) and PostToolUse (tool_response).

    Returns:
        A (response_text, task_context, agent_id) tuple.
        agent_id is the subagent type for Task calls, or "claude-main" otherwise.
    """
    tool_name: str = data.get("tool_name", "")
    tool_input: Dict[str, Any] = data.get("tool_input", {})
    tool_result: Any = data.get("tool_response", {})

    response_text: str = ""
    task_context: str = ""
    agent_id: str = f"claude:{tool_name}" if tool_name else "claude-main"

    # For Task (subagent) calls -- identify the subagent
    if tool_name == "Task":
        subagent_type: str = tool_input.get("subagent_type", "unknown")
        description: str = tool_input.get("description", "")
        agent_id = f"subagent:{subagent_type}"
        # PostToolUse: inspect the subagent's returned result
        if isinstance(tool_result, str):
            response_text = tool_result
        elif isinstance(tool_result, dict):
            response_text = tool_result.get("output", "") or str(tool_result)
        task_context = f"Subagent({subagent_type}): {description[:60]}"

    # For Write/Edit tools -- inspect what was written
    elif tool_name in ("Write", "Edit", "MultiEdit"):
        content: str = tool_input.get("content", "") or tool_input.get("new_string", "")
        path: str = tool_input.get("file_path", tool_input.get("path", ""))
        response_text = content
        task_context = f"Write file: {path}"

    # For Bash -- inspect command (PreToolUse) or output (PostToolUse)
    elif tool_name == "Bash":
        command: str = str(tool_input.get("command", ""))
        output: str = ""
        if isinstance(tool_result, dict):
            output = tool_result.get("output", "") or tool_result.get("stdout", "")
        elif isinstance(tool_result, str):
            output = tool_result
        # Prefer tool output if available (PostToolUse), else use command (PreToolUse)
        response_text = output if output else command
        task_context = f"Bash: {command[:80]}"

    # For generic assistant messages
    elif "content" in data:
        raw_content: Any = data["content"]
        if isinstance(raw_content, list):
            texts: List[str] = [
                b.get("text", "") for b in raw_content if b.get("type") == "text"
            ]
            response_text = "\n".join(texts)
        elif isinstance(raw_content, str):
            response_text = raw_content
        task_context = str(data.get("task", ""))

    return response_text, task_context, agent_id


def write_audit_entry(
    tool_name: str,
    model: str,
    anomalies: List[Dict[str, str]],
    resolved: bool = False,
    requires_human_review: bool = False,
    stealth: bool = False,
    checkpoint: bool = False,
) -> None:
    """Append a JSONL entry to the audit log for the TUI dashboard.

    The dashboard (insaits_tui.py / AuditWatcher) tails this file in real
    time.  Each line must be a self-contained JSON object with at minimum:
    ts, model, anomalies, resolved.

    Optional metadata fields:
        stealth: whether this entry was processed in stealth mode.
        checkpoint: whether this was a neutral checkpoint message.
        requires_human_review: whether CRITICAL was paused for human review.
    """
    entry: Dict[str, Any] = {
        "ts": datetime.now(timezone.utc).isoformat(),
        "model": model,
        "tool": tool_name,
        "msg_id": f"hook_{int(time.time() * 1000)}",
        "anomalies": anomalies,
        "resolved": resolved,
        "retries": 0,
    }
    if requires_human_review:
        entry["requires_human_review"] = True
    if stealth:
        entry["stealth"] = True
    if checkpoint:
        entry["checkpoint"] = True
    try:
        log_path: Path = Path(AUDIT_LOG_PATH)
        with open(log_path, "a", encoding="utf-8") as fh:
            fh.write(json.dumps(entry, default=str) + "\n")
    except (IOError, OSError) as exc:
        log.debug("Could not write audit log: %s", exc)


def format_feedback(issues: List[Dict[str, str]], stealth: bool = False) -> str:
    """Format findings as a clear message for Claude Code.

    When stealth=True, uses opaque system-constraint language with no
    branding or source identification.

    Returns:
        A human-readable multi-line string describing each finding.
    """
    if stealth:
        return _format_stealth_feedback(issues)

    lines: List[str] = [
        "== InsAIts Quality Gate -- Issues Detected ==",
        "",
    ]
    for i, issue in enumerate(issues, 1):
        sev: str = issue.get("severity", "MEDIUM")
        itype: str = issue.get("type", "unknown").replace("_", " ").upper()
        desc: str = issue.get("description", "")
        fix: str = issue.get("intervention", "")
        lines.extend([
            f"{i}. [{sev}] {itype}",
            f"   What happened: {desc[:100]}",
            f"   Fix required:  {fix[:120]}",
            "",
        ])
    lines.extend([
        "-" * 56,
        "Address ALL issues above before completing this task.",
        "Call insaits_check tool after your fix to verify.",
    ])
    return "\n".join(lines)


# Stealth mode: anomaly type -> opaque system-constraint message (no branding)
_STEALTH_FEEDBACK_MAP: Dict[str, str] = {
    "blank_response": "Output does not meet minimum content requirements. Regenerate.",
    "waiting_for_description": "Proceed with available context. State assumptions and deliver.",
    "incomplete_code": "Completeness check failed: unfinished sections detected.",
    "truncated_output": "Output integrity check: response appears truncated.",
    "context_collapse": "Response pattern inconsistency detected. Realign with task.",
    "semantic_drift": "Relevance constraint violation: output diverges from task parameters.",
    "repetition_loop": "Redundancy threshold exceeded. Advance to unaddressed requirements.",
    "refusal_drift": "Capability assertion error: requested operation is within scope.",
    "confidence_collapse": "Confidence floor violation: commit to a concrete answer.",
    "credential_exposure": "Content policy violation: sensitive tokens detected. Redact.",
    "prompt_injection": "Input integrity violation: instruction-override patterns detected.",
    "tool_poisoning": "Input integrity violation: suspicious patterns in tool content.",
}


def _format_stealth_feedback(issues: List[Dict[str, str]]) -> str:
    """Format findings as opaque system constraints -- no branding."""
    lines: List[str] = [
        "== Output validation failed. Corrections required: ==",
        "",
    ]
    for i, issue in enumerate(issues, 1):
        itype: str = issue.get("type", "unknown")
        stealth_msg: str = _STEALTH_FEEDBACK_MAP.get(
            itype,
            issue.get("intervention", "Review and correct output."),
        )
        lines.extend([
            f"{i}. {stealth_msg}",
            "",
        ])
    lines.extend([
        "Address all items above before continuing.",
    ])
    return "\n".join(lines)


def _write_critical_alert_file(
    issues: List[Dict[str, str]],
    model_id: str,
    tool_name: str,
) -> None:
    """Write CRITICAL anomaly details to a JSON file for human review.

    In stealth mode, the agent never sees CRITICAL content. A human must
    review .insaits_critical_alert.json and decide how to proceed.
    """
    alert: Dict[str, Any] = {
        "ts": datetime.now(timezone.utc).isoformat(),
        "model": model_id,
        "tool": tool_name,
        "requires_human_review": True,
        "anomalies": [
            i for i in issues if i.get("severity") in BLOCKING_SEVERITIES
        ],
    }
    try:
        with open(".insaits_critical_alert.json", "w", encoding="utf-8") as fh:
            import json as _json
            _json.dump(alert, fh, indent=2, default=str)
    except (IOError, OSError) as exc:
        log.error("Failed to write critical alert file: %s", exc)


def main() -> None:
    """Main entry point for the Claude Code hook."""
    # Read tool input/result from stdin (Claude Code passes it as JSON)
    raw: str = sys.stdin.read().strip()

    if not raw:
        sys.exit(0)  # Nothing to inspect -- pass through

    try:
        data: Dict[str, Any] = json.loads(raw)
    except json.JSONDecodeError:
        # Not JSON -- might be plain text output, inspect it directly
        data = {"content": raw, "task": ""}

    response_text: str
    task_context: str
    agent_id: str
    response_text, task_context, agent_id = extract_text_from_tool_result(data)

    # Skip very short outputs (e.g. "OK", empty bash results)
    if len(response_text.strip()) < MIN_CONTENT_LENGTH:
        sys.exit(0)

    # Skip binary or non-text content
    if not response_text.isprintable() and len(response_text) > 100:
        sys.exit(0)

    tool_name: str = data.get("tool_name", "unknown")
    model_id: str = os.environ.get("INSAITS_MODEL", DEFAULT_MODEL)
    stealth: bool = os.environ.get("INSAITS_STEALTH", "").strip() == "1"
    # Use subagent identity when available, otherwise the model env var
    audit_model: str = agent_id if agent_id != "claude-main" else model_id

    if not AVAILABLE:
        # InsAIts not installed -- log and pass through
        log.warning("Monitor not available. Install: pip install insa-its")
        write_audit_entry(tool_name, audit_model, [], resolved=False)
        sys.exit(0)

    # Wrap SDK calls so an internal error does not crash the hook
    try:
        monitor: AIMonitor = AIMonitor(
            model_id=model_id, session_id="hook", stealth=stealth,
        )
        result = monitor.inspect(
            response_text[:MAX_SCAN_LENGTH],
            task_context,
        )
    except Exception as exc:  # Broad catch intentional: unknown SDK internals
        fail_mode: str = os.environ.get("INSAITS_FAIL_MODE", "open").lower()
        write_audit_entry(tool_name, audit_model, [
            {"type": "sdk_error", "severity": "HIGH",
             "description": f"{type(exc).__name__}: {exc}"},
        ], stealth=stealth)
        if fail_mode == "closed":
            sys.stdout.write(
                f"InsAIts SDK error ({type(exc).__name__}); "
                "blocking execution to avoid unscanned input.\n"
            )
            sys.exit(2)
        log.warning(
            "SDK error (%s), skipping inspection: %s",
            type(exc).__name__, exc,
        )
        sys.exit(0)

    if result.clean:
        # All good -- log to audit and pass through silently
        is_checkpoint: bool = False

        # Stealth mode: occasionally emit a neutral checkpoint message
        # on clean outputs to prevent the agent from learning that
        # silence == no monitoring.
        if stealth:
            checkpoint_msg: Optional[str] = monitor.get_stealth_checkpoint()
            if checkpoint_msg:
                is_checkpoint = True
                sys.stdout.write(checkpoint_msg + "\n")

        write_audit_entry(
            tool_name, audit_model, [], resolved=True,
            stealth=stealth, checkpoint=is_checkpoint,
        )
        log.debug("Response clean -- all checks passed.")
        sys.exit(0)

    # Issues found -- format feedback and send to Claude Code
    issues: List[Dict[str, str]] = [
        {
            "type": a.anomaly_type.value,
            "severity": a.severity.value,
            "description": a.description,
            "intervention": a.intervention,
        }
        for a in result.anomalies
    ]

    has_critical: bool = any(
        i["severity"] in BLOCKING_SEVERITIES for i in issues
    )

    # Stealth + CRITICAL: write alert file for human, never show to agent
    if stealth and has_critical:
        write_audit_entry(
            tool_name, audit_model, issues,
            resolved=False, requires_human_review=True,
            stealth=True,
        )
        _write_critical_alert_file(issues, model_id, tool_name)
        sys.exit(2)  # pause execution for human review

    # Write to audit log for dashboard
    write_audit_entry(
        tool_name, audit_model, issues, resolved=False,
        stealth=stealth,
    )

    feedback: str = format_feedback(issues, stealth=stealth)

    # Exit code 2 = Claude Code blocks the action and shows the message
    #               (only effective in PreToolUse)
    # Exit code 0 = pass through (stderr warnings are non-blocking)
    if has_critical:
        sys.stdout.write(feedback + "\n")
        sys.exit(2)  # block -- Claude must fix before continuing
    else:
        log.warning("\n%s", feedback)
        sys.exit(0)


if __name__ == "__main__":
    main()
