#!/usr/bin/env python3
"""
pgtune_cli.py — Python port of pgtune.leopard.in.ua algorithm

Source: https://github.com/le0pard/pgtune (configurationSlice.js)
Generates PostgreSQL tuning recommendations identical to the web UI.

Usage:
    ./pgtune_cli.py --db-version 16 --os linux --db-type mixed \
                    --cpu 8 --ram 32 --ram-unit GB --connections 100 \
                    --hd-type hdd --db-size mid_ram

Output matches the web UI's "Configuration File" tab format.
"""
import argparse
import math
import sys

# ── Constants ────────────────────────────────────────────────────────────────
SIZE_UNIT_MAP = {
    "KB": 1024,
    "MB": 1024 * 1024,
    "GB": 1024 * 1024 * 1024,
    "TB": 1024 * 1024 * 1024 * 1024,
}
KB_PER_MB = 1024
KB_PER_GB = 1024 * 1024

DB_TYPE_WEB = "web"
DB_TYPE_OLTP = "oltp"
DB_TYPE_DW = "dw"
DB_TYPE_DESKTOP = "desktop"
DB_TYPE_MIXED = "mixed"

OS_LINUX = "linux"
OS_WINDOWS = "windows"
OS_MAC = "mac"

HD_SSD = "ssd"
HD_HDD = "hdd"
HD_SAN = "san"
HD_NVME = "nvme"

DB_SIZE_LESS = "less_ram"
DB_SIZE_MID = "mid_ram"
DB_SIZE_GREATER = "greater_ram"


# ── Helpers ──────────────────────────────────────────────────────────────────
def total_memory_kb(ram, ram_unit):
    return int(ram * SIZE_UNIT_MAP[ram_unit] / SIZE_UNIT_MAP["KB"])


def format_value(kb):
    """Mimics JS formatValue: prefer GB/MB only when exactly divisible, else kB."""
    if kb % KB_PER_GB == 0:
        return f"{kb // KB_PER_GB}GB"
    if kb % KB_PER_MB == 0:
        return f"{kb // KB_PER_MB}MB"
    return f"{kb}kB"


# ── Selectors (port of configurationSlice.js) ────────────────────────────────
def select_max_connections(connection_num, db_type):
    if connection_num:
        return connection_num
    return {DB_TYPE_WEB: 200, DB_TYPE_OLTP: 300, DB_TYPE_DW: 40,
            DB_TYPE_DESKTOP: 20, DB_TYPE_MIXED: 100}[db_type]


def select_shared_buffers(mem_kb, db_type, os_type, db_version):
    val = {
        DB_TYPE_WEB: mem_kb // 4,
        DB_TYPE_OLTP: mem_kb // 4,
        DB_TYPE_DW: mem_kb // 4,
        DB_TYPE_DESKTOP: mem_kb // 16,
        DB_TYPE_MIXED: mem_kb // 4,
    }[db_type]
    if db_version < 10 and os_type == OS_WINDOWS:
        win_limit = (512 * SIZE_UNIT_MAP["MB"]) // SIZE_UNIT_MAP["KB"]
        if val > win_limit:
            val = win_limit
    return val


def select_huge_pages(shared_buffers_kb, os_type):
    if os_type == OS_MAC:
        return "off"
    threshold = (2 * SIZE_UNIT_MAP["GB"]) // SIZE_UNIT_MAP["KB"]
    return "try" if shared_buffers_kb >= threshold else "off"


def select_effective_cache_size(mem_kb, db_type):
    return {
        DB_TYPE_WEB: mem_kb * 3 // 4,
        DB_TYPE_OLTP: mem_kb * 3 // 4,
        DB_TYPE_DW: mem_kb * 3 // 4,
        DB_TYPE_DESKTOP: mem_kb // 4,
        DB_TYPE_MIXED: mem_kb * 3 // 4,
    }[db_type]


def select_maintenance_work_mem(mem_kb, db_type, os_type, db_version):
    val = {
        DB_TYPE_WEB: mem_kb // 16,
        DB_TYPE_OLTP: mem_kb // 16,
        DB_TYPE_DW: mem_kb // 8,
        DB_TYPE_DESKTOP: mem_kb // 16,
        DB_TYPE_MIXED: mem_kb // 16,
    }[db_type]
    memory_limit = (8 * SIZE_UNIT_MAP["GB"]) // SIZE_UNIT_MAP["KB"]
    if os_type == OS_WINDOWS and db_version <= 17:
        memory_limit = (2 * SIZE_UNIT_MAP["GB"]) // SIZE_UNIT_MAP["KB"]
    if val >= memory_limit:
        if os_type == OS_WINDOWS and db_version <= 17:
            val = memory_limit - (1 * SIZE_UNIT_MAP["MB"]) // SIZE_UNIT_MAP["KB"]
        else:
            val = memory_limit
    return val


def select_checkpoint_segments(db_type):
    return [
        ("min_wal_size", {
            DB_TYPE_WEB: 1024 * KB_PER_MB,
            DB_TYPE_OLTP: 2048 * KB_PER_MB,
            DB_TYPE_DW: 4096 * KB_PER_MB,
            DB_TYPE_DESKTOP: 100 * KB_PER_MB,
            DB_TYPE_MIXED: 1024 * KB_PER_MB,
        }[db_type]),
        ("max_wal_size", {
            DB_TYPE_WEB: 4096 * KB_PER_MB,
            DB_TYPE_OLTP: 8192 * KB_PER_MB,
            DB_TYPE_DW: 16384 * KB_PER_MB,
            DB_TYPE_DESKTOP: 2048 * KB_PER_MB,
            DB_TYPE_MIXED: 4096 * KB_PER_MB,
        }[db_type]),
    ]


def select_wal_buffers(shared_buffers_kb):
    val = 3 * shared_buffers_kb // 100
    max_wal = 16 * KB_PER_MB
    if val > max_wal:
        val = max_wal
    near = 14 * KB_PER_MB
    if val > near and val < max_wal:
        val = max_wal
    if val < 32:
        val = 32
    return val


def select_default_statistics_target(db_type):
    return {DB_TYPE_WEB: 100, DB_TYPE_OLTP: 100, DB_TYPE_DW: 500,
            DB_TYPE_DESKTOP: 100, DB_TYPE_MIXED: 100}[db_type]


def select_random_page_cost(hd_type, db_size, db_type):
    if db_size == DB_SIZE_LESS:
        return 1.1
    if hd_type == HD_HDD:
        return 4
    if db_type == DB_TYPE_DW:
        return 4
    return 1.1


def select_effective_io_concurrency(os_type, hd_type):
    if os_type != OS_LINUX:
        return None
    return {HD_HDD: 2, HD_SSD: 200, HD_SAN: 300, HD_NVME: 1000}[hd_type]


def select_parallel_settings(db_version, db_type, cpu_num):
    if not cpu_num:
        return []
    if cpu_num < 4:
        cfg = [
            ("max_worker_processes", 8),
            ("max_parallel_workers_per_gather", 1),
        ]
        if db_version >= 10:
            cfg.append(("max_parallel_workers", max(cpu_num, 2)))
        if db_version >= 11:
            cfg.append(("max_parallel_maintenance_workers", 1))
        return cfg
    workers_per_gather = math.ceil(cpu_num / 2)
    if db_type != DB_TYPE_DW and workers_per_gather > 4:
        workers_per_gather = 4
    cfg = [
        ("max_worker_processes", max(cpu_num, 8)),
        ("max_parallel_workers_per_gather", workers_per_gather),
    ]
    if db_version >= 10:
        cfg.append(("max_parallel_workers", cpu_num))
    if db_version >= 11:
        p = math.ceil(cpu_num / 2)
        if p > 4:
            p = 4
        cfg.append(("max_parallel_maintenance_workers", p))
    return cfg


DEFAULT_DB_SETTINGS = {"default": {"max_worker_processes": 8,
                                   "max_parallel_workers_per_gather": 2,
                                   "max_parallel_workers": 8}}


def select_work_mem(mem_kb, shared_buffers_kb, max_conn, parallel_cfg,
                    db_type, os_type, db_version, db_size):
    # parallelForWorkMem
    parallel_for = 1
    if parallel_cfg:
        for k, v in parallel_cfg:
            if k == "max_worker_processes" and v and v > 0:
                parallel_for = v
                break
    else:
        default = DEFAULT_DB_SETTINGS.get(db_version, DEFAULT_DB_SETTINGS["default"])
        if default.get("max_worker_processes", 0) > 0:
            parallel_for = default["max_worker_processes"]

    work_mem_value = (mem_kb - shared_buffers_kb) / ((max_conn + parallel_for) * 3)

    result = {
        DB_TYPE_WEB: math.floor(work_mem_value),
        DB_TYPE_OLTP: math.floor(work_mem_value),
        DB_TYPE_DW: math.floor(work_mem_value / 2),
        DB_TYPE_DESKTOP: math.floor(work_mem_value / 6),
        DB_TYPE_MIXED: math.floor(work_mem_value / 2),
    }[db_type]

    if db_size == DB_SIZE_LESS:
        result = math.floor(result * 1.3)
    elif db_size == DB_SIZE_GREATER:
        result = math.floor(result * 0.9)

    min_wm = 4 * KB_PER_MB
    if result < min_wm:
        result = min_wm

    if os_type == OS_WINDOWS and db_version <= 17:
        win_limit = (2 * SIZE_UNIT_MAP["GB"] // SIZE_UNIT_MAP["KB"]) - (1 * KB_PER_MB)
        if result > win_limit:
            result = win_limit

    return result


def select_wal_level(db_type):
    if db_type == DB_TYPE_DESKTOP:
        return [("wal_level", "minimal"), ("max_wal_senders", "0")]
    return []


def select_jit(db_version, db_type):
    if db_version >= 12 and db_type in (DB_TYPE_WEB, DB_TYPE_OLTP, DB_TYPE_MIXED):
        return "off"
    return None


def select_wal_compression(db_version):
    if db_version >= 15:
        return "lz4"
    if db_version >= 10:
        return "on"
    return None


def select_autovacuum_max_workers(cpu_num):
    if not cpu_num:
        return None
    if cpu_num >= 32:
        return 5
    if cpu_num >= 16:
        return 4
    return None


def select_autovacuum_work_mem(maintenance_work_mem_kb, os_type, db_version):
    # 內部 patch（偏離 upstream pgtune leopard.in.ua）：
    # upstream 對 maintenance >= 2GB 機器一律給 2GB，但 OLTP 高峰是 3 worker × 2GB = 6GB，
    # 對 32GB 機就佔 18% RAM。改用 maintenance/4 並 cap 1GB，固定佔比 ~4.7% RAM，
    # 同時避免大機（128GB+）虛耗。下限 64MB 讓小機（4GB）也有合理推薦。
    val = min(maintenance_work_mem_kb // 4, KB_PER_GB)
    val = max(val, 64 * KB_PER_MB)
    if os_type == OS_WINDOWS and db_version <= 17:
        win_limit = (2 * SIZE_UNIT_MAP["GB"] // SIZE_UNIT_MAP["KB"]) - KB_PER_MB
        if val > win_limit:
            val = win_limit
    return val


def select_io_workers(db_version, cpu_num):
    if db_version < 18 or not cpu_num:
        return None
    val = min(32, max(3, cpu_num // 4))
    return val if val > 3 else None


def select_io_method(db_version, os_type):
    if db_version < 18:
        return None
    if os_type == OS_LINUX:
        return "io_uring"
    return "worker"


def select_warnings(total_memory_bytes, wal_comp, io_method, db_type, hd_type, db_size):
    w = []
    if total_memory_bytes < 256 * SIZE_UNIT_MAP["MB"]:
        w += ["this tool not being optimal", "for low memory systems"]
    elif total_memory_bytes > 100 * SIZE_UNIT_MAP["GB"]:
        w += ["this tool not being optimal", "for very high memory systems"]
    if wal_comp == "lz4":
        if w:
            w.append("")
        w += ["wal_compression = lz4 requires PostgreSQL", "to be compiled with --with-lz4"]
    if io_method == "io_uring":
        if w:
            w.append("")
        w += ["io_method = io_uring requires PostgreSQL", "to be compiled with --with-liburing"]
    if db_type == DB_TYPE_DW and hd_type != HD_HDD and db_size != DB_SIZE_LESS:
        name = {HD_SSD: "SSDs", HD_NVME: "NVMe drives", HD_SAN: "SAN storage"}[hd_type]
        if w:
            w.append("")
        w += [
            f"Cost parameters for Data Warehouses on {name} are left at defaults",
            "to avoid catastrophic index scan selections",
            "Monitor query planner behavior and adjust random_page_cost if necessary",
        ]
    return ["WARNING"] + w if w else []


# ── Main generator ───────────────────────────────────────────────────────────
def generate(args):
    mem_bytes = args.ram * SIZE_UNIT_MAP[args.ram_unit]
    mem_kb = int(mem_bytes / SIZE_UNIT_MAP["KB"])

    max_conn = select_max_connections(args.connections, args.db_type)
    shared_buffers = select_shared_buffers(mem_kb, args.db_type, args.os, args.db_version)
    huge_pages = select_huge_pages(shared_buffers, args.os)
    eff_cache = select_effective_cache_size(mem_kb, args.db_type)
    mwm = select_maintenance_work_mem(mem_kb, args.db_type, args.os, args.db_version)
    ckpt_seg = select_checkpoint_segments(args.db_type)
    wal_level = select_wal_level(args.db_type)
    ckpt_comp_target = 0.9
    wal_buffers = select_wal_buffers(shared_buffers)
    dst = select_default_statistics_target(args.db_type)
    rpc = select_random_page_cost(args.hd_type, args.db_size, args.db_type)
    eic = select_effective_io_concurrency(args.os, args.hd_type)
    parallel = select_parallel_settings(args.db_version, args.db_type, args.cpu)
    work_mem = select_work_mem(mem_kb, shared_buffers, max_conn, parallel,
                                args.db_type, args.os, args.db_version, args.db_size)
    jit = select_jit(args.db_version, args.db_type)
    wal_comp = select_wal_compression(args.db_version)
    avw = select_autovacuum_max_workers(args.cpu)
    avwm = select_autovacuum_work_mem(mwm, args.os, args.db_version)
    io_method = select_io_method(args.db_version, args.os)
    io_workers = select_io_workers(args.db_version, args.cpu)
    warnings = select_warnings(mem_bytes, wal_comp, io_method, args.db_type, args.hd_type, args.db_size)

    # Format config data (in web UI order)
    config = [
        ("max_connections", max_conn),
        ("shared_buffers", format_value(shared_buffers)),
        ("effective_cache_size", format_value(eff_cache)),
        ("maintenance_work_mem", format_value(mwm)),
        ("checkpoint_completion_target", ckpt_comp_target),
        ("wal_buffers", format_value(wal_buffers)),
        ("default_statistics_target", dst),
        ("random_page_cost", rpc),
        ("effective_io_concurrency", eic),
        ("work_mem", format_value(work_mem)),
        ("huge_pages", huge_pages),
        ("jit", jit),
        ("wal_compression", wal_comp),
        ("autovacuum_max_workers", avw),
        ("autovacuum_work_mem", format_value(avwm) if avwm else None),
        ("io_method", io_method),
        ("io_workers", io_workers),
    ]
    # Append checkpoint_segments (formatted) + parallel + wal_level
    for k, v in ckpt_seg:
        config.append((k, format_value(v)))
    for k, v in parallel:
        config.append((k, v))
    for k, v in wal_level:
        config.append((k, v))

    # Print output matching web UI format
    lines = []
    if warnings:
        for w in warnings:
            lines.append(f"# {w}")
        lines.append("")
    lines.append(f"# DB Version: {args.db_version}")
    lines.append(f"# OS Type: {args.os}")
    lines.append(f"# DB Type: {args.db_type}")
    ram_display = int(args.ram) if args.ram == int(args.ram) else args.ram
    lines.append(f"# Total Memory (RAM): {ram_display} {args.ram_unit}")
    lines.append(f"# CPUs num: {args.cpu}")
    if args.connections:
        lines.append(f"# Connections num: {args.connections}")
    lines.append(f"# Data Storage: {args.hd_type}")
    lines.append("")
    for k, v in config:
        if v is None or v == "":
            continue
        lines.append(f"{k} = {v}")
    return "\n".join(lines)


def main():
    p = argparse.ArgumentParser(description="PostgreSQL config tuning (pgtune algorithm)")
    p.add_argument("--db-version", type=int, default=16, choices=[10,11,12,13,14,15,16,17,18])
    p.add_argument("--os", default=OS_LINUX, choices=[OS_LINUX, OS_WINDOWS, OS_MAC])
    p.add_argument("--db-type", default=DB_TYPE_WEB,
                   choices=[DB_TYPE_WEB, DB_TYPE_OLTP, DB_TYPE_DW, DB_TYPE_DESKTOP, DB_TYPE_MIXED])
    p.add_argument("--ram", type=float, required=True, help="Total RAM (e.g. 32)")
    p.add_argument("--ram-unit", default="GB", choices=["MB", "GB"])
    p.add_argument("--cpu", type=int, required=True, help="CPU core count")
    p.add_argument("--connections", type=int, default=None, help="Max connections (blank = use default for db-type)")
    p.add_argument("--hd-type", default=HD_HDD, choices=[HD_SSD, HD_HDD, HD_SAN, HD_NVME])
    p.add_argument("--db-size", default=DB_SIZE_MID,
                   choices=[DB_SIZE_LESS, DB_SIZE_MID, DB_SIZE_GREATER])
    args = p.parse_args()
    print(generate(args))


if __name__ == "__main__":
    main()
