timmy-home/scripts/plan_laptop_fleet.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import yaml

DAYLIGHT_START = "10:00"
DAYLIGHT_END = "16:00"


def load_manifest(path: str | Path) -> dict[str, Any]:
    data = yaml.safe_load(Path(path).read_text()) or {}
    data.setdefault("machines", [])
    return data


def validate_manifest(data: dict[str, Any]) -> None:
    machines = data.get("machines", [])
    if not machines:
        raise ValueError("manifest must contain at least one machine")

    seen: set[str] = set()
    for machine in machines:
        hostname = machine.get("hostname", "").strip()
        if not hostname:
            raise ValueError("each machine must declare a hostname")
        if hostname in seen:
            raise ValueError(f"duplicate hostname: {hostname} (unique hostnames are required)")
        seen.add(hostname)

        for field in ("machine_type", "ram_gb", "cpu_cores", "os", "adapter_condition"):
            if field not in machine:
                raise ValueError(f"machine {hostname} missing required field: {field}")


def _laptops(machines: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [m for m in machines if m.get("machine_type") == "laptop"]


def _desktop(machines: list[dict[str, Any]]) -> dict[str, Any] | None:
    for machine in machines:
        if machine.get("machine_type") == "desktop":
            return machine
    return None


def choose_anchor_agents(machines: list[dict[str, Any]], count: int = 2) -> list[dict[str, Any]]:
    eligible = [
        m for m in _laptops(machines)
        if m.get("adapter_condition") in {"good", "ok"} and m.get("always_on_capable", True)
    ]
    eligible.sort(key=lambda m: (m.get("idle_watts", 9999), -m.get("ram_gb", 0), -m.get("cpu_cores", 0), m["hostname"]))
    return eligible[:count]


def assign_roles(machines: list[dict[str, Any]]) -> dict[str, Any]:
    anchors = choose_anchor_agents(machines, count=2)
    anchor_names = {m["hostname"] for m in anchors}
    desktop = _desktop(machines)

    mapping: dict[str, dict[str, Any]] = {}
    for machine in machines:
        hostname = machine["hostname"]
        if desktop and hostname == desktop["hostname"]:
            mapping[hostname] = {
                "role": "desktop_nas",
                "schedule": f"{DAYLIGHT_START}-{DAYLIGHT_END}",
                "duty_cycle": "daylight_only",
            }
        elif hostname in anchor_names:
            mapping[hostname] = {
                "role": "anchor_agent",
                "schedule": "24/7",
                "duty_cycle": "continuous",
            }
        else:
            mapping[hostname] = {
                "role": "daylight_agent",
                "schedule": f"{DAYLIGHT_START}-{DAYLIGHT_END}",
                "duty_cycle": "peak_solar",
            }
    return {
        "anchor_agents": [m["hostname"] for m in anchors],
        "desktop_nas": desktop["hostname"] if desktop else None,
        "role_mapping": mapping,
    }


def build_plan(data: dict[str, Any]) -> dict[str, Any]:
    validate_manifest(data)
    machines = data["machines"]
    role_plan = assign_roles(machines)
    return {
        "fleet_name": data.get("fleet_name", "timmy-laptop-fleet"),
        "machine_count": len(machines),
        "anchor_agents": role_plan["anchor_agents"],
        "desktop_nas": role_plan["desktop_nas"],
        "daylight_window": f"{DAYLIGHT_START}-{DAYLIGHT_END}",
        "role_mapping": role_plan["role_mapping"],
    }


def render_markdown(plan: dict[str, Any], data: dict[str, Any]) -> str:
    lines = [
        "# Laptop Fleet Deployment Plan",
        "",
        f"Fleet: {plan['fleet_name']}",
        f"Machine count: {plan['machine_count']}",
        f"24/7 anchor agents: {', '.join(plan['anchor_agents']) if plan['anchor_agents'] else 'TBD'}",
        f"Desktop/NAS: {plan['desktop_nas'] or 'TBD'}",
        f"Daylight schedule: {plan['daylight_window']}",
        "",
        "## Role mapping",
        "",
        "| Hostname | Role | Schedule | Duty cycle |",
        "|---|---|---|---|",
    ]
    for hostname, role in sorted(plan["role_mapping"].items()):
        lines.append(f"| {hostname} | {role['role']} | {role['schedule']} | {role['duty_cycle']} |")

    lines.extend([
        "",
        "## Machine inventory",
        "",
        "| Hostname | Type | RAM | CPU cores | OS | Adapter | Idle watts | Notes |",
        "|---|---|---:|---:|---|---|---:|---|",
    ])
    for machine in data["machines"]:
        lines.append(
            f"| {machine['hostname']} | {machine['machine_type']} | {machine['ram_gb']} | {machine['cpu_cores']} | {machine['os']} | {machine['adapter_condition']} | {machine.get('idle_watts', 'n/a')} | {machine.get('notes', '')} |"
        )
    return "\n".join(lines) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser(description="Plan LAB-005 laptop fleet deployment.")
    parser.add_argument("manifest", help="Path to laptop fleet manifest YAML")
    parser.add_argument("--markdown", action="store_true", help="Render a markdown deployment plan instead of JSON")
    args = parser.parse_args()

    data = load_manifest(args.manifest)
    plan = build_plan(data)
    if args.markdown:
        print(render_markdown(plan, data))
    else:
        print(json.dumps(plan, indent=2))
    return 0


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