from .performance_helpers import get_collector_endpoint

"""Memory leak validation for MCA SDK.

This module validates that the SDK does not have memory leaks over extended
operation. Uses accelerated aging tests with forced garbage collection.

Acceptance Criteria (AC 2):
- Memory usage remains flat over long-duration operation
- Total memory growth <10MB over 24-hour equivalent after warmup
- No unbounded growth in Span, Record, or Queue items
"""

import gc
import logging
import statistics
import sys
import time
import tracemalloc
from collections import defaultdict
from typing import Dict, List, Tuple

import objgraph
import pytest

from mca_sdk import MCAClient

# Mark as slow/memory test
pytestmark = [pytest.mark.slow, pytest.mark.memory]

logger = logging.getLogger(__name__)


def get_memory_usage_mb() -> float:
    """Get current process memory usage in MB.

    Returns:
        Memory usage in megabytes
    """
    tracemalloc.start()
    snapshot = tracemalloc.take_snapshot()
    total_bytes = sum(stat.size for stat in snapshot.statistics("lineno"))
    tracemalloc.stop()
    return total_bytes / 1024 / 1024


def analyze_memory_growth(snapshots: List[Tuple[int, int]]) -> Dict:
    """Analyze memory growth pattern from snapshots using linear regression.

    Args:
        snapshots: List of (iteration, memory_bytes) tuples

    Returns:
        Dict with growth analysis including regression statistics
    """
    if len(snapshots) < 2:
        return {
            "growth_mb": 0,
            "growth_per_iter_bytes": 0,
            "linear_growth": False,
            "slope": 0,
            "r_squared": 0,
        }

    iterations = [s[0] for s in snapshots]
    memory_bytes = [s[1] for s in snapshots]

    # Calculate total growth
    baseline_mb = memory_bytes[0] / 1024 / 1024
    final_mb = memory_bytes[-1] / 1024 / 1024
    total_growth_mb = final_mb - baseline_mb

    # Growth per iteration
    total_iterations = iterations[-1] - iterations[0]
    growth_per_iter = (
        (memory_bytes[-1] - memory_bytes[0]) / total_iterations if total_iterations > 0 else 0
    )

    # Perform linear regression to detect true linear growth
    # y = mx + b, where y = memory, x = iteration
    if len(snapshots) >= 10:
        n = len(snapshots)
        x = iterations
        y = memory_bytes

        # Calculate means
        x_mean = statistics.mean(x)
        y_mean = statistics.mean(y)

        # Calculate slope (m) and intercept (b)
        numerator = sum((x[i] - x_mean) * (y[i] - y_mean) for i in range(n))
        denominator = sum((x[i] - x_mean) ** 2 for i in range(n))
        slope = numerator / denominator if denominator != 0 else 0
        intercept = y_mean - slope * x_mean

        # Calculate R² (coefficient of determination)
        y_pred = [slope * x[i] + intercept for i in range(n)]
        ss_res = sum((y[i] - y_pred[i]) ** 2 for i in range(n))
        ss_tot = sum((y[i] - y_mean) ** 2 for i in range(n))
        r_squared = 1 - (ss_res / ss_tot) if ss_tot != 0 else 0

        # Detect linear growth:
        # - Positive slope (memory increasing)
        # - High R² (>0.8) indicating strong linear correlation
        # - Significant growth rate (>10 bytes per iteration or >1% total growth)
        min_growth_threshold = 10  # bytes per iteration
        min_relative_growth = 0.01  # 1% growth
        relative_growth = total_growth_mb / baseline_mb if baseline_mb > 0 else 0

        linear_growth = (
            slope > min_growth_threshold
            and r_squared > 0.8
            and relative_growth > min_relative_growth
        )
    else:
        slope = 0
        r_squared = 0
        linear_growth = False

    return {
        "baseline_mb": baseline_mb,
        "final_mb": final_mb,
        "growth_mb": total_growth_mb,
        "growth_per_iter_bytes": growth_per_iter,
        "linear_growth": linear_growth,
        "slope": slope,
        "r_squared": r_squared,
        "total_iterations": total_iterations,
    }


def check_object_leaks(baseline_snapshot, final_snapshot, top_n: int = 10) -> Dict:
    """Check for object leaks by comparing snapshots.

    Args:
        baseline_snapshot: Initial tracemalloc snapshot
        final_snapshot: Final tracemalloc snapshot
        top_n: Number of top allocations to analyze

    Returns:
        Dict with leak analysis
    """
    # Compare snapshots
    top_stats = final_snapshot.compare_to(baseline_snapshot, "lineno")

    # Identify potential leaks
    leaks = []
    for stat in top_stats[:top_n]:
        # Look for significant growth
        if stat.size_diff > 1024 * 1024:  # >1MB growth
            leaks.append(
                {
                    "file": stat.traceback.format()[0] if stat.traceback else "unknown",
                    "size_diff_mb": stat.size_diff / 1024 / 1024,
                    "count_diff": stat.count_diff,
                }
            )

    return {
        "potential_leaks": leaks,
        "leak_detected": len(leaks) > 0,
        "top_allocations": [
            {
                "file": stat.traceback.format()[0] if stat.traceback else "unknown",
                "size_mb": stat.size / 1024 / 1024,
                "count": stat.count,
            }
            for stat in top_stats[:top_n]
        ],
    }


def run_accelerated_aging_test(
    iterations: int = 1000,
    cycles: int = 100,
    force_gc_every: int = 100,
) -> Dict:
    """Run accelerated aging test with forced garbage collection.

    Simulates long-duration operation by running many short cycles with
    forced garbage collection to expose leaks quickly.

    Args:
        iterations: Predictions per cycle
        cycles: Number of cycles to run
        force_gc_every: Force GC every N iterations

    Returns:
        Dict containing test results
    """
    logger.info("Starting accelerated aging test:")
    logger.info(f"  Iterations per cycle: {iterations}")
    logger.info(f"  Cycles: {cycles}")
    logger.info(f"  Total predictions: {iterations * cycles:,}")

    # Start memory tracking
    tracemalloc.start()

    # Initialize client
    client = MCAClient(
        service_name="memory_aging_test",
        model_id="aging_test_model",
        team_name="performance_team",
        model_category="internal",
        model_type="regression",
        collector_endpoint=get_collector_endpoint(),
        buffering_enabled=True,
    )

    # Warmup phase (not counted in measurements)
    logger.info("Warmup phase (10 iterations)...")
    for _ in range(10):
        client.record_prediction(latency=0.001)
    gc.collect()
    time.sleep(0.5)

    # Baseline snapshot
    baseline_snapshot = tracemalloc.take_snapshot()
    baseline_memory = sum(stat.size for stat in baseline_snapshot.statistics("lineno"))
    snapshots = [(0, baseline_memory)]
    logger.info(f"Baseline memory: {baseline_memory / 1024 / 1024:.2f} MB")

    # Run cycles
    start_time = time.time()
    total_iterations = 0

    for cycle in range(1, cycles + 1):
        cycle_start = time.time()

        # Run predictions
        for i in range(iterations):
            client.record_prediction(
                latency=0.001,
                model_id=f"model_{i % 10}",  # Vary attributes
                model_version=f"v{i % 5}",
            )
            total_iterations += 1

            # Force GC periodically
            if total_iterations % force_gc_every == 0:
                gc.collect()

        # Take snapshot at end of cycle
        cycle_snapshot = tracemalloc.take_snapshot()
        cycle_memory = sum(stat.size for stat in cycle_snapshot.statistics("lineno"))
        snapshots.append((total_iterations, cycle_memory))

        cycle_elapsed = time.time() - cycle_start
        memory_mb = cycle_memory / 1024 / 1024
        growth_mb = (cycle_memory - baseline_memory) / 1024 / 1024

        # Report progress
        if cycle % 10 == 0 or cycle == 1:
            logger.info(
                f"Cycle {cycle}/{cycles}: "
                f"Memory: {memory_mb:.2f} MB "
                f"(+{growth_mb:.2f} MB) "
                f"| {cycle_elapsed:.2f}s"
            )

    elapsed_time = time.time() - start_time

    # Final snapshot and GC
    gc.collect()
    time.sleep(0.5)
    final_snapshot = tracemalloc.take_snapshot()
    final_memory = sum(stat.size for stat in final_snapshot.statistics("lineno"))
    snapshots.append((total_iterations, final_memory))

    # Shutdown client
    client.shutdown()

    # Analyze growth
    growth_analysis = analyze_memory_growth(snapshots)

    # Check for object leaks
    leak_analysis = check_object_leaks(baseline_snapshot, final_snapshot)

    # Results
    results = {
        "total_iterations": total_iterations,
        "cycles": cycles,
        "elapsed_seconds": elapsed_time,
        "memory": growth_analysis,
        "leaks": leak_analysis,
        "snapshots": snapshots,
        "acceptance_criteria": {
            "growth_under_10mb": growth_analysis["growth_mb"] < 10,
            "no_linear_growth": not growth_analysis["linear_growth"],
            "no_object_leaks": not leak_analysis["leak_detected"],
        },
    }

    # Print results
    logger.info("\n" + "=" * 60)
    logger.info("ACCELERATED AGING TEST RESULTS")
    logger.info("=" * 60)
    logger.info(f"Total Predictions: {total_iterations:,}")
    logger.info(f"Cycles: {cycles}")
    logger.info(f"Duration: {elapsed_time:.1f}s")
    logger.info("")
    logger.info("Memory Analysis:")
    logger.info(f"  Baseline:   {growth_analysis['baseline_mb']:.2f} MB")
    logger.info(f"  Final:      {growth_analysis['final_mb']:.2f} MB")
    logger.info(f"  Growth:     {growth_analysis['growth_mb']:.2f} MB")
    logger.info(f"  Per-iter:   {growth_analysis['growth_per_iter_bytes']:.2f} bytes")
    logger.info(
        f"  Linear:     {'YES (LEAK!)' if growth_analysis['linear_growth'] else 'NO (stable)'}"
    )
    logger.info("")
    logger.info("Object Leak Analysis:")
    logger.info(f"  Leaks found: {leak_analysis['leak_detected']}")
    if leak_analysis["potential_leaks"]:
        logger.info("  Potential leak locations:")
        for leak in leak_analysis["potential_leaks"]:
            logger.info(
                f"    - {leak['file']}: +{leak['size_diff_mb']:.2f} MB ({leak['count_diff']} objects)"
            )
    logger.info("")
    logger.info("Acceptance Criteria:")
    logger.info(
        f"  Growth < 10MB:     {'PASS' if results['acceptance_criteria']['growth_under_10mb'] else 'FAIL'} "
        f"({growth_analysis['growth_mb']:.2f} MB)"
    )
    logger.info(
        f"  No linear growth:  {'PASS' if results['acceptance_criteria']['no_linear_growth'] else 'FAIL'}"
    )
    logger.info(
        f"  No object leaks:   {'PASS' if results['acceptance_criteria']['no_object_leaks'] else 'FAIL'}"
    )
    logger.info("=" * 60)

    return results


def run_object_growth_check(iterations: int = 10000) -> Dict:
    """Check for object growth in specific SDK types.

    Focuses on: Span, Record, and Queue items.

    Args:
        iterations: Number of predictions to make

    Returns:
        Dict with object count analysis
    """
    logger.info(f"Running object growth check ({iterations:,} iterations)...")

    # Track object counts
    object_types = [
        "Span",
        "Record",
        "LogRecord",
        "dict",
        "deque",
        "list",
    ]

    # Baseline counts
    gc.collect()
    baseline_counts = {obj_type: len(objgraph.by_type(obj_type)) for obj_type in object_types}

    # Initialize client
    client = MCAClient(
        service_name="object_growth_test",
        model_id="object_test_model",
        team_name="performance_team",
        model_category="internal",
        model_type="regression",
        collector_endpoint=get_collector_endpoint(),
        buffering_enabled=True,
    )

    # Run predictions
    for i in range(iterations):
        client.record_prediction(latency=0.001)

    # Force GC and measure
    gc.collect()
    time.sleep(0.5)

    final_counts = {obj_type: len(objgraph.by_type(obj_type)) for obj_type in object_types}

    # Shutdown
    client.shutdown()

    # Calculate growth
    growth = {
        obj_type: final_counts[obj_type] - baseline_counts[obj_type] for obj_type in object_types
    }

    # Results
    results = {
        "iterations": iterations,
        "baseline": baseline_counts,
        "final": final_counts,
        "growth": growth,
        "growth_per_iter": {obj_type: growth[obj_type] / iterations for obj_type in object_types},
    }

    logger.info("\nObject Growth Analysis:")
    for obj_type in object_types:
        logger.info(
            f"  {obj_type:12s}: {baseline_counts[obj_type]:6d} -> {final_counts[obj_type]:6d} "
            f"(+{growth[obj_type]:6d}, {results['growth_per_iter'][obj_type]:.4f} per iter)"
        )

    return results


@pytest.mark.slow
@pytest.mark.memory
def test_accelerated_aging_no_memory_leak():
    """Validate no memory leaks using accelerated aging test.

    Acceptance Criteria:
    - AC 2.1: Memory usage remains flat (no unbounded growth)
    - AC 2.2: Total memory growth <10MB over test duration after warmup
    """
    results = run_accelerated_aging_test(
        iterations=1000,
        cycles=100,
        force_gc_every=100,
    )

    # Assertions
    assert results["acceptance_criteria"][
        "growth_under_10mb"
    ], f"Memory growth exceeds 10MB: {results['memory']['growth_mb']:.2f} MB"

    assert results["acceptance_criteria"][
        "no_linear_growth"
    ], "Linear memory growth detected (unbounded leak pattern)"

    assert results["acceptance_criteria"][
        "no_object_leaks"
    ], f"Object leaks detected: {len(results['leaks']['potential_leaks'])} leak sources"


@pytest.mark.slow
@pytest.mark.memory
def test_specific_object_types_no_leak():
    """Validate specific SDK object types don't leak.

    Checks: Span, Record, Queue items.

    Acceptance Criteria:
    - AC 2.3: No unbounded growth in Span, Record, or Queue items
    """
    results = run_object_growth_check(iterations=10000)

    # Check critical object types
    critical_types = ["Span", "Record", "deque"]

    for obj_type in critical_types:
        growth_per_iter = results["growth_per_iter"].get(obj_type, 0)

        # Stricter threshold: <0.001 objects per iteration (1000 iterations = 1 object)
        # Or absolute limit: max 10 objects total growth over entire test
        max_growth_per_iter = 0.001
        max_absolute_growth = 10

        assert (
            growth_per_iter < max_growth_per_iter
            or results["growth"][obj_type] < max_absolute_growth
        ), (
            f"Unbounded growth detected in {obj_type}: "
            f"{growth_per_iter:.6f} objects per iteration "
            f"(total growth: {results['growth'][obj_type]}, threshold: {max_growth_per_iter} per iter or {max_absolute_growth} total)"
        )

    logger.info("Object leak check: PASS (all critical types stable)")


if __name__ == "__main__":
    # Configure logging
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

    # Run accelerated aging test
    results = run_accelerated_aging_test(
        iterations=1000,
        cycles=100,
        force_gc_every=100,
    )

    # Run object growth check
    object_results = run_object_growth_check(iterations=10000)

    # Exit with status
    all_pass = all(results["acceptance_criteria"].values())
    exit(0 if all_pass else 1)
