#!/usr/bin/env python3
"""Composite safety formula for opendeviationbar-py parallel worker calculation.

Determines the safe number of concurrent populate workers on bigblack by taking
the minimum across four capacity dimensions: ClickHouse merge capacity, network
bandwidth, available memory, and CPU cores.

Usage:
    # With defaults (bigblack: 32 cores, 61 GB, 1 Gbps, light threshold)
    python capacity_model.py

    # Override for heavy threshold (250 dbps, BTC/ETH)
    python capacity_model.py --per-worker-mem 5.0

    # Override for very heavy threshold (100 dbps)
    python capacity_model.py --per-worker-mem 8.0

    # Custom hardware
    python capacity_model.py --total-ram 128 --cores 64 --link-bw 250

    # With current service load already measured
    python capacity_model.py --service-mem 5.2 --service-cores 3
"""

from __future__ import annotations

import argparse
import math
import sys


def calculate_safe_workers(
    *,
    # ClickHouse
    merge_pool_size: int = 16,
    merge_efficiency: float = 0.6,
    # Network
    link_bw_mbps: float = 125.0,  # 1 Gbps = 125 MB/s
    per_worker_bw_mbps: float = 15.0,
    net_safety_margin: float = 0.8,
    # Memory
    total_ram_gb: float = 61.0,
    service_mem_gb: float = 4.817,  # CH 2.3 + sidecar 0.326 + kintsugi 1.6 + backfill 0.591
    headroom_gb: float = 4.0,
    per_worker_mem_gb: float = 0.2,  # Light threshold; 5.0 for BTC@250, 8.0 for @100
    # CPU
    total_cores: int = 32,
    service_cores: float = 2.0,
    cpu_reserve_factor: float = 0.5,  # Reserve 50% for ClickHouse merges
    per_worker_cores: float = 0.5,
) -> dict[str, int | str]:
    """Calculate safe worker count across all four dimensions."""

    # Dimension 1: ClickHouse merge capacity
    ch_workers = math.floor(merge_pool_size * merge_efficiency)

    # Dimension 2: Network bandwidth
    net_workers = math.floor(link_bw_mbps / per_worker_bw_mbps * net_safety_margin)

    # Dimension 3: Memory
    available_ram = total_ram_gb - service_mem_gb - headroom_gb
    if available_ram <= 0:
        mem_workers = 0
    else:
        mem_workers = math.floor(available_ram / per_worker_mem_gb)

    # Dimension 4: CPU
    available_cores = total_cores * cpu_reserve_factor - service_cores
    if available_cores <= 0:
        cpu_workers = 0
    else:
        cpu_workers = math.floor(available_cores / per_worker_cores)

    # Composite: minimum across all dimensions
    safe = min(ch_workers, net_workers, mem_workers, cpu_workers)
    safe = max(safe, 1)  # At least 1 worker

    # Identify binding constraint
    values = {
        "ClickHouse merge": ch_workers,
        "Network bandwidth": net_workers,
        "Memory": mem_workers,
        "CPU": cpu_workers,
    }
    binding = min(values, key=values.get)  # type: ignore[arg-type]

    return {
        "safe_workers": safe,
        "binding_constraint": binding,
        "ch_workers": ch_workers,
        "net_workers": net_workers,
        "mem_workers": mem_workers,
        "cpu_workers": cpu_workers,
        "available_ram_gb": round(available_ram, 1),
    }


def estimate_eta(
    total_symbol_days: int,
    num_workers: int,
    blended_rate_s_per_day: float = 1.8,
    contingency: float = 0.2,
) -> dict[str, float]:
    """Estimate backfill completion time."""
    base_seconds = total_symbol_days * blended_rate_s_per_day / num_workers
    with_contingency = base_seconds * (1 + contingency)
    return {
        "base_hours": round(base_seconds / 3600, 1),
        "with_contingency_hours": round(with_contingency / 3600, 1),
        "total_symbol_days": total_symbol_days,
        "num_workers": num_workers,
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Calculate safe parallel worker count for opendeviationbar backfill"
    )

    # Hardware
    parser.add_argument("--total-ram", type=float, default=61.0, help="Total RAM in GB (default: 61)")
    parser.add_argument("--cores", type=int, default=32, help="Total CPU cores (default: 32)")
    parser.add_argument("--link-bw", type=float, default=125.0, help="Link bandwidth in MB/s (default: 125)")

    # ClickHouse
    parser.add_argument("--merge-pool", type=int, default=16, help="background_pool_size (default: 16)")
    parser.add_argument("--merge-efficiency", type=float, default=0.6, help="Merge efficiency factor (default: 0.6)")

    # Per-worker resources
    parser.add_argument(
        "--per-worker-mem",
        type=float,
        default=0.2,
        help="Per-worker memory in GB (default: 0.2 for light threshold; use 5.0 for BTC@250, 8.0 for @100)",
    )
    parser.add_argument("--per-worker-bw", type=float, default=15.0, help="Per-worker bandwidth MB/s (default: 15)")
    parser.add_argument("--per-worker-cores", type=float, default=0.5, help="Per-worker CPU cores (default: 0.5)")

    # Service overhead
    parser.add_argument(
        "--service-mem",
        type=float,
        default=4.817,
        help="Total service memory in GB (default: 4.817 = CH + sidecar + kintsugi + backfill)",
    )
    parser.add_argument("--service-cores", type=float, default=2.0, help="Service CPU cores (default: 2)")
    parser.add_argument("--headroom", type=float, default=4.0, help="OS headroom in GB (default: 4)")

    # ETA
    parser.add_argument("--eta-days", type=int, default=0, help="Total symbol-days for ETA estimate")
    parser.add_argument("--eta-rate", type=float, default=1.8, help="Blended rate s/day (default: 1.8)")

    args = parser.parse_args()

    result = calculate_safe_workers(
        merge_pool_size=args.merge_pool,
        merge_efficiency=args.merge_efficiency,
        link_bw_mbps=args.link_bw,
        per_worker_bw_mbps=args.per_worker_bw,
        total_ram_gb=args.total_ram,
        service_mem_gb=args.service_mem,
        headroom_gb=args.headroom,
        per_worker_mem_gb=args.per_worker_mem,
        total_cores=args.cores,
        service_cores=args.service_cores,
        per_worker_cores=args.per_worker_cores,
    )

    print("=" * 60)
    print("  Composite Safety Formula — Safe Worker Count")
    print("=" * 60)
    print()
    print(f"  ClickHouse merge capacity:  {result['ch_workers']:>3} workers")
    print(f"  Network bandwidth:          {result['net_workers']:>3} workers")
    print(f"  Memory ({result['available_ram_gb']} GB available): {result['mem_workers']:>3} workers")
    print(f"  CPU:                        {result['cpu_workers']:>3} workers")
    print()
    print(f"  >>> SAFE WORKERS: {result['safe_workers']}  (binding: {result['binding_constraint']})")
    print()

    if args.per_worker_mem > 1.0:
        print(f"  NOTE: Using heavy threshold memory ({args.per_worker_mem} GB/worker)")
        print()

    if args.eta_days > 0:
        eta = estimate_eta(args.eta_days, result["safe_workers"], args.eta_rate)
        print("-" * 60)
        print(f"  ETA Estimate:")
        print(f"    Symbol-days:       {eta['total_symbol_days']:,}")
        print(f"    Workers:           {eta['num_workers']}")
        print(f"    Base estimate:     {eta['base_hours']} hours")
        print(f"    With contingency:  {eta['with_contingency_hours']} hours (+20%)")
        print()


if __name__ == "__main__":
    main()
