src/tools.py

"""
P08 · SRE Agent Tools
All tools the agent can call. Mix of real (HTTP) and mock (realistic data).
No API keys required for any tool.
"""

import json
import os
import time
from dataclasses import dataclass
from typing import Any

import httpx


# ── Tool result dataclass ─────────────────────────────────────────────────────
@dataclass
class ToolResult:
    tool: str
    success: bool
    data: Any
    error: str = ""
    latency_ms: int = 0


# ── Tool 1: SLO Status (mock Prometheus) ─────────────────────────────────────
MOCK_SERVICES = {
    "api-gateway": {
        "slo_target": 99.9,
        "error_rate_1h": 0.08,
        "error_rate_6h": 0.12,
        "total_requests_1h": 125000,
    },
    "payment-service": {
        "slo_target": 99.95,
        "error_rate_1h": 0.02,
        "error_rate_6h": 0.01,
        "total_requests_1h": 45000,
    },
    "auth-service": {
        "slo_target": 99.9,
        "error_rate_1h": 0.45,
        "error_rate_6h": 0.38,
        "total_requests_1h": 280000,
    },
    "notification-service": {
        "slo_target": 99.5,
        "error_rate_1h": 0.03,
        "error_rate_6h": 0.02,
        "total_requests_1h": 15000,
    },
}


def check_slo_status(service_name: str) -> ToolResult:
    """
    Check SLO burn rate and error budget for a service.
    Uses mock Prometheus data — realistic for portfolio demo.
    """
    start = time.time()
    service_name = service_name.lower().strip()

    if service_name not in MOCK_SERVICES:
        available = list(MOCK_SERVICES.keys())
        return ToolResult(
            tool="check_slo_status",
            success=False,
            data=None,
            error=f"Service '{service_name}' not found. Available: {available}",
            latency_ms=int((time.time() - start) * 1000),
        )

    svc = MOCK_SERVICES[service_name]
    slo_target = svc["slo_target"]
    error_budget_pct = 1 - (slo_target / 100)

    # Burn rate = actual error rate / allowed error rate
    burn_rate_1h = svc["error_rate_1h"] / (error_budget_pct * 100)
    burn_rate_6h = svc["error_rate_6h"] / (error_budget_pct * 100)

    # Budget consumed this month (30d window)
    budget_consumed_pct = min(
        (svc["error_rate_6h"] / (error_budget_pct * 100)) * (6 / (30 * 24)) * 100,
        100,
    )
    budget_remaining_pct = max(0, 100 - budget_consumed_pct)

    # Status
    if burn_rate_1h > 14.4:
        status = "CRITICAL"
        action = "Page on-call immediately — exhausts budget in <2h"
    elif burn_rate_1h > 6:
        status = "WARNING"
        action = "Investigate — exhausts budget in <5 days"
    elif burn_rate_1h > 1:
        status = "ELEVATED"
        action = "Monitor closely — burning faster than replenishment"
    else:
        status = "OK"
        action = "Within normal bounds"

    return ToolResult(
        tool="check_slo_status",
        success=True,
        data={
            "service": service_name,
            "slo_target_pct": slo_target,
            "status": status,
            "action_required": action,
            "burn_rate_1h": round(burn_rate_1h, 2),
            "burn_rate_6h": round(burn_rate_6h, 2),
            "error_rate_1h_pct": svc["error_rate_1h"],
            "error_rate_6h_pct": svc["error_rate_6h"],
            "budget_remaining_pct": round(budget_remaining_pct, 1),
            "total_requests_1h": svc["total_requests_1h"],
            "data_source": "mock-prometheus",
        },
        latency_ms=int((time.time() - start) * 1000),
    )


# ── Tool 2: Runbook Fetcher (real — reads local files) ────────────────────────
RUNBOOK_PATH = os.path.join(
    os.path.dirname(__file__), "..", "data", "raw", "runbooks.jsonl"
)


def get_runbook(alert_name: str) -> ToolResult:
    """
    Fetch runbook for an alert name.
    Reads from local runbooks.jsonl — real file lookup.
    """
    start = time.time()
    alert_name = alert_name.lower().strip()

    try:
        runbooks = {}
        runbook_file = os.path.abspath(RUNBOOK_PATH)
        with open(runbook_file) as f:
            for line in f:
                line = line.strip()
                if line:
                    rb = json.loads(line)
                    runbooks[rb["alert"].lower()] = rb

        # Exact match first
        if alert_name in runbooks:
            return ToolResult(
                tool="get_runbook",
                success=True,
                data=runbooks[alert_name],
                latency_ms=int((time.time() - start) * 1000),
            )

        # Fuzzy match — find runbooks containing the search term
        matches = [
            rb for key, rb in runbooks.items()
            if alert_name in key or key in alert_name
        ]
        if matches:
            return ToolResult(
                tool="get_runbook",
                success=True,
                data=matches[0],
                latency_ms=int((time.time() - start) * 1000),
            )

        available = list(runbooks.keys())
        return ToolResult(
            tool="get_runbook",
            success=False,
            data=None,
            error=f"No runbook found for '{alert_name}'. Available: {available}",
            latency_ms=int((time.time() - start) * 1000),
        )

    except FileNotFoundError:
        return ToolResult(
            tool="get_runbook",
            success=False,
            data=None,
            error=f"Runbook file not found at {RUNBOOK_PATH}",
            latency_ms=int((time.time() - start) * 1000),
        )


# ── Tool 3: Recent Alerts (mock PagerDuty) ────────────────────────────────────
MOCK_ALERTS = [
    {
        "id": "PD-001",
        "service": "auth-service",
        "alert": "HighErrorRate",
        "severity": "critical",
        "started_at": "2026-05-30T03:12:00Z",
        "duration_min": 47,
        "status": "firing",
        "summary": "auth-service error rate 0.45% exceeds SLO threshold of 0.1%",
    },
    {
        "id": "PD-002",
        "service": "api-gateway",
        "alert": "SlowResponseTime",
        "severity": "warning",
        "started_at": "2026-05-30T03:45:00Z",
        "duration_min": 14,
        "status": "firing",
        "summary": "api-gateway p99 latency 850ms exceeds 500ms threshold",
    },
    {
        "id": "PD-003",
        "service": "payment-service",
        "alert": "CertificateExpiringSoon",
        "severity": "warning",
        "started_at": "2026-05-29T10:00:00Z",
        "duration_min": 1032,
        "status": "firing",
        "summary": "TLS certificate expires in 7 days",
    },
]


def query_recent_alerts(severity: str = "all") -> ToolResult:
    """
    Query recent firing alerts. Mock PagerDuty data.
    severity: 'all', 'critical', 'warning'
    """
    start = time.time()
    severity = severity.lower().strip()

    if severity == "all":
        alerts = MOCK_ALERTS
    else:
        alerts = [a for a in MOCK_ALERTS if a["severity"] == severity]

    return ToolResult(
        tool="query_recent_alerts",
        success=True,
        data={
            "total_firing": len(alerts),
            "filter": severity,
            "alerts": alerts,
            "data_source": "mock-pagerduty",
        },
        latency_ms=int((time.time() - start) * 1000),
    )


# ── Tool 4: Service Health Check (real HTTP) ─────────────────────────────────
def check_service_health(url: str) -> ToolResult:
    """
    Real HTTP health check on a URL.
    Works on any public URL — no API key needed.
    """
    start = time.time()

    # Safety: only allow http/https
    if not url.startswith(("http://", "https://")):
        url = "https://" + url

    try:
        resp = httpx.get(url, timeout=10, follow_redirects=True)
        latency_ms = int((time.time() - start) * 1000)
        return ToolResult(
            tool="check_service_health",
            success=True,
            data={
                "url": url,
                "status_code": resp.status_code,
                "healthy": resp.status_code < 400,
                "latency_ms": latency_ms,
                "response_size_bytes": len(resp.content),
            },
            latency_ms=latency_ms,
        )
    except httpx.TimeoutException:
        return ToolResult(
            tool="check_service_health",
            success=False,
            data={"url": url},
            error="Request timed out after 10s",
            latency_ms=int((time.time() - start) * 1000),
        )
    except Exception as e:
        return ToolResult(
            tool="check_service_health",
            success=False,
            data={"url": url},
            error=str(e),
            latency_ms=int((time.time() - start) * 1000),
        )


# ── Tool 5: Error Budget Calculator (pure math) ───────────────────────────────
def calculate_error_budget(
    slo_pct: float,
    window_days: int = 30,
    consumed_pct: float = 0.0,
) -> ToolResult:
    """
    Calculate error budget for an SLO.
    Pure math — always works, no dependencies.
    """
    start = time.time()

    if not (0 < slo_pct <= 100):
        return ToolResult(
            tool="calculate_error_budget",
            success=False,
            data=None,
            error="slo_pct must be between 0 and 100",
        )

    error_budget_pct = 1 - (slo_pct / 100)
    total_minutes = window_days * 24 * 60
    allowed_downtime_min = total_minutes * error_budget_pct
    consumed_min = allowed_downtime_min * (consumed_pct / 100)
    remaining_min = allowed_downtime_min - consumed_min

    return ToolResult(
        tool="calculate_error_budget",
        success=True,
        data={
            "slo_target_pct": slo_pct,
            "window_days": window_days,
            "error_budget_pct": round(error_budget_pct * 100, 4),
            "allowed_downtime_minutes": round(allowed_downtime_min, 1),
            "allowed_downtime_hours": round(allowed_downtime_min / 60, 2),
            "consumed_pct": consumed_pct,
            "consumed_minutes": round(consumed_min, 1),
            "remaining_minutes": round(remaining_min, 1),
            "remaining_hours": round(remaining_min / 60, 2),
        },
        latency_ms=int((time.time() - start) * 1000),
    )


# ── Tool registry ─────────────────────────────────────────────────────────────
TOOLS = {
    "check_slo_status": {
        "fn": check_slo_status,
        "description": "Check SLO burn rate and error budget for a service. Args: service_name (str)",
        "requires_approval": False,
    },
    "get_runbook": {
        "fn": get_runbook,
        "description": "Fetch runbook steps for an alert. Args: alert_name (str)",
        "requires_approval": False,
    },
    "query_recent_alerts": {
        "fn": query_recent_alerts,
        "description": "Get recent firing alerts. Args: severity ('all', 'critical', 'warning')",
        "requires_approval": False,
    },
    "check_service_health": {
        "fn": check_service_health,
        "description": "Real HTTP health check on a URL. Args: url (str). REQUIRES APPROVAL — makes external network call.",
        "requires_approval": True,
    },
    "calculate_error_budget": {
        "fn": calculate_error_budget,
        "description": "Calculate error budget. Args: slo_pct (float), window_days (int, default 30), consumed_pct (float, default 0)",
        "requires_approval": False,
    },
}