from .performance_helpers import get_collector_endpoint

"""CPU and Memory profiling script for MCA SDK.

This script profiles the SDK under high-throughput load to identify
performance bottlenecks and memory hotspots.

Usage:
    # CPU profiling with cProfile
    python -m cProfile -o profile.stats tests/performance/profile_sdk.py

    # Memory profiling with memory_profiler
    python -m memory_profiler tests/performance/profile_sdk.py

    # Analyze cProfile output
    python -m pstats profile.stats
"""

import cProfile
import os
import pstats
import sys
import time
import tracemalloc
from pathlib import Path

# Import SDK from installed package (assumes 'pip install -e .' was run)
# Do not use sys.path manipulation - breaks when SDK is properly installed
from mca_sdk import MCAClient


def profile_sdk_under_load(iterations: int = 10000):
    """Profile SDK with multiple predictions to simulate high load.

    Args:
        iterations: Number of predictions to record (default: 10,000)
    """
    print(f"Starting SDK profiling with {iterations} predictions...")

    # Use environment variable or default endpoint
    # Allows override without code changes: export MCA_OTEL_ENDPOINT="http://custom:4318"
    collector_endpoint = os.environ.get("MCA_OTEL_ENDPOINT", "http://localhost:4318")
    print(f"Using collector endpoint: {collector_endpoint}")

    # NOTE: Use production default validation for realistic profiling
    client = MCAClient(
        service_name="profile-test",
        model_id="profile-model-001",
        team_name="performance-team",
        model_type="internal",
        collector_endpoint=collector_endpoint,
        strict_validation=True,  # Use production default
    )

    start_time = time.time()

    # Simulate high-throughput prediction recording
    for i in range(iterations):
        client.record_prediction(
            latency=0.1 + (i % 100) / 1000,  # Vary latency: 0.1 - 0.199s
            model_version="v1.0",
            input_size=1024 + (i % 512),  # Vary input size
            output_size=512 + (i % 256),  # Vary output size
            prediction_id=f"pred-{i:06d}",
        )

        # Add some variance to simulate real-world usage
        if i % 100 == 0:
            client.record_metric("model.accuracy", 0.95 + (i % 5) / 100)

        if i % 1000 == 0:
            print(f"Progress: {i}/{iterations} predictions recorded...")

    elapsed = time.time() - start_time
    throughput = iterations / elapsed

    print(f"\nCompleted {iterations} predictions in {elapsed:.2f}s")
    print(f"Throughput: {throughput:.0f} predictions/second")

    client.shutdown()
    print("SDK shutdown complete.")


def profile_memory_with_tracemalloc(iterations: int = 10000):
    """Profile memory usage using built-in tracemalloc.

    NOTE: tracemalloc only tracks Python object allocations, not total process
    memory (RSS) which includes Python VM, shared libraries, C extensions, etc.
    For complete memory analysis, use external tools like memory_profiler or psutil.

    Args:
        iterations: Number of predictions to record
    """
    print(f"\n=== Memory Profiling with tracemalloc ({iterations} iterations) ===\n")
    print("NOTE: tracemalloc tracks Python allocations only, not total process memory (RSS)")

    tracemalloc.start()

    # Baseline memory
    baseline_mem = tracemalloc.get_traced_memory()
    print(f"Baseline memory (Python objects): {baseline_mem[0] / 1024 / 1024:.2f} MB")

    # Use environment variable or default endpoint
    collector_endpoint = os.environ.get("MCA_OTEL_ENDPOINT", "http://localhost:4318")

    # Run SDK operations with production defaults
    client = MCAClient(
        service_name="memory-test",
        model_id="memory-model-001",
        team_name="performance-team",
        collector_endpoint=collector_endpoint,
        strict_validation=True,  # Use production default
    )

    for i in range(iterations):
        client.record_prediction(
            latency=0.1, model_version="v1.0", input_size=1024, output_size=512
        )

        # Check memory every 1000 iterations
        if i > 0 and i % 1000 == 0:
            current, peak = tracemalloc.get_traced_memory()
            print(
                f"After {i} predictions - Current: {current / 1024 / 1024:.2f} MB, "
                f"Peak: {peak / 1024 / 1024:.2f} MB"
            )

    client.shutdown()

    # Final memory snapshot
    current, peak = tracemalloc.get_traced_memory()
    print(
        f"\nFinal memory - Current: {current / 1024 / 1024:.2f} MB, "
        f"Peak: {peak / 1024 / 1024:.2f} MB"
    )

    # Get top memory allocations
    snapshot = tracemalloc.take_snapshot()
    top_stats = snapshot.statistics("lineno")

    print("\n[ Top 10 Memory Allocations ]")
    for stat in top_stats[:10]:
        print(stat)

    tracemalloc.stop()


def main():
    """Run profiling suite."""
    print("=" * 80)
    print("MCA SDK Performance Profiling Suite")
    print("=" * 80)

    # Check if running under cProfile
    if any("cProfile" in arg for arg in sys.argv):
        print("\nRunning under cProfile...")
        profile_sdk_under_load()
    else:
        print("\nRunning standalone with built-in profiling...")

        # CPU profiling
        profiler = cProfile.Profile()
        profiler.enable()
        profile_sdk_under_load()
        profiler.disable()

        # Save and display CPU stats
        profiler.dump_stats("sdk_profile.stats")
        print("\n" + "=" * 80)
        print("TOP 20 FUNCTIONS BY CUMULATIVE TIME")
        print("=" * 80)
        stats = pstats.Stats("sdk_profile.stats")
        stats.strip_dirs()
        stats.sort_stats("cumulative")
        stats.print_stats(20)

        print("\n" + "=" * 80)
        print("TOP 20 FUNCTIONS BY TIME PER CALL (TOTTIME)")
        print("=" * 80)
        stats.sort_stats("tottime")
        stats.print_stats(20)

        # Memory profiling
        profile_memory_with_tracemalloc()

    print("\n" + "=" * 80)
    print("Profiling complete!")
    print("=" * 80)


if __name__ == "__main__":
    main()
