from __future__ import annotations

import argparse
import json
import shlex
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

MANAGED_LABEL = "hochguertel.startup.managed"
EXCLUDE_LABEL = "hochguertel.startup.exclude"
GROUP_LABEL = "hochguertel.startup.group"
PRIORITY_LABEL = "hochguertel.startup.priority"
WAIT_FOR_HEALTH_LABEL = "hochguertel.startup.wait-for-health"

DEFAULT_GROUP_ORDER = {
    "edge": 500,
    "auth": 400,
    "data": 300,
    "platform": 200,
    "apps": 100,
    "background": 0,
    "default": -100,
}

SCHEMA_TEXT = f"""Suggested compose label schema
================================

Recommended minimal labels:

  {MANAGED_LABEL}: "true"
  {GROUP_LABEL}: "edge|auth|data|platform|apps|background"
  {PRIORITY_LABEL}: "integer; higher starts earlier within the group"
  {WAIT_FOR_HEALTH_LABEL}: "true|false"

Optional safety label:

  {EXCLUDE_LABEL}: "true"

Recommended meaning:

- {MANAGED_LABEL}=true
    Opts the service into the controller. Keep this explicit so ad hoc or per-user
    containers are not pulled into boot management accidentally.
- {GROUP_LABEL}
    Coarse startup wave. Use a small fixed vocabulary so the boot order stays readable.
- {PRIORITY_LABEL}
    Fine-grained ordering inside the group. Higher numbers start first.
- {WAIT_FOR_HEALTH_LABEL}=true
    The controller waits until the batch reaches a healthy state before continuing.
- {EXCLUDE_LABEL}=true
    Explicitly keep a container out of startup management even if other defaults change later.

Example compose snippet:

services:
  traefik:
    restart: always
    labels:
      {MANAGED_LABEL}: "true"
      {GROUP_LABEL}: "edge"
      {PRIORITY_LABEL}: "1000"
      {WAIT_FOR_HEALTH_LABEL}: "true"

  authelia:
    restart: always
    labels:
      {MANAGED_LABEL}: "true"
      {GROUP_LABEL}: "auth"
      {PRIORITY_LABEL}: "900"
      {WAIT_FOR_HEALTH_LABEL}: "true"

  homepage:
    restart: always
    labels:
      {MANAGED_LABEL}: "true"
      {GROUP_LABEL}: "apps"
      {PRIORITY_LABEL}: "200"
      {WAIT_FOR_HEALTH_LABEL}: "false"
"""


@dataclass(slots=True)
class TargetContainer:
    container_id: str
    name: str
    project: str | None
    service: str | None
    restart_policy: str | None
    status: str
    running: bool
    health: str | None
    healthcheck_defined: bool
    group: str
    group_rank: int
    priority: int
    wait_for_health: bool

    @property
    def batch_key(self) -> tuple[int, int]:
        return (self.group_rank, self.priority)


@dataclass(slots=True)
class PressureSnapshot:
    cpu: float | None
    memory: float | None
    io: float | None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Reference Python implementation of the label-driven Podman startup controller."
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    schema = subparsers.add_parser("schema", help="Print the suggested compose label schema.")
    schema.set_defaults(func=command_schema)

    for name, help_text in (
        ("plan", "Show the current startup plan without starting containers."),
        ("run", "Run the startup controller and start matching containers."),
    ):
        subparser = subparsers.add_parser(name, help=help_text)
        subparser.add_argument(
            "--groups",
            default="",
            help="Comma-separated subset of groups to process (default: all groups).",
        )
        subparser.add_argument(
            "--allowed-restart-policies",
            default="always",
            help="Comma-separated restart policies allowed into the plan (default: always).",
        )
        subparser.add_argument(
            "--batch-size",
            type=int,
            default=2,
            help="Maximum number of containers started in a single batch (default: %(default)s).",
        )
        subparser.add_argument(
            "--settle-seconds",
            type=float,
            default=3.0,
            help="Sleep between completed batches (default: %(default)s).",
        )
        subparser.add_argument(
            "--health-timeout",
            type=float,
            default=300.0,
            help="Maximum seconds to wait for a health-gated batch (default: %(default)s).",
        )
        subparser.add_argument(
            "--poll-interval",
            type=float,
            default=2.0,
            help="Polling interval for health and pressure checks (default: %(default)s).",
        )
        subparser.add_argument(
            "--cpu-pressure-avg10-max",
            type=float,
            default=None,
            help="Maximum allowed /proc/pressure/cpu avg10 before a batch starts.",
        )
        subparser.add_argument(
            "--memory-pressure-avg10-max",
            type=float,
            default=None,
            help="Maximum allowed /proc/pressure/memory avg10 before a batch starts.",
        )
        subparser.add_argument(
            "--io-pressure-avg10-max",
            type=float,
            default=None,
            help="Maximum allowed /proc/pressure/io avg10 before a batch starts.",
        )
        subparser.add_argument(
            "--pressure-timeout",
            type=float,
            default=300.0,
            help=(
                "Maximum seconds to wait for pressure to fall below thresholds "
                "(default: %(default)s)."
            ),
        )
        subparser.set_defaults(func=command_plan if name == "plan" else command_run)

    return parser.parse_args()


def run_command(command: list[str]) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(command, capture_output=True, text=True, check=False)
    if result.returncode != 0:
        joined = shlex.join(command)
        raise RuntimeError(
            f"Command failed ({result.returncode}): {joined}\n{result.stderr.strip()}"
        )
    return result


def log(message: str) -> None:
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {message}")


def parse_bool(value: Any, *, default: bool = False) -> bool:
    if value is None:
        return default
    if isinstance(value, bool):
        return value
    normalized = str(value).strip().lower()
    if normalized in {"1", "true", "yes", "on"}:
        return True
    if normalized in {"0", "false", "no", "off"}:
        return False
    return default


def parse_int(value: Any, *, default: int = 0) -> int:
    if value is None:
        return default
    try:
        return int(str(value).strip())
    except ValueError:
        return default


def chunked(items: list[TargetContainer], size: int) -> list[list[TargetContainer]]:
    return [items[index : index + size] for index in range(0, len(items), size)]


def read_pressure_avg10(resource: str) -> float | None:
    pressure_path = Path("/proc/pressure") / resource
    if not pressure_path.exists():
        return None

    for line in pressure_path.read_text().splitlines():
        if not line.startswith("some "):
            continue
        fields = dict(field.split("=", 1) for field in line.split()[1:] if "=" in field)
        raw = fields.get("avg10")
        if raw is None:
            return None
        return float(raw)
    return None


def read_pressure_snapshot() -> PressureSnapshot:
    return PressureSnapshot(
        cpu=read_pressure_avg10("cpu"),
        memory=read_pressure_avg10("memory"),
        io=read_pressure_avg10("io"),
    )


def pressure_ok(snapshot: PressureSnapshot, args: argparse.Namespace) -> tuple[bool, list[str]]:
    failures: list[str] = []
    checks = (
        ("cpu", snapshot.cpu, args.cpu_pressure_avg10_max),
        ("memory", snapshot.memory, args.memory_pressure_avg10_max),
        ("io", snapshot.io, args.io_pressure_avg10_max),
    )
    for name, current, maximum in checks:
        if maximum is None or current is None:
            continue
        if current > maximum:
            failures.append(f"{name} avg10={current:.2f} > {maximum:.2f}")
    return (not failures, failures)


def wait_for_pressure(args: argparse.Namespace) -> None:
    if (
        args.cpu_pressure_avg10_max is None
        and args.memory_pressure_avg10_max is None
        and args.io_pressure_avg10_max is None
    ):
        return

    deadline = time.monotonic() + args.pressure_timeout
    while True:
        snapshot = read_pressure_snapshot()
        ok, failures = pressure_ok(snapshot, args)
        if ok:
            return
        if time.monotonic() >= deadline:
            details = ", ".join(failures) or "pressure above threshold"
            raise RuntimeError(f"Timed out waiting for host pressure to settle: {details}")
        log(f"Waiting for pressure to settle: {', '.join(failures)}")
        time.sleep(args.poll_interval)


def load_targets(args: argparse.Namespace) -> tuple[list[TargetContainer], dict[str, int]]:
    container_ids = [
        line.strip()
        for line in run_command(["podman", "ps", "-aq"]).stdout.splitlines()
        if line.strip()
    ]
    if not container_ids:
        return (
            [],
            {
                "all": 0,
                "managed": 0,
                "excluded": 0,
                "restart_policy_mismatch": 0,
                "group_filtered": 0,
            },
        )

    inspect_data = json.loads(run_command(["podman", "inspect", *container_ids]).stdout)
    allowed_policies = {
        item.strip() for item in args.allowed_restart_policies.split(",") if item.strip()
    }
    allowed_groups = {item.strip() for item in args.groups.split(",") if item.strip()}
    stats = {
        "all": 0,
        "managed": 0,
        "excluded": 0,
        "restart_policy_mismatch": 0,
        "group_filtered": 0,
    }
    targets: list[TargetContainer] = []

    for item in inspect_data:
        stats["all"] += 1
        labels = ((item.get("Config") or {}).get("Labels")) or {}
        managed = parse_bool(labels.get(MANAGED_LABEL), default=False)
        if not managed:
            continue
        stats["managed"] += 1

        if parse_bool(labels.get(EXCLUDE_LABEL), default=False):
            stats["excluded"] += 1
            continue

        restart_policy = (((item.get("HostConfig") or {}).get("RestartPolicy")) or {}).get("Name")
        if allowed_policies and restart_policy not in allowed_policies:
            stats["restart_policy_mismatch"] += 1
            continue

        group = str(labels.get(GROUP_LABEL) or "default").strip()
        if allowed_groups and group not in allowed_groups:
            stats["group_filtered"] += 1
            continue

        state = item.get("State") or {}
        health = ((state.get("Health") or {}).get("Status")) or None
        targets.append(
            TargetContainer(
                container_id=item.get("Id", ""),
                name=str(item.get("Name", "")).removeprefix("/"),
                project=labels.get("io.podman.compose.project")
                or labels.get("com.docker.compose.project"),
                service=labels.get("io.podman.compose.service")
                or labels.get("com.docker.compose.service"),
                restart_policy=restart_policy,
                status=str(state.get("Status") or "unknown"),
                running=bool(state.get("Running")),
                health=health,
                healthcheck_defined=bool((item.get("Config") or {}).get("Healthcheck")),
                group=group,
                group_rank=DEFAULT_GROUP_ORDER.get(group, DEFAULT_GROUP_ORDER["default"]),
                priority=parse_int(labels.get(PRIORITY_LABEL), default=0),
                wait_for_health=parse_bool(labels.get(WAIT_FOR_HEALTH_LABEL), default=False),
            )
        )

    ordered = sorted(
        targets,
        key=lambda target: (
            -target.group_rank,
            -target.priority,
            target.project or "",
            target.service or "",
            target.name,
        ),
    )
    return (ordered, stats)


def build_batches(targets: list[TargetContainer], batch_size: int) -> list[list[TargetContainer]]:
    if batch_size < 1:
        raise ValueError("--batch-size must be at least 1")

    buckets: list[list[TargetContainer]] = []
    current: list[TargetContainer] = []
    current_key: tuple[int, int] | None = None

    for target in targets:
        if current_key is None:
            current_key = target.batch_key
        if target.batch_key != current_key:
            buckets.extend(chunked(current, batch_size))
            current = []
            current_key = target.batch_key
        current.append(target)

    if current:
        buckets.extend(chunked(current, batch_size))

    return buckets


def render_plan(targets: list[TargetContainer], stats: dict[str, int], batch_size: int) -> str:
    lines: list[str] = []
    lines.append(f"Targets discovered: {len(targets)}")
    lines.append(
        "Selection summary: "
        f"all={stats['all']} managed={stats['managed']} excluded={stats['excluded']} "
        "restart-policy-mismatch="
        f"{stats['restart_policy_mismatch']} group-filtered={stats['group_filtered']}"
    )

    if not targets:
        lines.append("")
        lines.append(
            "No containers matched. Add "
            f'`{MANAGED_LABEL}: "true"` to the services you want the controller to manage.'
        )
        return "\n".join(lines)

    batches = build_batches(targets, batch_size)
    lines.append(f"Batches: {len(batches)} (batch size {batch_size})")
    lines.append("")

    for index, batch in enumerate(batches, start=1):
        group = batch[0].group
        priority = batch[0].priority
        lines.append(f"Batch {index}: group={group} priority={priority}")
        for target in batch:
            health_text = "wait-health" if target.wait_for_health else "no-health-wait"
            running_text = "running" if target.running else target.status
            compose_bits = "/".join(part for part in (target.project, target.service) if part)
            compose_text = f" ({compose_bits})" if compose_bits else ""
            lines.append(f"  - {target.name}{compose_text} [{running_text}; {health_text}]")
        lines.append("")

    return "\n".join(lines).rstrip()


def inspect_health_status(names: list[str]) -> dict[str, tuple[str | None, bool]]:
    if not names:
        return {}
    inspect_data = json.loads(run_command(["podman", "inspect", *names]).stdout)
    result: dict[str, tuple[str | None, bool]] = {}
    for item in inspect_data:
        name = str(item.get("Name", "")).removeprefix("/")
        state = item.get("State") or {}
        health = ((state.get("Health") or {}).get("Status")) or None
        result[name] = (health, bool(state.get("Running")))
    return result


def wait_for_health(batch: list[TargetContainer], args: argparse.Namespace) -> None:
    required = [target for target in batch if target.wait_for_health]
    if not required:
        return

    missing = [target.name for target in required if not target.healthcheck_defined]
    if missing:
        raise RuntimeError(
            "Health wait requested for containers without a defined healthcheck: "
            + ", ".join(sorted(missing))
        )

    deadline = time.monotonic() + args.health_timeout
    required_names = [target.name for target in required]
    while True:
        status_map = inspect_health_status(required_names)
        unhealthy: list[str] = []
        pending: list[str] = []

        for target in required:
            health, running = status_map.get(target.name, (None, False))
            if not running:
                pending.append(f"{target.name}=not-running")
                continue
            if health == "healthy":
                continue
            if health is None:
                pending.append(f"{target.name}=no-health")
                continue
            if health == "starting":
                pending.append(f"{target.name}=starting")
                continue
            unhealthy.append(f"{target.name}={health}")

        if not pending and not unhealthy:
            return
        if time.monotonic() >= deadline:
            details = ", ".join(unhealthy + pending)
            raise RuntimeError(f"Timed out waiting for healthy batch: {details}")
        detail = ", ".join(unhealthy + pending)
        log(f"Waiting for health: {detail}")
        time.sleep(args.poll_interval)


def start_batch(batch: list[TargetContainer]) -> None:
    names_to_start = [target.name for target in batch if not target.running]
    if not names_to_start:
        log("All containers in batch are already running.")
        return
    log(f"Starting batch: {', '.join(names_to_start)}")
    run_command(["podman", "start", *names_to_start])


def command_schema(_args: argparse.Namespace) -> int:
    print(SCHEMA_TEXT)
    return 0


def command_plan(args: argparse.Namespace) -> int:
    targets, stats = load_targets(args)
    print(render_plan(targets, stats, args.batch_size))
    return 0


def command_run(args: argparse.Namespace) -> int:
    targets, stats = load_targets(args)
    log(render_plan(targets, stats, args.batch_size))

    if not targets:
        return 0

    batches = build_batches(targets, args.batch_size)
    for index, batch in enumerate(batches, start=1):
        batch_description = ", ".join(target.name for target in batch)
        log(
            f"Processing batch {index}/{len(batches)} "
            f"(group={batch[0].group} priority={batch[0].priority}): {batch_description}"
        )
        wait_for_pressure(args)
        start_batch(batch)
        wait_for_health(batch, args)
        if index != len(batches) and args.settle_seconds > 0:
            log(f"Sleeping {args.settle_seconds:.1f}s before the next batch.")
            time.sleep(args.settle_seconds)

    log("Startup controller completed successfully.")
    return 0


def main() -> int:
    args = parse_args()
    try:
        return int(args.func(args))
    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
        return 130
    except Exception as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
