#!/usr/bin/env python3
"""
Genesis Multi-Agent Observability Logger (v2 - Full Lifecycle)
==============================================================
Universal hook that captures tool calls, agent lifecycle, errors,
and session events across all agents.

This is Pillar 2 of the Three Pillars (IndyDevDan):
  1. Multi-Agent Orchestration ✓
  2. Multi-Agent Observability ✓ (Full Lifecycle)
  3. Agent Sandboxes

Supported hook events:
  PostToolUse        - Every tool call (primary event stream)
  PostToolUseFailure - Failed tool calls (error tracking)
  SubagentStart      - Agent spawn events (team lifecycle)
  SubagentStop       - Agent completion events (team lifecycle)
  Stop               - Session end events (session tracking)
  Notification       - Inter-agent notifications

Usage:
  python3 observability_logger.py                    # PostToolUse (default)
  python3 observability_logger.py --event tool_failure
  python3 observability_logger.py --event subagent_start
  python3 observability_logger.py --event subagent_stop
  python3 observability_logger.py --event session_stop
  python3 observability_logger.py --event notification

Logs to: /mnt/e/genesis-system/data/observability/events.jsonl

Source: Alpha Evolve Cycle 3+4, axiom MAO-005:
"Observability is prerequisite - you can't improve what you can't see"
"""

import sys
import json
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path

# Event log location
EVENTS_DIR = Path("/mnt/e/genesis-system/data/observability")
EVENTS_FILE = EVENTS_DIR / "events.jsonl"
SESSION_FILE = EVENTS_DIR / "sessions.jsonl"
METRICS_FILE = EVENTS_DIR / "metrics.json"
CORRELATION_FILE = EVENTS_DIR / "correlations.json"

# Max size for arg values in logs (prevent huge payloads)
MAX_ARG_LENGTH = 500
MAX_CONTENT_LENGTH = 200

# Tools to always log (even if they seem boring)
ALWAYS_LOG = {
    "Task", "TaskCreate", "TaskUpdate", "TaskList", "TaskGet",
    "TeamCreate", "TeamDelete", "SendMessage",
    "Bash", "Read", "Write", "Edit", "Glob", "Grep",
    "WebSearch", "WebFetch", "Skill",
    "EnterPlanMode", "ExitPlanMode", "AskUserQuestion",
}

# Sensitive fields to redact
REDACT_FIELDS = {"api_key", "token", "password", "secret", "credential"}


def sanitize_args(args: dict) -> dict:
    """Sanitize tool arguments for safe logging."""
    sanitized = {}
    for key, value in args.items():
        # Redact sensitive fields
        if any(r in key.lower() for r in REDACT_FIELDS):
            sanitized[key] = "[REDACTED]"
            continue

        # Truncate long string values
        if isinstance(value, str):
            if key in ("content", "new_source", "prompt"):
                sanitized[key] = value[:MAX_CONTENT_LENGTH] + ("..." if len(value) > MAX_CONTENT_LENGTH else "")
            elif len(value) > MAX_ARG_LENGTH:
                sanitized[key] = value[:MAX_ARG_LENGTH] + "..."
            else:
                sanitized[key] = value
        elif isinstance(value, (list, dict)):
            # Summarize complex structures
            json_str = json.dumps(value)
            if len(json_str) > MAX_ARG_LENGTH:
                sanitized[key] = f"[{type(value).__name__} len={len(value)}]"
            else:
                sanitized[key] = value
        else:
            sanitized[key] = value

    return sanitized


def get_correlation_id() -> str:
    """Generate or retrieve correlation ID for cross-agent request tracing.

    Correlation ID hierarchy (P0 gap fix from observability audit):
    1. GENESIS_CORRELATION_ID env var (set by parent agent)
    2. CLAUDE_SESSION_ID (session-level tracing)
    3. Generate new UUID (standalone execution)

    This enables tracing requests across Claude -> sub-agent -> Gemini boundaries.
    """
    # Check for explicit correlation ID from parent agent
    corr_id = os.environ.get("GENESIS_CORRELATION_ID", "")
    if corr_id:
        return corr_id

    # Fall back to session ID for session-level correlation
    session_id = os.environ.get("CLAUDE_SESSION_ID", "")
    if session_id:
        return f"sess-{session_id[:12]}"

    # Generate new ID for standalone execution
    return f"gen-{uuid.uuid4().hex[:12]}"


def get_parent_trace() -> dict:
    """Extract parent-child trace information for distributed tracing.

    Returns trace context that links child events to parent events,
    enabling reconstruction of the full call tree.
    """
    return {
        "correlation_id": get_correlation_id(),
        "parent_agent_id": os.environ.get("CLAUDE_AGENT_ID", ""),
        "parent_session_id": os.environ.get("CLAUDE_SESSION_ID", ""),
        "trace_depth": int(os.environ.get("GENESIS_TRACE_DEPTH", "0")),
    }


def extract_agent_info(hook_input: dict) -> dict:
    """Extract agent identity from hook context."""
    return {
        "session_id": os.environ.get("CLAUDE_SESSION_ID", "unknown"),
        "agent_id": os.environ.get("CLAUDE_AGENT_ID", "primary"),
        "agent_name": os.environ.get("CLAUDE_AGENT_NAME", "primary"),
        "team_name": os.environ.get("CLAUDE_TEAM_NAME", ""),
    }


def classify_tool_call(tool_name: str, args: dict) -> str:
    """Classify tool call into category for metrics."""
    categories = {
        "orchestration": {"Task", "TaskCreate", "TaskUpdate", "TaskList",
                          "TaskGet", "TeamCreate", "TeamDelete", "SendMessage"},
        "file_ops": {"Read", "Write", "Edit", "Glob", "Grep"},
        "execution": {"Bash"},
        "research": {"WebSearch", "WebFetch"},
        "planning": {"EnterPlanMode", "ExitPlanMode", "AskUserQuestion"},
        "skills": {"Skill"},
    }
    for category, tools in categories.items():
        if tool_name in tools:
            return category
    return "other"


def update_metrics(tool_name: str, category: str, success: bool):
    """Update rolling metrics file (best-effort)."""
    try:
        metrics = {}
        if METRICS_FILE.exists():
            with open(METRICS_FILE, "r") as f:
                metrics = json.load(f)

        # Update counters
        metrics.setdefault("total_tool_calls", 0)
        metrics["total_tool_calls"] += 1

        metrics.setdefault("by_tool", {})
        metrics["by_tool"].setdefault(tool_name, 0)
        metrics["by_tool"][tool_name] += 1

        metrics.setdefault("by_category", {})
        metrics["by_category"].setdefault(category, 0)
        metrics["by_category"][category] += 1

        metrics.setdefault("success_count", 0)
        metrics.setdefault("failure_count", 0)
        if success:
            metrics["success_count"] += 1
        else:
            metrics["failure_count"] += 1

        metrics["last_updated"] = datetime.now(timezone.utc).isoformat()

        with open(METRICS_FILE, "w") as f:
            json.dump(metrics, f, indent=2)
    except Exception:
        pass  # Never fail on metrics


def log_event(event: dict):
    """Append event to JSONL log with automatic correlation ID injection.

    Every event gets a correlation_id and trace context for cross-agent
    request tracing (P0 gap fix from observability audit).
    """
    try:
        # Inject correlation/trace context into every event
        trace = get_parent_trace()
        event["correlation_id"] = trace["correlation_id"]
        event["trace"] = trace

        EVENTS_DIR.mkdir(parents=True, exist_ok=True)
        with open(EVENTS_FILE, "a") as f:
            f.write(json.dumps(event) + "\n")
    except Exception:
        pass  # Never fail on logging


def handle_tool_call(hook_input: dict):
    """Handle PostToolUse event - tool call tracking."""
    tool_name = hook_input.get("tool_name", hook_input.get("tool", ""))
    tool_args = hook_input.get("tool_input", hook_input.get("args", {}))
    tool_result = hook_input.get("result", {})

    # Skip if not a tool we care about
    if tool_name not in ALWAYS_LOG:
        return

    now = datetime.now(timezone.utc)
    category = classify_tool_call(tool_name, tool_args)
    agent_info = extract_agent_info(hook_input)

    # Determine success from result
    result_str = str(tool_result)
    success = "error" not in result_str.lower()[:500]

    event = {
        "timestamp": now.isoformat(),
        "event_type": "tool_call",
        "tool": tool_name,
        "category": category,
        "args": sanitize_args(tool_args),
        "success": success,
        "agent": agent_info,
    }

    # Add tool-specific metadata
    if tool_name == "Task":
        event["subagent_type"] = tool_args.get("subagent_type", "")
        event["background"] = tool_args.get("run_in_background", False)
        event["description"] = tool_args.get("description", "")[:100]
    elif tool_name in ("TaskCreate", "TaskUpdate"):
        event["task_subject"] = tool_args.get("subject", "")[:100]
        event["task_status"] = tool_args.get("status", "")
    elif tool_name == "SendMessage":
        event["message_type"] = tool_args.get("type", "")
        event["recipient"] = tool_args.get("recipient", "")
    elif tool_name in ("TeamCreate", "TeamDelete"):
        event["team_name"] = tool_args.get("team_name", "")
    elif tool_name == "Bash":
        cmd = tool_args.get("command", "")
        event["command_preview"] = cmd[:200] + ("..." if len(cmd) > 200 else "")
    elif tool_name in ("Read", "Write", "Edit"):
        event["file_path"] = tool_args.get("file_path", "")
    elif tool_name in ("Glob", "Grep"):
        event["pattern"] = tool_args.get("pattern", "")[:100]

    log_event(event)
    update_metrics(tool_name, category, success)


def handle_tool_failure(hook_input: dict):
    """Handle PostToolUseFailure - failed tool call tracking."""
    tool_name = hook_input.get("tool_name", hook_input.get("tool", "unknown"))
    tool_args = hook_input.get("tool_input", hook_input.get("args", {}))
    error = hook_input.get("error", hook_input.get("result", ""))

    now = datetime.now(timezone.utc)
    category = classify_tool_call(tool_name, tool_args)
    agent_info = extract_agent_info(hook_input)

    event = {
        "timestamp": now.isoformat(),
        "event_type": "tool_failure",
        "tool": tool_name,
        "category": category,
        "args": sanitize_args(tool_args) if isinstance(tool_args, dict) else {},
        "success": False,
        "error_preview": str(error)[:300],
        "agent": agent_info,
    }

    log_event(event)
    update_metrics(tool_name, category, False)


def handle_subagent_start(hook_input: dict):
    """Handle SubagentStart - agent spawn lifecycle tracking.

    Also records parent-child relationship for distributed tracing,
    enabling reconstruction of the full agent call tree.
    """
    now = datetime.now(timezone.utc)
    agent_info = extract_agent_info(hook_input)
    trace = get_parent_trace()

    event = {
        "timestamp": now.isoformat(),
        "event_type": "agent_spawn",
        "category": "lifecycle",
        "spawned_agent_id": hook_input.get("agent_id", ""),
        "spawned_agent_name": hook_input.get("agent_name", hook_input.get("name", "")),
        "spawned_agent_type": hook_input.get("agent_type", hook_input.get("subagent_type", "")),
        "team_name": hook_input.get("team_name", ""),
        "description": str(hook_input.get("description", ""))[:200],
        "agent": agent_info,
        "parent_child": {
            "parent_agent": agent_info.get("agent_name", "primary"),
            "child_agent": hook_input.get("agent_name", hook_input.get("name", "")),
            "child_depth": trace["trace_depth"] + 1,
        },
    }

    log_event(event)

    # Update lifecycle metrics
    try:
        metrics = {}
        if METRICS_FILE.exists():
            with open(METRICS_FILE, "r") as f:
                metrics = json.load(f)
        metrics.setdefault("agent_spawns", 0)
        metrics["agent_spawns"] += 1
        metrics.setdefault("agents_active", 0)
        metrics["agents_active"] += 1
        metrics["last_updated"] = now.isoformat()
        with open(METRICS_FILE, "w") as f:
            json.dump(metrics, f, indent=2)
    except Exception:
        pass


def handle_subagent_stop(hook_input: dict):
    """Handle SubagentStop - agent completion lifecycle tracking."""
    now = datetime.now(timezone.utc)
    agent_info = extract_agent_info(hook_input)

    event = {
        "timestamp": now.isoformat(),
        "event_type": "agent_stop",
        "category": "lifecycle",
        "stopped_agent_id": hook_input.get("agent_id", ""),
        "stopped_agent_name": hook_input.get("agent_name", hook_input.get("name", "")),
        "stopped_agent_type": hook_input.get("agent_type", ""),
        "reason": str(hook_input.get("reason", hook_input.get("stop_reason", "")))[:200],
        "agent": agent_info,
    }

    log_event(event)

    # Update lifecycle metrics
    try:
        metrics = {}
        if METRICS_FILE.exists():
            with open(METRICS_FILE, "r") as f:
                metrics = json.load(f)
        metrics.setdefault("agent_stops", 0)
        metrics["agent_stops"] += 1
        active = metrics.get("agents_active", 1)
        metrics["agents_active"] = max(0, active - 1)
        metrics["last_updated"] = now.isoformat()
        with open(METRICS_FILE, "w") as f:
            json.dump(metrics, f, indent=2)
    except Exception:
        pass


def handle_session_stop(hook_input: dict):
    """Handle Stop - session end tracking."""
    now = datetime.now(timezone.utc)
    agent_info = extract_agent_info(hook_input)

    event = {
        "timestamp": now.isoformat(),
        "event_type": "session_stop",
        "category": "lifecycle",
        "reason": str(hook_input.get("reason", ""))[:200],
        "agent": agent_info,
    }

    log_event(event)

    # Log session end to sessions.jsonl
    try:
        session_event = {
            "timestamp": now.isoformat(),
            "type": "session_end",
            "session_id": agent_info.get("session_id", "unknown"),
            "agent_name": agent_info.get("agent_name", "primary"),
        }
        EVENTS_DIR.mkdir(parents=True, exist_ok=True)
        with open(SESSION_FILE, "a") as f:
            f.write(json.dumps(session_event) + "\n")
    except Exception:
        pass


def handle_notification(hook_input: dict):
    """Handle Notification - inter-agent notification tracking."""
    now = datetime.now(timezone.utc)
    agent_info = extract_agent_info(hook_input)

    event = {
        "timestamp": now.isoformat(),
        "event_type": "notification",
        "category": "orchestration",
        "notification_type": hook_input.get("type", hook_input.get("notification_type", "")),
        "content_preview": str(hook_input.get("content", hook_input.get("message", "")))[:200],
        "sender": hook_input.get("sender", hook_input.get("from", "")),
        "agent": agent_info,
    }

    log_event(event)


# Event type dispatch table
EVENT_HANDLERS = {
    "tool_call": handle_tool_call,       # PostToolUse (default)
    "tool_failure": handle_tool_failure,  # PostToolUseFailure
    "subagent_start": handle_subagent_start,  # SubagentStart
    "subagent_stop": handle_subagent_stop,    # SubagentStop
    "session_stop": handle_session_stop,      # Stop
    "notification": handle_notification,      # Notification
}


def main():
    """Universal hook entry point. Dispatches based on --event argument."""
    try:
        hook_input = json.loads(sys.stdin.read())
    except (json.JSONDecodeError, Exception):
        print(json.dumps({}))
        return

    # Determine event type from CLI args (default: tool_call for PostToolUse)
    event_type = "tool_call"
    for i, arg in enumerate(sys.argv):
        if arg == "--event" and i + 1 < len(sys.argv):
            event_type = sys.argv[i + 1]
            break

    # Dispatch to appropriate handler
    handler = EVENT_HANDLERS.get(event_type, handle_tool_call)
    try:
        handler(hook_input)
    except Exception:
        pass  # OBS-003: Never block execution

    # Return without blocking
    print(json.dumps({}))


if __name__ == "__main__":
    main()