import time
import psutil
import statistics


def benchmark_model(model_callable, input_data, runs=100, warmup=10):
    process = psutil.Process()

    # Warmup
    for _ in range(warmup):
        model_callable(input_data)

    latencies = []

    mem_before = process.memory_info().rss / (1024 ** 2)

    start_total = time.perf_counter()

    for _ in range(runs):
        start = time.perf_counter()
        model_callable(input_data)
        end = time.perf_counter()
        latencies.append((end - start) * 1000)

    end_total = time.perf_counter()

    mem_after = process.memory_info().rss / (1024 ** 2)

    total_time = end_total - start_total

    return {
        "avg_latency_ms": round(statistics.mean(latencies), 4),
        "max_latency_ms": round(max(latencies), 4),
        "min_latency_ms": round(min(latencies), 4),
        "memory_usage_mb": round(mem_after - mem_before, 4),
        "throughput_req_per_sec": round(runs / total_time, 2),
    }