"""Module resource estimation for cluster capacity planning.

Provides per-module resource estimates based on Helm chart default requests.
Used for informational warnings only — never blocks user operations.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ResourceEstimate:
    """Estimated resource consumption of a module."""

    memory_mi: int  # MiB (e.g. 512 = 512 MiB)
    cpu_millicores: int  # (e.g. 500 = 0.5 CPU)


# Per-module estimates based on Helm chart defaults (requests, not limits).
MODULE_RESOURCES: dict[str, ResourceEstimate] = {
    "core": ResourceEstimate(memory_mi=2048, cpu_millicores=1000),
    "ui": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "mongodb": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "minio": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "postgresql": ResourceEstimate(memory_mi=256, cpu_millicores=250),
    "jupyter": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "gitea": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "sql": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "hive-metastore": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "superset": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "airflow": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "trino": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "monitoring": ResourceEstimate(memory_mi=768, cpu_millicores=250),
    "loki": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "marquez": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "kafka": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "clickhouse": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "mlflow": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "langfuse": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "jupyterhub": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "livy-proxy": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "zeppelin": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "unity-catalog": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "mageai": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "streamlit": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "nifi": ResourceEstimate(memory_mi=1024, cpu_millicores=500),
    "n8n": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "kestra": ResourceEstimate(memory_mi=512, cpu_millicores=250),
    "nessie": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "hydra": ResourceEstimate(memory_mi=256, cpu_millicores=100),
    "openldap": ResourceEstimate(memory_mi=128, cpu_millicores=100),
}

_DEFAULT_ESTIMATE = ResourceEstimate(memory_mi=128, cpu_millicores=100)


def get_module_estimate(module_name: str) -> ResourceEstimate:
    """Return the resource estimate for a module, with a safe default."""
    return MODULE_RESOURCES.get(module_name, _DEFAULT_ESTIMATE)


def estimate_total(modules: list[str]) -> ResourceEstimate:
    """Sum resource estimates for a list of modules (no double-counting)."""
    unique = set(modules)
    total_mem = sum(get_module_estimate(m).memory_mi for m in unique)
    total_cpu = sum(get_module_estimate(m).cpu_millicores for m in unique)
    return ResourceEstimate(memory_mi=total_mem, cpu_millicores=total_cpu)


def format_estimate(est: ResourceEstimate) -> str:
    """Format a ResourceEstimate as a human-readable string.

    Examples: '~2.0 GiB RAM, 1.0 CPU' or '~512 MiB RAM, 0.5 CPU'
    """
    if est.memory_mi >= 1024:
        mem_str = f"~{est.memory_mi / 1024:.1f} GiB"
    else:
        mem_str = f"~{est.memory_mi} MiB"
    cpu_str = f"{est.cpu_millicores / 1000:.1f} CPU"
    return f"{mem_str} RAM, {cpu_str}"


@dataclass
class ResourceCheckResult:
    """Result of a resource check against cluster capacity."""

    estimated: ResourceEstimate
    allocatable_memory_mi: int | None = None  # None if cluster unreachable
    allocatable_cpu_millicores: int | None = None
    sufficient: bool = True  # True if available > estimated (or unknown)
    utilization_pct: float | None = None  # e.g. 71.0
    message: str = ""


def check_resources_for_modules(
    modules: list[str],
    cluster_resources: dict | None = None,
) -> ResourceCheckResult:
    """Check if cluster has sufficient resources for the given modules.

    Parameters
    ----------
    modules:
        List of module names to check.
    cluster_resources:
        Dict with 'allocatable_memory_bytes' and 'allocatable_cpu_millicores' keys.
        Pass None if cluster is unreachable.

    Returns
    -------
    ResourceCheckResult with estimated usage and capacity assessment.
    """
    estimated = estimate_total(modules)

    if cluster_resources is None:
        return ResourceCheckResult(
            estimated=estimated,
            sufficient=True,
            message=f"Estimated: {format_estimate(estimated)} (cluster capacity unknown)",
        )

    alloc_mem_mi = cluster_resources.get("allocatable_memory_bytes", 0) // (1024 * 1024)
    alloc_cpu = cluster_resources.get("allocatable_cpu_millicores", 0)

    if alloc_mem_mi == 0:
        return ResourceCheckResult(
            estimated=estimated,
            sufficient=True,
            message=f"Estimated: {format_estimate(estimated)} (cluster capacity unknown)",
        )

    mem_pct = (estimated.memory_mi / alloc_mem_mi) * 100 if alloc_mem_mi > 0 else 0
    sufficient = mem_pct < 80  # Leave 20% headroom

    alloc_str = f"{alloc_mem_mi / 1024:.1f} GiB" if alloc_mem_mi >= 1024 else f"{alloc_mem_mi} MiB"

    if sufficient:
        message = (
            f"Estimated: {format_estimate(estimated)} / {alloc_str} allocatable ({mem_pct:.0f}%)"
        )
    else:
        message = (
            f"Estimated resource usage: {format_estimate(estimated)}. "
            f"Cluster allocatable: {alloc_str} RAM. "
            f"Modules may experience memory pressure. Consider disabling "
            f"non-essential modules or adding cluster resources."
        )

    return ResourceCheckResult(
        estimated=estimated,
        allocatable_memory_mi=alloc_mem_mi,
        allocatable_cpu_millicores=alloc_cpu,
        sufficient=sufficient,
        utilization_pct=round(mem_pct, 1),
        message=message,
    )
