#!/usr/bin/env python3
"""Plan and apply guarded self-healing actions from Jeannie status findings."""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]


@dataclass(frozen=True)
class HealRule:
    check: str
    command: str
    risk: str
    reason: str
    auto: bool
    priority: int


RULES = [
    HealRule("Kubernetes API", "./jeannie start-cluster", "low", "API is down; start-cluster starts kubelet/containerd and worker VMs without destroying state.", True, 100),
    HealRule("Pimox workers running", "./jeannie start-cluster", "low", "Worker VMs are expected cluster capacity and start-cluster is idempotent.", True, 95),
    HealRule("Kubernetes nodes Ready", "./jeannie start-cluster", "low", "NotReady worker nodes commonly recover by starting the saved cluster runtime.", True, 90),
    HealRule("Pi-hole DNS", "./jeannie rpi-services", "medium", "Reapplies RPi DNS services; safe for the lab but can briefly disrupt DNS.", True, 80),
    HealRule("RPi Docker root state", "./jeannie rpi-services", "medium", "May repair service runtime but should be reviewed if storage is failing.", True, 78),
    HealRule("Uptime Kuma HTTP", "./jeannie rpi-services", "medium", "Reapplies RPi service Compose stack.", True, 70),
    HealRule("Gitea container", "./jeannie deploy-gitea", "medium", "Reapplies the Debian-hosted Gitea Compose service.", True, 75),
    HealRule("Gitea local HTTP", "./jeannie deploy-gitea", "medium", "Local Gitea HTTP is down; redeploying Compose is usually safe but still mutates Git service runtime.", True, 74),
    HealRule("Traefik deployment", "./jeannie promote validate", "diagnostic", "Validate gates first; deployment rollout failures are usually downstream of cluster/node health.", False, 65),
    HealRule("Website deployment", "./jeannie promote validate", "diagnostic", "Validate gates first; deployment rollout failures are usually downstream of cluster/node health.", False, 64),
    HealRule("Traefik LoadBalancer HTTP", "./jeannie doctor-edge", "diagnostic", "Edge/LB failures need diagnosis after cluster nodes are healthy.", False, 60),
    HealRule("Website public URL", "./jeannie doctor-edge", "diagnostic", "Public URL failures need edge and cluster diagnosis after local services recover.", False, 58),
    HealRule("Gitea public route", "./jeannie doctor-gitea", "diagnostic", "Public Gitea route needs Gitea and edge diagnosis.", False, 55),
    HealRule("No problem pods", "./jeannie explain status", "diagnostic", "Pod states need root-cause context; avoid blind deletes.", False, 50),
    HealRule("Recent deployments healthy", "./jeannie explain status", "diagnostic", "Deployment drift can be a symptom of node health, image pulls, or scheduling.", False, 45),
    HealRule("Pod restart pressure", "./jeannie explain status", "diagnostic", "Restart pressure needs workload-specific diagnosis before mutation.", False, 40),
    HealRule("Traefik 5xx/404 signals", "./jeannie doctor-edge", "diagnostic", "Use edge logs and route checks before changing config.", False, 35),
]


def load_status(path: Path | None) -> dict[str, object]:
    if path:
        return json.loads(path.read_text(encoding="utf-8"))

    process = subprocess.run(
        [str(REPO_ROOT / "jeannie"), "status", "--json"],
        cwd=REPO_ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if not process.stdout.strip():
        print(process.stderr.strip() or "status produced no JSON", file=sys.stderr)
        raise SystemExit(1)
    return json.loads(process.stdout)


def status_rows(status: dict[str, object]) -> list[dict[str, str]]:
    rows = status.get("rows", [])
    if not isinstance(rows, list):
        return []
    return [row for row in rows if isinstance(row, dict) and row.get("status") in {"fail", "warn"}]


def rule_for(row: dict[str, str]) -> HealRule | None:
    check = str(row.get("check", ""))
    for rule in RULES:
        if rule.check == check:
            return rule
    return None


def pimox_worker_index(row: dict[str, str]) -> int | None:
    text = " ".join(str(row.get(key, "")) for key in ("summary", "detail", "check"))
    match = re.search(r"\bpimox-worker-(\d{1,2})\b", text)
    if not match:
        return None
    return int(match.group(1))


def command_for(row: dict[str, str], rule: HealRule) -> tuple[str, str, int]:
    if rule.check == "Kubernetes nodes Ready":
        index = pimox_worker_index(row)
        if index is not None:
            return (
                f"./jeannie workers restart {index}",
                f"Restart Pimox VM for pimox-worker-{index:02d}; a running VM with a NotReady kubelet can be wedged.",
                98,
            )
    return rule.command, rule.reason, rule.priority


def build_plan(status: dict[str, object]) -> list[dict[str, object]]:
    planned: list[dict[str, object]] = []
    seen_commands: set[str] = set()
    for row in status_rows(status):
        rule = rule_for(row)
        if rule is None:
            planned.append(
                {
                    "check": row.get("check", ""),
                    "area": row.get("area", ""),
                    "status": row.get("status", ""),
                    "command": "./jeannie explain status",
                    "risk": "diagnostic",
                    "auto": False,
                    "priority": 1,
                    "reason": "No specific heal rule exists yet.",
                    "summary": row.get("summary", ""),
                }
            )
            continue
        command, reason, priority = command_for(row, rule)
        if command in seen_commands and rule.auto:
            continue
        seen_commands.add(command)
        planned.append(
            {
                "check": row.get("check", ""),
                "area": row.get("area", ""),
                "status": row.get("status", ""),
                "command": command,
                "risk": rule.risk,
                "auto": rule.auto,
                "priority": priority,
                "reason": reason,
                "summary": row.get("summary", ""),
            }
        )
    return planned


def impact_sort_key(item: dict[str, object]) -> tuple[int, int, str]:
    status_score = 10 if item.get("status") == "fail" else 0
    return (int(item.get("priority") or 0) + status_score, 1 if item.get("auto") else 0, str(item.get("check") or ""))


def highest_impact(plan: list[dict[str, object]], *, auto_only: bool) -> dict[str, object] | None:
    candidates = [item for item in plan if item.get("auto")] if auto_only else list(plan)
    if not candidates:
        return None
    return sorted(candidates, key=impact_sort_key, reverse=True)[0]


def print_plan(plan: list[dict[str, object]], ai: bool) -> None:
    target = highest_impact(plan, auto_only=True)
    top_overall = highest_impact(plan, auto_only=False)
    deferred = [item for item in sorted(plan, key=impact_sort_key, reverse=True) if item is not target]
    print("Jeannie Heal Plan")
    print("=================")
    print("mode: one-at-a-time")
    print(f"findings={len(plan)}")
    print()
    if target:
        print("Next heal target:")
        print(f"  command: {target['command']}")
        print(f"  finding: {target['area']} / {target['check']} ({target['status']})")
        print(f"  impact:  {target['priority']}")
        print(f"  why:     {target['reason']}")
    else:
        print("Next heal target: none")
        if top_overall:
            print(f"highest finding is diagnostic-only: {top_overall['area']} / {top_overall['check']}")
            print(f"next diagnostic command: {top_overall['command']}")
    print()
    if deferred:
        print("Deferred until next status/heal cycle:")
        for item in deferred[:8]:
            auto_marker = "auto" if item.get("auto") else "diagnostic"
            print(f"  - {item['area']} / {item['check']} ({auto_marker}, impact={item['priority']})")
    if ai:
        print_ai_context([target] if target else ([top_overall] if top_overall else []))


def print_ai_context(plan: list[dict[str, object]]) -> None:
    query = " ".join(str(item["check"]) for item in plan[:5])
    if not query:
        return
    process = subprocess.run(
        [str(REPO_ROOT / "scripts" / "query-homelab-ai-index"), "--citations", "--limit", "3", query],
        cwd=REPO_ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        check=False,
    )
    if process.returncode == 0 and process.stdout.strip():
        print()
        print("AI/RAG Context")
        print("--------------")
        print(process.stdout.strip())


def apply_plan(plan: list[dict[str, object]], yes: bool) -> int:
    target = highest_impact(plan, auto_only=True)
    if not target:
        print("No auto-eligible heal actions.")
        return 0

    print("Jeannie Heal Apply")
    print("==================")
    command = str(target["command"])
    sandbox = subprocess.run(
        [str(REPO_ROOT / "scripts" / "agent-sandbox"), "check", command, "--json"],
        cwd=REPO_ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        check=False,
    )
    try:
        decision = json.loads(sandbox.stdout).get("decision")
    except json.JSONDecodeError:
        decision = "approval"
    if decision != "allow":
        print(f"skip: {command} blocked by agent sandbox decision={decision}")
        return 1
    print(f"target: {target['area']} / {target['check']}")
    if not yes:
        print(f"would run: {command}")
    else:
        print(f"run: {command}")
        process = subprocess.run(command.split(), cwd=REPO_ROOT, text=True, check=False)
        if process.returncode != 0:
            return process.returncode
    if not yes:
        print()
        print("Dry run only. Re-run with: ./jeannie heal apply --yes")
        print("After it runs, use ./jeannie status again before healing the next finding.")
    else:
        print("Run ./jeannie status again before healing the next finding.")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    plan_parser = subparsers.add_parser("plan")
    plan_parser.add_argument("--status-json", type=Path)
    plan_parser.add_argument("--ai", action="store_true")
    apply_parser = subparsers.add_parser("apply")
    apply_parser.add_argument("--status-json", type=Path)
    apply_parser.add_argument("--yes", action="store_true")
    apply_parser.add_argument("--ai", action="store_true")
    args = parser.parse_args()

    status = load_status(getattr(args, "status_json", None))
    plan = build_plan(status)
    if args.command == "plan":
        print_plan(plan, args.ai)
        return 0
    if args.command == "apply":
        print_plan(plan, args.ai)
        print()
        return apply_plan(plan, args.yes)
    return 2


if __name__ == "__main__":
    raise SystemExit(main())
