"""Obol Agent SDK integration -- security hooks for the Claude Agent SDK.

Provides PreToolUse and PostToolUse hooks that plug directly into the
Claude Agent SDK's hook system, giving any agent built with the SDK
automatic Obol security monitoring.

PreToolUse: runs detection on proposed actions, can deny in enforce mode.
PostToolUse: runs full detection pipeline, sends telemetry, never blocks.

Zero external dependencies (stdlib only). Detection imports from
obol.hooks.detection (also stdlib-only).

Usage:

    from claude_agent_sdk import query, ClaudeAgentOptions
    from obol.agent_sdk import ObolAgent

    obol = ObolAgent(api_key="ob_live_xxx")

    async for msg in query(
        prompt="Analyze the codebase",
        options=ClaudeAgentOptions(hooks=obol.hooks)
    ):
        print(msg)

Minimal usage (reads config from ~/.obol/config.json):

    from obol.agent_sdk import get_obol_hooks

    hooks = get_obol_hooks()
    # Pass to ClaudeAgentOptions(hooks=hooks)

The hooks dict has the shape:
    {
        "PreToolUse": [pre_tool_use_fn],
        "PostToolUse": [post_tool_use_fn],
    }

The Claude Agent SDK expects each hook to be an async callable with
signature: async def hook(input_data, tool_use_id, context) -> dict
"""
import datetime
import hashlib
import json
import os
import sys
import time
import urllib.request
import urllib.error


# ---- Config loading --------------------------------------------------------

_CONFIG_PATH = os.path.expanduser("~/.obol/config.json")
_LOCKDOWN_PATH = os.path.expanduser("~/.obol/lockdown")


def _load_config():
    """Load config from ~/.obol/config.json with type validation."""
    try:
        with open(_CONFIG_PATH) as f:
            config = json.load(f)
        if not isinstance(config, dict):
            return {}
        if "api_key" in config and not isinstance(config["api_key"], str):
            config["api_key"] = ""
        if "api_url" in config and not isinstance(config["api_url"], str):
            config["api_url"] = ""
        if "enforcement_mode" in config:
            if config["enforcement_mode"] not in ("audit", "alert", "enforce"):
                config["enforcement_mode"] = "audit"
        if "telemetry_mode" in config:
            if config["telemetry_mode"] not in ("full", "shapes", "minimal", "off"):
                config["telemetry_mode"] = "shapes"
        return config
    except Exception:
        return {}


def _check_lockdown():
    """Check if agent is locked down. Returns (is_locked, reason)."""
    try:
        with open(_LOCKDOWN_PATH) as f:
            info = json.load(f)
        reason = info.get("reason", "Unknown") if isinstance(info, dict) else "Unknown"
        return True, reason
    except FileNotFoundError:
        return False, ""
    except Exception:
        return True, "Lockdown file unreadable"


# ---- Target extraction (mirrors claude_code.py) ----------------------------

def _extract_target(tool_name, tool_input):
    """Extract a meaningful target string from tool input."""
    if not tool_input or not isinstance(tool_input, dict):
        return None
    if tool_name in ("Read", "Write", "Edit"):
        return tool_input.get("file_path")
    if tool_name == "Bash":
        cmd = tool_input.get("command", "")
        return cmd[:500] if cmd else None
    if tool_name == "Glob":
        return tool_input.get("pattern")
    if tool_name == "Grep":
        return tool_input.get("pattern")
    if tool_name == "WebFetch":
        return tool_input.get("url")
    if tool_name == "WebSearch":
        return tool_input.get("query")
    if tool_name == "NotebookEdit":
        return tool_input.get("notebook_path")
    # MCP tools
    if tool_name and tool_name.startswith("mcp__"):
        for key in ("url", "query", "path", "selector", "key"):
            if key in tool_input:
                return tool_input[key]
    # Generic fallback
    try:
        s = json.dumps(tool_input)
        return s[:500]
    except Exception:
        return None


def _content_hash(tool_input):
    """Produce a SHA-256 hash of tool input for deduplication."""
    if not tool_input:
        return None
    try:
        raw = json.dumps(tool_input, sort_keys=True).encode()
        return hashlib.sha256(raw).hexdigest()[:16]
    except Exception:
        return None


# ---- Severity classification -----------------------------------------------

# Flags from detection.py that should be treated as critical
_CRITICAL_FLAGS = frozenset({
    "escape_attempt", "privilege_escalation", "container_escape",
    "process_injection", "kernel_module_loading", "null_byte_injection",
    "supply_chain_risk", "agent_self_modification", "reverse_shell",
    "piped_exfiltration", "persistence_install", "crypto_miner",
    "delayed_exfiltration", "variable_exfiltration", "env_harvesting",
    "dns_exfiltration",
})


def _has_critical_flags(flags):
    """Return True if any flag in the list is critical severity."""
    return any(f in _CRITICAL_FLAGS for f in flags)


# ---- API helpers (non-blocking, fire-and-forget) ----------------------------

_BACKOFF_PATH = os.path.expanduser("~/.obol/.api_backoff")


def _should_skip_api():
    """Check if we should skip the API call due to recent failures."""
    try:
        with open(_BACKOFF_PATH) as f:
            parts = f.read().strip().split(":")
        if len(parts) == 2:
            retry_after = float(parts[1])
            if time.time() < retry_after:
                return True
            os.unlink(_BACKOFF_PATH)
    except (FileNotFoundError, ValueError, OSError):
        pass
    return False


def _record_api_failure():
    """Record an API failure and set exponential backoff (30s..480s)."""
    failures = 1
    try:
        with open(_BACKOFF_PATH) as f:
            parts = f.read().strip().split(":")
        if len(parts) == 2:
            failures = int(parts[0]) + 1
    except (FileNotFoundError, ValueError, OSError):
        pass
    delay = min(30 * (2 ** (failures - 1)), 480)
    retry_after = time.time() + delay
    try:
        fd = os.open(_BACKOFF_PATH, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "w") as f:
            f.write(f"{failures}:{retry_after:.0f}")
    except Exception:
        pass


def _clear_api_backoff():
    """Clear the backoff state on successful API call."""
    try:
        os.unlink(_BACKOFF_PATH)
    except FileNotFoundError:
        pass
    except Exception:
        pass


def _send_event(event, api_key, api_url):
    """Send event to Obol API. Non-blocking best-effort.

    Uses urllib (stdlib). Fails silently -- monitoring should never
    break the monitored application.
    """
    if not api_key or not api_url:
        return
    if _should_skip_api():
        return

    try:
        body = json.dumps(event).encode()
        req = urllib.request.Request(
            f"{api_url.rstrip('/')}/events",
            data=body,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "User-Agent": "obol-sdk/0.2.0",
            },
            method="POST",
        )
        urllib.request.urlopen(req, timeout=10)
        _clear_api_backoff()
    except Exception:
        _record_api_failure()


# ---- PreToolUse hook --------------------------------------------------------

async def obol_pre_tool_use(input_data, tool_use_id, context,
                            *, _enforcement_mode="audit", _config=None):
    """PreToolUse hook for the Claude Agent SDK.

    Runs on-device detection on the proposed tool action. In enforce mode,
    returns deny for critical threats. In audit/alert mode, always allows.

    Args:
        input_data: Dict with tool_name, tool_input from the Agent SDK.
        tool_use_id: Unique identifier for this tool invocation.
        context: Agent SDK context object.
        _enforcement_mode: Override for testing (normally read from config).
        _config: Override config dict for testing.

    Returns:
        Dict with hookSpecificOutput containing permission decision.
    """
    config = _config or _load_config()
    mode = _enforcement_mode or config.get("enforcement_mode", "audit")

    tool_name = input_data.get("tool_name", "")
    tool_input = input_data.get("tool_input") or {}

    # Layer 1: Lockdown check
    is_locked, lock_reason = _check_lockdown()
    if is_locked and mode == "enforce":
        return {
            "hookSpecificOutput": {
                "hookEventName": "PreToolUse",
                "permissionDecision": "deny",
                "permissionDecisionReason": (
                    f"[Obol] Agent locked down: {lock_reason}. "
                    "Run 'obol unlock' to restore access."
                ),
            }
        }

    # Layer 2: Pre-flight detection
    target = _extract_target(tool_name, tool_input)
    flags = []

    try:
        from obol.hooks.detection import detect_anomalies
        flags = detect_anomalies(tool_name, target or "", "pending")
    except ImportError:
        pass  # Detection module not available
    except Exception:
        pass  # Never block on detection errors

    # Enforce mode: deny on critical findings
    if flags and mode == "enforce" and _has_critical_flags(flags):
        critical = [f for f in flags if f in _CRITICAL_FLAGS]
        return {
            "hookSpecificOutput": {
                "hookEventName": "PreToolUse",
                "permissionDecision": "deny",
                "permissionDecisionReason": (
                    f"[Obol] Blocked: {', '.join(critical[:3])}. "
                    f"Tool '{tool_name}' denied by security policy. "
                    "Set enforcement_mode to 'audit' to allow."
                ),
            }
        }

    # Audit and alert modes: always allow
    return {
        "hookSpecificOutput": {
            "hookEventName": "PreToolUse",
            "permissionDecision": "allow",
            "permissionDecisionReason": "",
        }
    }


# ---- PostToolUse hook -------------------------------------------------------

async def obol_post_tool_use(input_data, tool_use_id, context,
                             *, _api_key=None, _api_url=None, _config=None):
    """PostToolUse hook for the Claude Agent SDK.

    Runs full detection pipeline (anomalies, content scanning, baseline
    tracking) and sends a telemetry event to the Obol API.

    Never blocks. Returns empty dict.

    Args:
        input_data: Dict with tool_name, tool_input, tool_response, etc.
        tool_use_id: Unique identifier for this tool invocation.
        context: Agent SDK context object.
        _api_key: Override for testing.
        _api_url: Override for testing.
        _config: Override config dict for testing.

    Returns:
        Empty dict (PostToolUse hooks don't produce decisions).
    """
    config = _config or _load_config()

    api_key = _api_key or os.environ.get("OBOL_API_KEY") or config.get("api_key", "")
    api_url = _api_url or os.environ.get("OBOL_API_URL") or config.get(
        "api_url", "https://obol.grayowlresearch.com/v1"
    )
    agent_id = config.get("claude_code_agent_id") or config.get("agent_id")

    if not api_key or not agent_id:
        return {}  # Not configured

    tool_name = input_data.get("tool_name", "")
    tool_input = input_data.get("tool_input") or {}
    tool_response = input_data.get("tool_response") or {}
    session_id = input_data.get("session_id", "unknown")

    action = tool_name or "unknown_tool"
    result = "success"  # PostToolUse means it ran successfully
    target = _extract_target(tool_name, tool_input)

    # ---- On-device detection (Layer 1+2) ----
    local_detection_flags = []
    detection_failed = False

    # Load CDC rules
    extra_rules = {}
    try:
        from obol.hooks.rule_update import get_rules
        extra_rules = get_rules(config) or {}
    except (ImportError, Exception):
        pass

    try:
        from obol.hooks.detection import detect_anomalies
        local_detection_flags = detect_anomalies(
            action, target or "", result,
            extra_rules=extra_rules,
        )
    except ImportError:
        pass
    except Exception:
        detection_failed = True

    # ---- Target categorization ----
    target_category = None
    try:
        from obol.hooks.categories import categorize_target
        target_category = categorize_target(
            tool_name, target,
            api_url=api_url,
            allowed_destinations=config.get("allowed_destinations", []),
        )
    except (ImportError, Exception):
        pass

    # ---- Content scanning (secrets/PII) ----
    content_scan = {}
    try:
        from obol.hooks.content_scanner import scan_response, scan_input
        resp_scan = scan_response(tool_name, tool_response)
        if resp_scan.get("secret_count", 0) > 0 or resp_scan.get("pii_count", 0) > 0:
            content_scan["response"] = {
                "secrets": resp_scan["secrets_found"],
                "pii": resp_scan["pii_found"],
            }
        inp_scan = scan_input(tool_name, tool_input)
        if inp_scan.get("secret_count", 0) > 0 or inp_scan.get("pii_count", 0) > 0:
            content_scan["input"] = {
                "secrets": inp_scan["secrets_found"],
                "pii": inp_scan["pii_found"],
            }
    except (ImportError, Exception):
        pass

    # ---- Behavioral baseline ----
    baseline_deviations = []
    baseline_summary = None
    try:
        from obol.hooks.baseline import LocalBaseline
        bl = LocalBaseline(str(agent_id))
        baseline_deviations = bl.check_and_update(
            action, result,
            detected_flags=local_detection_flags,
        )
        if bl.is_ready:
            baseline_summary = bl.summary
    except (ImportError, Exception):
        pass

    # ---- Combine all detected patterns ----
    content_flags = []
    if content_scan.get("response", {}).get("secrets"):
        content_flags.append("secret_in_response")
    if content_scan.get("response", {}).get("pii"):
        content_flags.append("pii_in_response")
    if content_scan.get("input", {}).get("secrets"):
        content_flags.append("secret_in_input")

    all_detected_patterns = list(dict.fromkeys(
        local_detection_flags + baseline_deviations + content_flags
    ))

    # ---- Build event payload ----
    event = {
        "agent_id": agent_id,
        "session_id": str(session_id)[:200],
        "action": action[:100],
        "target": str(target)[:2000] if target else None,
        "target_hash": _content_hash(tool_input),
        "result": result,
        "detected_patterns": all_detected_patterns,
        "metadata": {
            "hook_event": "PostToolUse",
            "sdk_version": "0.2.0",
            "source": "agent_sdk",
        },
    }

    if target_category:
        event["metadata"]["target_category"] = target_category
    if local_detection_flags:
        event["metadata"]["local_detection"] = local_detection_flags
    if baseline_deviations:
        event["metadata"]["baseline_deviations"] = baseline_deviations
    if baseline_summary:
        event["metadata"]["baseline_summary"] = baseline_summary
    if content_scan:
        event["metadata"]["content_scan"] = content_scan
    if detection_failed:
        event["metadata"]["detection_status"] = "failed"
    if tool_use_id:
        event["metadata"]["tool_use_id"] = tool_use_id

    # ---- Shape for privacy (telemetry v2) ----
    try:
        from obol.hooks.telemetry import shape_event
        shaped = shape_event(event, config)
        if shaped is None:
            return {}  # telemetry_mode == "off"
        event = shaped
    except (ImportError, Exception):
        pass  # Send unshaped if shaping fails

    # ---- Send to API (fire-and-forget, non-blocking) ----
    _send_event(event, api_key, api_url)

    return {}


# ---- Convenience function ---------------------------------------------------

def get_obol_hooks(api_key=None, enforcement_mode="audit"):
    """Return a hooks dict ready for ClaudeAgentOptions.

    The returned dict has the shape expected by the Claude Agent SDK:
        {
            "PreToolUse": [pre_hook_fn],
            "PostToolUse": [post_hook_fn],
        }

    The caller wraps these in HookMatcher objects if needed, or passes
    them directly depending on the Agent SDK version.

    Args:
        api_key: Obol API key. If None, reads from OBOL_API_KEY env
                 or ~/.obol/config.json.
        enforcement_mode: "audit" (log only), "alert" (log + notify),
                         or "enforce" (log + notify + block critical).

    Returns:
        Dict mapping hook event names to lists of async hook functions.

    Example:
        from obol.agent_sdk import get_obol_hooks
        hooks = get_obol_hooks(enforcement_mode="alert")
        options = ClaudeAgentOptions(hooks=hooks)
    """
    config = _load_config()
    if api_key:
        config["api_key"] = api_key
    if enforcement_mode:
        config["enforcement_mode"] = enforcement_mode

    api_key_resolved = api_key or os.environ.get("OBOL_API_KEY") or config.get("api_key", "")
    api_url = os.environ.get("OBOL_API_URL") or config.get(
        "api_url", "https://obol.grayowlresearch.com/v1"
    )

    async def pre(input_data, tool_use_id=None, context=None):
        return await obol_pre_tool_use(
            input_data, tool_use_id, context,
            _enforcement_mode=enforcement_mode,
            _config=config,
        )

    async def post(input_data, tool_use_id=None, context=None):
        return await obol_post_tool_use(
            input_data, tool_use_id, context,
            _api_key=api_key_resolved,
            _api_url=api_url,
            _config=config,
        )

    return {
        "PreToolUse": [pre],
        "PostToolUse": [post],
    }


# ---- ObolAgent class -------------------------------------------------------

class ObolAgent:
    """Higher-level wrapper for Obol + Claude Agent SDK integration.

    Manages lifecycle, holds config, and exposes hooks as a property.

    Usage:
        from claude_agent_sdk import query, ClaudeAgentOptions
        from obol.agent_sdk import ObolAgent

        obol = ObolAgent(api_key="ob_live_xxx")

        async for msg in query(
            prompt="Analyze the codebase",
            options=ClaudeAgentOptions(hooks=obol.hooks)
        ):
            print(msg)

        obol.session_end()

    With enforcement:
        obol = ObolAgent(enforcement_mode="enforce")
        # Critical threats will now cause tool denials

    Config resolution order:
        1. Explicit api_key parameter
        2. OBOL_API_KEY environment variable
        3. ~/.obol/config.json api_key field
    """

    def __init__(self, api_key=None, enforcement_mode="audit"):
        """Initialize ObolAgent.

        Args:
            api_key: Obol API key. If None, reads from env/config.
            enforcement_mode: "audit", "alert", or "enforce".
        """
        self._config = _load_config()
        self._api_key = (
            api_key
            or os.environ.get("OBOL_API_KEY")
            or self._config.get("api_key", "")
        )
        self._api_url = (
            os.environ.get("OBOL_API_URL")
            or self._config.get("api_url", "https://obol.grayowlresearch.com/v1")
        )
        self._enforcement_mode = enforcement_mode
        self._session_id = None
        self._session_start_time = None

        # Merge overrides into config copy
        if api_key:
            self._config["api_key"] = api_key
        self._config["enforcement_mode"] = enforcement_mode

    @property
    def hooks(self):
        """Return hooks dict for ClaudeAgentOptions.

        Returns:
            Dict with "PreToolUse" and "PostToolUse" hook lists.
        """
        return get_obol_hooks(
            api_key=self._api_key,
            enforcement_mode=self._enforcement_mode,
        )

    def session_start(self, session_id=None):
        """Signal start of an agent session.

        Sends a session.start event to the Obol API for tracking.

        Args:
            session_id: Optional session identifier. Generated if not provided.
        """
        self._session_id = session_id or hashlib.sha256(
            f"{time.time()}:{os.getpid()}".encode()
        ).hexdigest()[:16]
        self._session_start_time = time.time()

        agent_id = self._config.get("claude_code_agent_id") or self._config.get("agent_id")
        if not agent_id or not self._api_key:
            return

        event = {
            "agent_id": agent_id,
            "session_id": self._session_id,
            "action": "session.start",
            "result": "success",
            "detected_patterns": [],
            "metadata": {
                "hook_event": "SessionStart",
                "sdk_version": "0.2.0",
                "source": "agent_sdk",
            },
        }

        # Shape for privacy
        try:
            from obol.hooks.telemetry import shape_event
            shaped = shape_event(event, self._config)
            if shaped is None:
                return
            event = shaped
        except (ImportError, Exception):
            pass

        _send_event(event, self._api_key, self._api_url)

    def session_end(self):
        """Signal end of an agent session.

        Sends a session.end event with duration metadata.
        """
        agent_id = self._config.get("claude_code_agent_id") or self._config.get("agent_id")
        if not agent_id or not self._api_key:
            return

        duration_ms = None
        if self._session_start_time:
            duration_ms = int((time.time() - self._session_start_time) * 1000)

        event = {
            "agent_id": agent_id,
            "session_id": self._session_id or "unknown",
            "action": "session.end",
            "result": "success",
            "detected_patterns": [],
            "metadata": {
                "hook_event": "SessionEnd",
                "sdk_version": "0.2.0",
                "source": "agent_sdk",
            },
        }
        if duration_ms is not None:
            event["metadata"]["duration_ms"] = duration_ms

        # Shape for privacy
        try:
            from obol.hooks.telemetry import shape_event
            shaped = shape_event(event, self._config)
            if shaped is None:
                return
            event = shaped
        except (ImportError, Exception):
            pass

        _send_event(event, self._api_key, self._api_url)
        self._session_id = None
        self._session_start_time = None
