ajarbot/config/pulse_brain_config.py

"""
Custom Pulse & Brain configuration.

Define your own pulse checks (zero cost) and brain tasks (uses tokens).
"""

import subprocess
from typing import Any, Dict, List

import requests

from pulse_brain import BrainTask, CheckType, PulseCheck


# === PULSE CHECKS (Pure Python, Zero Cost) ===


def check_server_uptime() -> Dict[str, Any]:
    """Check if server is responsive (pure Python, no agent)."""
    try:
        response = requests.get(
            "http://localhost:8000/health", timeout=5
        )
        status = "ok" if response.status_code == 200 else "error"
        return {
            "status": status,
            "message": f"Server responded: {response.status_code}",
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"Server unreachable: {e}",
        }


def check_docker_containers() -> Dict[str, Any]:
    """Check Docker container status (pure Python)."""
    try:
        result = subprocess.run(
            ["docker", "ps", "--format", "{{.Status}}"],
            capture_output=True,
            text=True,
            timeout=5,
        )

        if result.returncode != 0:
            return {
                "status": "error",
                "message": "Docker check failed",
            }

        unhealthy = sum(
            1
            for line in result.stdout.split("\n")
            if "unhealthy" in line.lower()
        )

        if unhealthy > 0:
            message = f"{unhealthy} unhealthy container(s)"
        else:
            message = "All containers healthy"

        return {
            "status": "error" if unhealthy > 0 else "ok",
            "unhealthy_count": unhealthy,
            "message": message,
        }
    except Exception as e:
        return {"status": "error", "message": str(e)}


def check_plex_server() -> Dict[str, Any]:
    """Check if Plex is running (pure Python)."""
    try:
        response = requests.get(
            "http://localhost:32400/identity", timeout=5
        )
        is_ok = response.status_code == 200
        return {
            "status": "ok" if is_ok else "warn",
            "message": (
                "Plex server is running"
                if is_ok
                else "Plex unreachable"
            ),
        }
    except Exception as e:
        return {
            "status": "warn",
            "message": f"Plex check failed: {e}",
        }


def check_unifi_controller() -> Dict[str, Any]:
    """Check UniFi controller (pure Python)."""
    try:
        requests.get(
            "https://localhost:8443", verify=False, timeout=5
        )
        return {
            "status": "ok",
            "message": "UniFi controller responding",
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"UniFi unreachable: {e}",
        }


def check_gpu_temperature() -> Dict[str, Any]:
    """Check GPU temperature (pure Python, requires nvidia-smi)."""
    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=temperature.gpu",
                "--format=csv,noheader",
            ],
            capture_output=True,
            text=True,
            timeout=5,
        )

        if result.returncode != 0:
            return {"status": "ok", "message": "GPU check skipped"}

        temp = int(result.stdout.strip())

        if temp > 85:
            status = "error"
        elif temp > 75:
            status = "warn"
        else:
            status = "ok"

        return {
            "status": status,
            "temperature": temp,
            "message": f"GPU temperature: {temp}C",
        }
    except Exception:
        return {"status": "ok", "message": "GPU check skipped"}


def check_star_citizen_patch() -> Dict[str, Any]:
    """Check for Star Citizen patches (pure Python, placeholder)."""
    return {
        "status": "ok",
        "new_patch": False,
        "message": "No new Star Citizen patches",
    }


# === CUSTOM PULSE CHECKS ===

CUSTOM_PULSE_CHECKS: List[PulseCheck] = [
    PulseCheck(
        "server-uptime", check_server_uptime,
        interval_seconds=60,
    ),
    PulseCheck(
        "docker-health", check_docker_containers,
        interval_seconds=120,
    ),
    PulseCheck(
        "plex-status", check_plex_server,
        interval_seconds=300,
    ),
    PulseCheck(
        "unifi-controller", check_unifi_controller,
        interval_seconds=300,
    ),
    PulseCheck(
        "gpu-temp", check_gpu_temperature,
        interval_seconds=60,
    ),
    PulseCheck(
        "star-citizen", check_star_citizen_patch,
        interval_seconds=3600,
    ),
]


# === BRAIN TASKS (Agent/SDK, Uses Tokens) ===

CUSTOM_BRAIN_TASKS: List[BrainTask] = [
    BrainTask(
        name="server-medic",
        check_type=CheckType.CONDITIONAL,
        prompt_template=(
            "Server is down!\n\n"
            "Status: $message\n\n"
            "Please analyze:\n"
            "1. What could cause this?\n"
            "2. What should I check first?\n"
            "3. Should I restart services?\n\n"
            "Be concise and actionable."
        ),
        condition_func=lambda data: data.get("status") == "error",
        send_to_platform="slack",
        send_to_channel="C_ALERTS",
    ),
    BrainTask(
        name="docker-diagnostician",
        check_type=CheckType.CONDITIONAL,
        prompt_template=(
            "Docker containers unhealthy!\n\n"
            "Unhealthy count: $unhealthy_count\n\n"
            "Please diagnose:\n"
            "1. What might cause container health issues?\n"
            "2. Should I restart them?\n"
            "3. What logs should I check?"
        ),
        condition_func=lambda data: (
            data.get("unhealthy_count", 0) > 0
        ),
        send_to_platform="telegram",
        send_to_channel="123456789",
    ),
    BrainTask(
        name="gpu-thermal-advisor",
        check_type=CheckType.CONDITIONAL,
        prompt_template=(
            "GPU temperature is high!\n\n"
            "Current: $temperatureC\n\n"
            "Please advise:\n"
            "1. Is this dangerous?\n"
            "2. What can I do to cool it down?\n"
            "3. Should I stop current workloads?"
        ),
        condition_func=lambda data: (
            data.get("temperature", 0) > 80
        ),
    ),
    BrainTask(
        name="homelab-briefing",
        check_type=CheckType.SCHEDULED,
        schedule_time="08:00",
        prompt_template=(
            "Good morning! Homelab status report:\n\n"
            "Server: $server_message\n"
            "Docker: $docker_message\n"
            "Plex: $plex_message\n"
            "UniFi: $unifi_message\n"
            "Star Citizen: $star_citizen_message\n\n"
            "Overnight summary:\n"
            "1. Any services restart?\n"
            "2. Notable events?\n"
            "3. Action items for today?\n\n"
            "Keep it brief and friendly."
        ),
        send_to_platform="slack",
        send_to_channel="C_HOMELAB",
    ),
    BrainTask(
        name="homelab-evening-report",
        check_type=CheckType.SCHEDULED,
        schedule_time="22:00",
        prompt_template=(
            "Evening homelab report:\n\n"
            "Today's status:\n"
            "- Server uptime: $server_message\n"
            "- Docker health: $docker_message\n"
            "- GPU temp: $gpu_message\n\n"
            "Summary:\n"
            "1. Any issues today?\n"
            "2. Services that needed attention?\n"
            "3. Overnight monitoring notes?"
        ),
        send_to_platform="telegram",
        send_to_channel="123456789",
    ),
    BrainTask(
        name="patch-notifier",
        check_type=CheckType.CONDITIONAL,
        prompt_template=(
            "New Star Citizen patch detected!\n\n"
            "Please:\n"
            "1. Summarize patch notes (if available)\n"
            "2. Note any breaking changes\n"
            "3. Recommend if I should update now or wait"
        ),
        condition_func=lambda data: data.get("new_patch", False),
        send_to_platform="discord",
        send_to_channel="GAMING_CHANNEL",
    ),
]


def apply_custom_config(pulse_brain: Any) -> None:
    """Apply custom configuration to PulseBrain instance."""
    existing_pulse_names = {c.name for c in pulse_brain.pulse_checks}
    for check in CUSTOM_PULSE_CHECKS:
        if check.name not in existing_pulse_names:
            pulse_brain.pulse_checks.append(check)

    existing_brain_names = {t.name for t in pulse_brain.brain_tasks}
    for task in CUSTOM_BRAIN_TASKS:
        if task.name not in existing_brain_names:
            pulse_brain.brain_tasks.append(task)

    print(
        f"Applied custom config: "
        f"{len(CUSTOM_PULSE_CHECKS)} pulse checks, "
        f"{len(CUSTOM_BRAIN_TASKS)} brain tasks"
    )