#!/usr/bin/env python3
"""Simulate placement for AI-ish homelab workloads."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
SCHED_DIR = REPO_ROOT / "infra" / "ai-scheduler"
NODES_FILE = SCHED_DIR / "nodes.tsv"
WORKLOADS_FILE = SCHED_DIR / "workloads.tsv"


def read_tsv(path: Path) -> list[dict[str, str]]:
    with path.open(encoding="utf-8", newline="") as handle:
        return list(csv.DictReader(handle, delimiter="\t"))


def as_bool(value: str) -> bool:
    return value.strip().lower() in {"true", "yes", "1"}


def score_node(workload: dict[str, str], node: dict[str, str]) -> tuple[int, list[str], list[str]]:
    score = 100
    reasons: list[str] = []
    blockers: list[str] = []
    cpu_needed = int(workload["cpu_m"])
    mem_needed = int(workload["memory_mib"])
    cpu_have = int(node["cpu_m"])
    mem_have = int(node["memory_mib"])

    if cpu_have < cpu_needed:
        blockers.append(f"cpu {cpu_have}m < {cpu_needed}m")
    else:
        reasons.append("cpu fits")
        score += min((cpu_have - cpu_needed) // 250, 10)
    if mem_have < mem_needed:
        blockers.append(f"memory {mem_have}Mi < {mem_needed}Mi")
    else:
        reasons.append("memory fits")
        score += min((mem_have - mem_needed) // 512, 10)
    if as_bool(workload["requires_tailnet"]) and not as_bool(node["tailnet"]):
        blockers.append("tailnet required")
    elif as_bool(node["tailnet"]):
        reasons.append("tailnet available")
        score += 8
    if workload["prefer_arch"] != "any" and workload["prefer_arch"] != node["arch"]:
        score -= 15
        reasons.append(f"non-preferred arch {node['arch']}")
    else:
        reasons.append(f"arch {node['arch']} ok")
    if as_bool(node["control_plane"]) and not as_bool(workload["allow_control_plane"]):
        score -= 30
        reasons.append("control-plane penalty")
    if node["power_class"] == "low":
        score += 5
        reasons.append("low-power node")
    return score, reasons, blockers


def recommend(workload_name: str) -> int:
    workloads = {row["workload"]: row for row in read_tsv(WORKLOADS_FILE)}
    nodes = read_tsv(NODES_FILE)
    if workload_name not in workloads:
        print(f"Unknown workload: {workload_name}")
        print("Known workloads:")
        for name in workloads:
            print(f"  {name}")
        return 2
    workload = workloads[workload_name]
    rows = []
    for node in nodes:
        score, reasons, blockers = score_node(workload, node)
        rows.append((score, node["node"], reasons, blockers))
    rows.sort(reverse=True)

    print("AI Workload Scheduler")
    print("=====================")
    print(f"workload: {workload_name}")
    print()
    for score, node_name, reasons, blockers in rows:
        status = "blocked" if blockers else "candidate"
        print(f"{status:9} {node_name:18} score={score}")
        if blockers:
            print(f"          blockers: {', '.join(blockers)}")
        print(f"          reasons: {', '.join(reasons)}")
    best = next((row for row in rows if not row[3]), None)
    if best:
        print()
        print(f"Recommendation: {best[1]}")
        return 0
    print()
    print("Recommendation: add capacity or relax workload requirements")
    return 1


def simulate() -> int:
    failures = 0
    for workload in read_tsv(WORKLOADS_FILE):
        print()
        rc = recommend(workload["workload"])
        failures += 1 if rc else 0
    return 1 if failures else 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    subparsers.add_parser("simulate")
    rec = subparsers.add_parser("recommend")
    rec.add_argument("workload")
    args = parser.parse_args()
    if args.command == "simulate":
        return simulate()
    if args.command == "recommend":
        return recommend(args.workload)
    return 2


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