"""MPS unified memory monitoring, pressure callbacks, and chunk size estimation."""

from __future__ import annotations

import subprocess
from dataclasses import dataclass

import psutil
from loguru import logger

# Minimum hardware requirements
MIN_RAM_4K = 16 * 1024**3      # 16 GB for 4K
MIN_RAM_1080P = 8 * 1024**3    # 8 GB for 1080p
CHUNK_SIZE_LOWER_BOUND = 8     # frames


@dataclass
class MemoryInfo:
    allocated_bytes: int
    total_bytes: int
    utilization: float
    recommended_free_bytes: int


class MemoryPressureError(RuntimeError):
    pass


class GPUManager:
    SAFETY_THRESHOLD: float = 0.90

    def __init__(self, threshold: float = 0.90) -> None:
        self.threshold = threshold
        self._total_bytes: int | None = None

    def get_total_memory(self) -> int:
        """Get total system memory via sysctl (cached)."""
        if self._total_bytes is None:
            result = subprocess.run(
                ["sysctl", "-n", "hw.memsize"],
                capture_output=True, text=True, check=True,
            )
            self._total_bytes = int(result.stdout.strip())
        return self._total_bytes

    def get_gpu_budget(self) -> int:
        """Estimate GPU-available memory from unified memory.

        Uses Metal's heuristic: ~75% of total, cross-validated with
        psutil available memory for dynamic accuracy.
        """
        total = self.get_total_memory()
        static_budget = int(total * 0.75)

        vm = psutil.virtual_memory()
        dynamic_budget = vm.available
        return min(static_budget, dynamic_budget)

    def get_memory_usage(self) -> MemoryInfo:
        """Return current MPS memory usage."""
        try:
            import torch
            allocated = torch.mps.current_allocated_memory()
        except Exception:
            allocated = 0

        budget = self.get_gpu_budget()
        utilization = allocated / budget if budget > 0 else 0.0
        safe_cap = int(budget * self.threshold)
        recommended_free = max(0, safe_cap - allocated)

        return MemoryInfo(
            allocated_bytes=allocated,
            total_bytes=budget,
            utilization=utilization,
            recommended_free_bytes=recommended_free,
        )

    def estimate_chunk_size(
        self,
        frame_height: int,
        frame_width: int,
        model: str = "propainter",
    ) -> int:
        """Estimate safe chunk frame count based on resolution and available memory.

        Uses a x6 multiplier for ProPainter intermediate tensors (conservative).
        """
        multiplier = 6 if model == "propainter" else 3
        per_frame_bytes = frame_height * frame_width * 3 * 4 * multiplier

        if per_frame_bytes == 0:
            return CHUNK_SIZE_LOWER_BOUND

        budget = self.get_gpu_budget()
        safe_budget = int(budget * self.threshold)

        # SAM ~2.3GB, ProPainter ~1.5GB
        model_overhead = 1_500_000_000 if model == "propainter" else 2_300_000_000
        available = safe_budget - model_overhead

        if available <= 0:
            return CHUNK_SIZE_LOWER_BOUND

        chunk_size = available // per_frame_bytes
        return max(CHUNK_SIZE_LOWER_BOUND, chunk_size)

    def check_pressure(self) -> bool:
        """Return True if memory utilization exceeds threshold."""
        info = self.get_memory_usage()
        return info.utilization > self.threshold

    def clear_cache(self) -> None:
        """Call torch.mps.empty_cache() and synchronize."""
        try:
            import torch
            torch.mps.synchronize()
            torch.mps.empty_cache()
        except Exception as e:
            logger.warning(f"clear_cache failed: {e}")

    def adaptive_reduce(self, current_chunk_size: int) -> int:
        """Reduce chunk size under memory pressure (x0.5, lower bound 8)."""
        reduced = max(CHUNK_SIZE_LOWER_BOUND, current_chunk_size // 2)
        logger.warning(f"Adaptive reduce: {current_chunk_size} → {reduced} frames")
        return reduced

    def check_hardware_requirements(self, frame_height: int, frame_width: int) -> tuple[bool, float | None]:
        """Check if hardware meets minimum requirements for given resolution.

        Returns (meets_requirement, suggested_scale_factor).
        scale_factor is None if requirements are met, otherwise 0.5 or lower.
        """
        total = self.get_total_memory()
        is_4k = frame_height >= 2160 or frame_width >= 3840

        if is_4k and total < MIN_RAM_4K:
            scale = 0.5
            logger.warning(
                f"4K requires {MIN_RAM_4K // 1024**3}GB RAM, "
                f"found {total / 1024**3:.0f}GB. Suggesting {scale}x downscale."
            )
            return False, scale

        if total < MIN_RAM_1080P:
            scale = 0.25
            logger.warning(
                f"Minimum {MIN_RAM_1080P // 1024**3}GB RAM required, "
                f"found {total / 1024**3:.0f}GB. Suggesting {scale}x downscale."
            )
            return False, scale

        return True, None
