"""Hardware & model profiling: system info, VRAM/RAM/HDD, GFLOPs/MACs.

Uses fvcore, thop, psutil, and torch.cuda for comprehensive reporting.
"""

from __future__ import annotations

import logging
import platform
import shutil
from pathlib import Path
from typing import Any

import psutil
import torch
import torch.nn as nn

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Formatting helpers
# ---------------------------------------------------------------------------

def _fmt_bytes(b: float) -> str:
    """Human-readable byte size."""
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if abs(b) < 1024:
            return f"{b:.1f} {unit}"
        b /= 1024
    return f"{b:.1f} PB"


def _fmt_number(n: float) -> str:
    """Human-readable large number (1.2M, 3.4B, etc.)."""
    if n >= 1e9:
        return f"{n / 1e9:.2f} B"
    if n >= 1e6:
        return f"{n / 1e6:.2f} M"
    if n >= 1e3:
        return f"{n / 1e3:.1f} K"
    return str(int(n))


# ---------------------------------------------------------------------------
# System info
# ---------------------------------------------------------------------------

def log_system_info() -> dict[str, Any]:
    """Log and return system hardware summary."""
    info: dict[str, Any] = {}

    # CPU
    cpu_name = platform.processor() or platform.machine()
    cpu_cores_phys = psutil.cpu_count(logical=False) or 0
    cpu_cores_logic = psutil.cpu_count(logical=True) or 0
    info["cpu"] = {
        "name": cpu_name,
        "cores_physical": cpu_cores_phys,
        "cores_logical": cpu_cores_logic,
    }

    # RAM
    mem = psutil.virtual_memory()
    info["ram"] = {
        "total": mem.total,
        "available": mem.available,
        "used": mem.used,
        "percent": mem.percent,
    }

    # GPU
    if torch.cuda.is_available():
        gpu_id = torch.cuda.current_device()
        gpu_name = torch.cuda.get_device_name(gpu_id)
        gpu_total = torch.cuda.get_device_properties(gpu_id).total_memory
        gpu_reserved = torch.cuda.memory_reserved(gpu_id)
        gpu_allocated = torch.cuda.memory_allocated(gpu_id)
        gpu_free = gpu_total - gpu_reserved
        info["gpu"] = {
            "name": gpu_name,
            "total": gpu_total,
            "allocated": gpu_allocated,
            "reserved": gpu_reserved,
            "free": gpu_free,
        }
    else:
        info["gpu"] = None

    # Log it
    logger.info("🖥️  System info:")
    logger.info("  🧠 CPU: %s (%d physical / %d logical cores)",
                cpu_name, cpu_cores_phys, cpu_cores_logic)
    logger.info("  💾 RAM: %s used / %s total (%.1f%% used)",
                _fmt_bytes(mem.used), _fmt_bytes(mem.total), mem.percent)

    if info["gpu"]:
        g = info["gpu"]
        logger.info("  🎮 GPU: %s", g["name"])
        logger.info("     VRAM: %s allocated / %s reserved / %s total (%s free)",
                     _fmt_bytes(g["allocated"]), _fmt_bytes(g["reserved"]),
                     _fmt_bytes(g["total"]), _fmt_bytes(g["free"]))
    else:
        logger.info("  🎮 GPU: not available")

    return info


def log_disk_info(*paths: Path) -> dict[str, Any]:
    """Log disk usage for given paths (or / by default)."""
    targets = paths or (Path("/"),)
    info: dict[str, Any] = {}

    for p in targets:
        try:
            usage = shutil.disk_usage(p)
            key = str(p)
            info[key] = {
                "total": usage.total,
                "used": usage.used,
                "free": usage.free,
            }
            pct = usage.used / usage.total * 100
            logger.info("  💿 Disk [%s]: %s used / %s total (%.1f%%, %s free)",
                        p, _fmt_bytes(usage.used), _fmt_bytes(usage.total),
                        pct, _fmt_bytes(usage.free))
        except OSError as e:
            logger.warning("  💿 Disk [%s]: unavailable (%s)", p, e)
    return info


# ---------------------------------------------------------------------------
# VRAM snapshot
# ---------------------------------------------------------------------------

def log_vram_snapshot(label: str = "") -> dict[str, float] | None:
    """Log current VRAM usage. Returns dict or None if no GPU."""
    if not torch.cuda.is_available():
        return None

    gpu_id = torch.cuda.current_device()
    allocated = torch.cuda.memory_allocated(gpu_id)
    reserved = torch.cuda.memory_reserved(gpu_id)
    total = torch.cuda.get_device_properties(gpu_id).total_memory
    free = total - reserved

    prefix = f"[{label}] " if label else ""
    logger.info("  🎮 %sVRAM: %s allocated / %s reserved / %s free",
                prefix, _fmt_bytes(allocated), _fmt_bytes(reserved),
                _fmt_bytes(free))
    return {"allocated": allocated, "reserved": reserved, "free": free, "total": total}


def log_ram_snapshot(label: str = "") -> dict[str, float]:
    """Log current RAM usage."""
    mem = psutil.virtual_memory()
    prefix = f"[{label}] " if label else ""
    logger.info("  💾 %sRAM: %s used / %s available / %s total (%.1f%%)",
                prefix, _fmt_bytes(mem.used), _fmt_bytes(mem.available),
                _fmt_bytes(mem.total), mem.percent)
    return {"used": mem.used, "available": mem.available,
            "total": mem.total, "percent": mem.percent}


# ---------------------------------------------------------------------------
# Model profiling: GFLOPs / MACs / params
# ---------------------------------------------------------------------------

def profile_model(
    model: nn.Module,
    input_size: tuple[int, ...] | None,
    device: torch.device | None = None,
    model_name: str = "",
) -> dict[str, Any]:
    """Profile model FLOPs/MACs/params using fvcore + thop.

    Args:
        model: PyTorch model.
        input_size: (B, C, H, W) input shape, or None to skip FLOPs/MACs.
        device: Device for dummy input.
        model_name: Label for log messages.

    Returns:
        Dict with params, flops_fvcore, macs_thop, etc.
    """
    label = model_name or model.__class__.__name__
    result: dict[str, Any] = {"name": label}

    # Params
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    result["total_params"] = total_params
    result["trainable_params"] = trainable_params
    result["param_bytes"] = param_bytes

    logger.info("📊 Model profile: %s", label)
    logger.info("  📐 Parameters: %s total (%s trainable), %s in memory",
                _fmt_number(total_params), _fmt_number(trainable_params),
                _fmt_bytes(param_bytes))

    if input_size is None:
        result["flops_fvcore"] = None
        result["macs_thop"] = None
        return result

    # Determine device and dtype from model
    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    dtype = next(model.parameters()).dtype if total_params > 0 else torch.float32

    dummy = torch.randn(*input_size, device=device, dtype=dtype)

    # fvcore FLOPs
    try:
        from fvcore.nn import FlopCountAnalysis
        flop_counter = FlopCountAnalysis(model, dummy)
        flop_counter.unsupported_ops_warnings(False)
        flop_counter.uncalled_modules_warnings(False)
        flops = flop_counter.total()
        result["flops_fvcore"] = flops
        result["gflops_fvcore"] = flops / 1e9
        logger.info("  ⚡ FLOPs (fvcore): %s (%.2f GFLOPs)",
                     _fmt_number(flops), flops / 1e9)
    except Exception as e:
        logger.warning("  ⚡ fvcore FLOPs failed: %s", e)
        result["flops_fvcore"] = None

    # thop MACs
    try:
        from thop import profile as thop_profile
        macs, params = thop_profile(model, inputs=(dummy,), verbose=False)
        result["macs_thop"] = macs
        result["gmacs_thop"] = macs / 1e9
        logger.info("  ⚡ MACs (thop):   %s (%.2f GMACs)",
                     _fmt_number(macs), macs / 1e9)
    except Exception as e:
        logger.warning("  ⚡ thop MACs failed: %s", e)
        result["macs_thop"] = None

    return result
