"""test_gpu_memory.py — GPU memory budget tests for social-bot stack.

Requirements (Orin Nano Super 8GB shared memory):
  - Full stack allocation must stay under 6GB total GPU usage
  - Each major model budget:
    Face detection (SCRFD-10GF TRT):     < 100 MB
    Face recognition (ArcFace TRT):      < 300 MB
    Speaker embedding (ECAPA-TDNN TRT):  < 200 MB
    Whisper STT (CTranslate2):           < 700 MB
    LLM (Phi-3-mini Q4 GGUF, 20 layers): < 2500 MB
    Piper TTS (ONNX CPU):                < 200 MB
    Tracking (Kalman, CPU):              < 10 MB
  ─────────────────────────────────────────────────
  Total budget:                          < 6000 MB  (6 GB)
  Reserve for OS + ROS2:                  2000 MB  (2 GB)

Tests:
  1. GPU available check (skip gracefully in CI without GPU)
  2. Used memory at baseline (before stack loads) < 1GB
  3. Used memory during full stack < 6GB
  4. Memory does not grow continuously (no leak over 30s)
  5. After shutdown, memory returns to near-baseline
"""

from __future__ import annotations

import os
import sys
import time
import unittest

import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from saltybot_social_tests.test_helpers import GpuMemoryChecker

# ── Budget constants ──────────────────────────────────────────────────────────
GPU_TOTAL_MB_EXPECTED = 8192.0    # Orin Nano Super 8GB
GPU_STACK_BUDGET_MB   = 6000.0    # full stack must stay under 6GB used
GPU_BASELINE_MB       = 1024.0    # baseline (OS + CUDA context) < 1GB
GPU_LEAK_TOLERANCE_MB =  200.0    # acceptable growth over 30s measurement window

# Per-model expected budgets (for documentation / CI without full stack)
MODEL_BUDGETS_MB = {
    "scrfd_10g_trt":   100,
    "arcface_r100_trt": 300,
    "ecapa_tdnn_trt":  200,
    "whisper_small":   700,
    "phi3_mini_gguf": 2500,
    "piper_tts":       200,
    "tracking":         10,
}
TOTAL_MODEL_BUDGET_MB = sum(MODEL_BUDGETS_MB.values())  # 4010 MB


def _skip_if_no_gpu():
    if not GpuMemoryChecker.gpu_available():
        pytest.skip("No GPU available — skipping GPU memory tests")


class TestGpuMemoryBudget(unittest.TestCase):
    """GPU memory budget tests for the full social-bot stack."""

    def test_gpu_is_available_or_ci(self):
        """GPU check: either GPU is available, or we're in CI (skip gracefully)."""
        if os.environ.get("CI", ""):
            # In CI without GPU, just log and pass
            print("[INFO] CI mode detected — GPU memory tests will be skipped if no GPU")
        available = GpuMemoryChecker.gpu_available()
        if not available:
            self.skipTest("No GPU available")
        total = GpuMemoryChecker.total_mb()
        self.assertIsNotNone(total)
        print(f"[INFO] GPU total: {total:.0f} MB")

    def test_gpu_total_memory_matches_orin(self):
        """Total GPU memory should be ~8GB on Orin Nano Super."""
        _skip_if_no_gpu()
        total = GpuMemoryChecker.total_mb()
        # Allow ±10% from 8192MB (shared memory varies by config)
        self.assertGreater(total, GPU_TOTAL_MB_EXPECTED * 0.5,
                           f"GPU total {total:.0f}MB seems too low for Orin")

    def test_baseline_memory_under_threshold(self):
        """Baseline GPU usage (before models load) must be < 1GB."""
        _skip_if_no_gpu()
        used = GpuMemoryChecker.used_mb()
        self.assertIsNotNone(used)
        print(f"[INFO] Baseline GPU used: {used:.0f} MB")
        # Note: if running after stack is up, this test would need stack to be down
        # In standalone mode (no stack running), assert low baseline
        if os.environ.get("SOCIAL_STACK_RUNNING", "0") == "0":
            self.assertLess(
                used, GPU_BASELINE_MB,
                f"Baseline GPU usage {used:.0f}MB >= {GPU_BASELINE_MB}MB"
            )

    def test_full_stack_memory_under_budget(self):
        """Full stack GPU usage must stay under 6GB."""
        _skip_if_no_gpu()
        if os.environ.get("SOCIAL_STACK_RUNNING", "0") != "1":
            self.skipTest("Set SOCIAL_STACK_RUNNING=1 to run with stack active")

        used = GpuMemoryChecker.used_mb()
        self.assertIsNotNone(used)
        print(f"[INFO] Stack GPU used: {used:.0f} MB / {GPU_STACK_BUDGET_MB:.0f} MB budget")
        self.assertLess(
            used, GPU_STACK_BUDGET_MB,
            f"GPU usage {used:.0f}MB exceeds {GPU_STACK_BUDGET_MB:.0f}MB budget"
        )

    def test_no_memory_leak_over_30s(self):
        """GPU memory must not grow by more than 200MB over 30 seconds."""
        _skip_if_no_gpu()
        if os.environ.get("SOCIAL_STACK_RUNNING", "0") != "1":
            self.skipTest("Set SOCIAL_STACK_RUNNING=1 to run with stack active")

        used_start = GpuMemoryChecker.used_mb()
        print(f"[INFO] GPU used at start: {used_start:.0f} MB")

        # Wait 30s while stack processes mock sensor data
        time.sleep(30.0)

        used_end = GpuMemoryChecker.used_mb()
        growth = used_end - used_start
        print(f"[INFO] GPU used after 30s: {used_end:.0f} MB (growth: {growth:+.0f} MB)")

        self.assertLess(
            growth, GPU_LEAK_TOLERANCE_MB,
            f"GPU memory grew by {growth:.0f}MB over 30s — possible leak"
        )

    def test_memory_frees_after_context_release(self):
        """Test that pycuda context releases memory properly."""
        _skip_if_no_gpu()
        try:
            import pycuda.driver as drv
            drv.init()

            ctx = drv.Device(0).make_context()
            free_in, total_in = drv.mem_get_info()

            # Allocate 100MB
            alloc = drv.mem_alloc(100 * 1024 * 1024)
            free_after_alloc, _ = drv.mem_get_info()
            growth = (free_in - free_after_alloc) / 1024 / 1024
            self.assertGreater(growth, 50,
                               "100MB allocation didn't register in GPU memory")

            # Free it
            del alloc
            free_after_free, _ = drv.mem_get_info()
            recovered = (free_after_free - free_after_alloc) / 1024 / 1024
            ctx.pop()

            self.assertGreater(recovered, 50,
                               f"Only {recovered:.0f}MB recovered after freeing 100MB alloc")
        except ImportError:
            self.skipTest("pycuda not available")


class TestModelBudgetDocumentation(unittest.TestCase):
    """Verify total model budget is within spec (pure math, no GPU needed)."""

    def test_total_model_budget_under_6gb(self):
        """Sum of all model budgets must be < 6000MB."""
        total = sum(MODEL_BUDGETS_MB.values())
        self.assertLess(
            total, GPU_STACK_BUDGET_MB,
            f"Model budgets sum to {total}MB, exceeds {GPU_STACK_BUDGET_MB}MB limit"
        )

    def test_individual_model_budgets_reasonable(self):
        """No single model should claim more than 3GB."""
        for model, budget in MODEL_BUDGETS_MB.items():
            self.assertLess(
                budget, 3000,
                f"Model {model} budget {budget}MB seems unreasonably high"
            )

    def test_budget_includes_all_models(self):
        """Verify all expected models are accounted for in the budget."""
        expected_keys = {
            "scrfd_10g_trt", "arcface_r100_trt", "ecapa_tdnn_trt",
            "whisper_small", "phi3_mini_gguf", "piper_tts", "tracking"
        }
        self.assertEqual(set(MODEL_BUDGETS_MB.keys()), expected_keys)

    def test_budget_summary(self):
        """Print human-readable budget summary."""
        print("\n[GPU Budget Summary]")
        for model, mb in sorted(MODEL_BUDGETS_MB.items(), key=lambda x: -x[1]):
            bar = "█" * (mb // 100)
            print(f"  {model:<25} {mb:>5} MB  {bar}")
        print(f"  {'TOTAL':<25} {TOTAL_MODEL_BUDGET_MB:>5} MB")
        print(f"  {'LIMIT':<25} {GPU_STACK_BUDGET_MB:>5} MB")
        print(f"  {'RESERVE':<25} {GPU_STACK_BUDGET_MB - TOTAL_MODEL_BUDGET_MB:>5} MB")
