import argparse
import asyncio
import io
import json
import os
import pickle
import socket
import sys
import tempfile
import threading
import time

import numpy as np
import psutil

import tenso

# Optional dependencies for comparison
try:
    import msgpack
    import pyarrow as pa
    from safetensors.numpy import load as st_load
    from safetensors.numpy import save as st_save
except ImportError:
    print(
        "Warning: Missing benchmark dependencies (msgpack, pyarrow, safetensors). Some tests skipped."
    )


try:
    from tenso.async_core import awrite_stream

    HAS_ASYNC = True
except ImportError:
    HAS_ASYNC = False

# Global state for integrity check
USE_INTEGRITY = False

# --- HELPER: Resource Monitoring ---


class ResourceMonitor:
    """Monitor CPU and memory during benchmarks."""

    def __init__(self):
        """Initialize the resource monitor."""
        self.process = psutil.Process()
        self.baseline_memory = self.process.memory_info().rss / 1024 / 1024

    def snapshot(self):
        """Take a snapshot of current resource usage."""
        return {
            "memory_mb": self.process.memory_info().rss / 1024 / 1024,
            "cpu_percent": self.process.cpu_percent(),
        }

    def memory_delta(self):
        """Calculate memory usage delta from baseline."""
        current = self.process.memory_info().rss / 1024 / 1024
        return current - self.baseline_memory


# --- 1. SERIALIZATION BENCHMARK HELPERS ---


def bench_json(data):
    """
    Benchmark JSON serialization and deserialization.

    JSON is a text-based format that converts arrays to lists, which can be
    inefficient for large numerical arrays due to string conversion overhead.

    Parameters
    ----------
    data : array_like
        The input data to benchmark. Will be converted to list for JSON.

    Returns
    -------
    enc : callable
        Encoder function that serializes to JSON bytes.
    dec : callable
        Decoder function that deserializes from JSON bytes.
    """
    enc = lambda x: json.dumps(x.tolist()).encode("utf-8")  # noqa
    dec = lambda x: np.array(json.loads(x), dtype=data.dtype)  # noqa
    return enc, dec


def bench_pickle(data):
    """
    Benchmark pickle serialization and deserialization.

    Pickle is Python's native serialization format, efficient for Python objects
    but not cross-language compatible and potentially insecure for untrusted data.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function that serializes with highest protocol.
    dec : callable
        Decoder function that deserializes from pickle bytes.
    """
    enc = lambda x: pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)  # noqa
    dec = lambda x: pickle.loads(x)  # noqa
    return enc, dec


def bench_msgpack(data):
    """
    Benchmark msgpack serialization and deserialization.

    Msgpack is a binary serialization format that's more efficient than JSON
    for numerical data, but requires converting arrays to bytes.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function that serializes to msgpack bytes.
    dec : callable
        Decoder function that deserializes and reshapes from msgpack bytes.
    """
    enc = lambda x: msgpack.packb(x.tobytes())  # noqa
    dec = lambda x: np.frombuffer(msgpack.unpackb(x), dtype=data.dtype).reshape(
        data.shape
    )  # noqa
    return enc, dec


def bench_safetensors(data):
    """
    Benchmark safetensors serialization and deserialization.

    Safetensors is a fast, safe serialization format for tensors, designed
    for machine learning models with memory mapping support.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function that saves to safetensors format.
    dec : callable
        Decoder function that loads from safetensors format.
    """

    def enc(x):
        return st_save({"t": x})

    def dec(x):
        return st_load(x)["t"]

    return enc, dec


def bench_arrow(data):
    """
    Benchmark Apache Arrow serialization and deserialization.

    Apache Arrow provides columnar in-memory analytics with efficient
    serialization for big data processing.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function that serializes to Arrow IPC format.
    dec : callable
        Decoder function that deserializes from Arrow IPC format.
    """

    def enc(x):
        arr = pa.array(x.flatten())
        batch = pa.RecordBatch.from_arrays([arr], names=["t"])
        sink = pa.BufferOutputStream()
        with pa.ipc.new_stream(sink, batch.schema) as writer:
            writer.write_batch(batch)
        return sink.getvalue()

    def dec(x):
        with pa.ipc.open_stream(x) as reader:
            batch = reader.read_next_batch()
            return batch.column(0).to_numpy(zero_copy_only=True).reshape(data.shape)

    return enc, dec


def bench_tenso(data):
    """
    Benchmark Tenso serialization and deserialization.

    Tenso is a high-performance tensor serialization format optimized for
    numpy arrays with optional integrity checking and compression.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function that serializes to Tenso format.
    dec : callable
        Decoder function that deserializes from Tenso format.
    """

    # Wrap dumps to respect the global USE_INTEGRITY flag
    def enc(x):
        return tenso.dumps(x, check_integrity=USE_INTEGRITY)

    return enc, tenso.loads


def bench_tenso_vectored(data):
    """
    Benchmark Tenso vectored serialization.

    Measures the time to prepare chunks for zero-copy transmission.
    Vectored I/O allows sending multiple buffers without copying.

    Parameters
    ----------
    data : array_like
        The input data to benchmark.

    Returns
    -------
    enc : callable
        Encoder function.
    dec : callable
        Decoder function (placeholder, returns original data).
    """
    # This prepares the packet metadata but yields the original tensor.data memoryview
    enc = lambda x: list(tenso.iter_dumps(x, check_integrity=USE_INTEGRITY))  # noqa
    # Deserialization from chunks happens at the I/O layer, so this is a placeholder
    dec = lambda x: data  # noqa
    return enc, dec


# --- BENCHMARK RUNNERS ---


def run_serialization():
    """Run in-memory serialization benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 1: IN-MEMORY SERIALIZATION (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    SCENARIOS = [
        {"name": "API Vector", "shape": (1536,), "dtype": np.float32},
        {"name": "CV Batch", "shape": (32, 256, 256, 3), "dtype": np.uint8},
        {"name": "LLM Layer", "shape": (4096, 4096), "dtype": np.float32},
        {"name": "Bundle", "type": "dict", "size": 10},
    ]

    print(
        f"{'SCENARIO':<15} | {'FORMAT':<12} | {'SIZE':<10} | {'SERIALIZE':<10} | {'DESERIALIZE':<10}"
    )
    print("-" * 80)

    for scen in SCENARIOS:
        if scen.get("type") == "dict":
            # Create a bundle of small tensors
            data = {f"key_{i}": np.random.rand(100, 100).astype(np.float32) for i in range(scen["size"])}
            # Helper for size calculation since data is a dict
            class Dummy:
                pass
            dummy = Dummy()
            dummy.nbytes = sum(x.nbytes for x in data.values())
            # We can't use existing bench functions directly for dicts because they expect array-like in some places
            # But bench_tenso and bench_pickle work with dicts.
            # Define specific wrappers for this iteration if needed, or rely on duck typing.
            # Let's adjust the loop to handle dicts.
        elif scen["dtype"] == np.uint8:
            data = np.random.randint(0, 255, scen["shape"]).astype(np.uint8)
        else:
            data = np.random.rand(*scen["shape"]).astype(scen["dtype"])

        competitors = {
            "Pickle": bench_pickle(data),
            "Tenso": bench_tenso(data),
        }
        
        # Only add array-specific formats if not a dict
        if not isinstance(data, dict):
            competitors["Tenso (Vect)"] = bench_tenso_vectored(data)
            if "msgpack" in globals():
                competitors["MsgPack"] = bench_msgpack(data)
            if "st_save" in globals():
                competitors["Safetensors"] = bench_safetensors(data)
            if "pa" in globals():
                competitors["Arrow"] = bench_arrow(data)
        
        for name, (enc_func, dec_func) in competitors.items():
            try:
                # Warmup
                encoded = enc_func(data)
                _ = dec_func(encoded)

                ITERATIONS = 10
                t0 = time.perf_counter()
                for _ in range(ITERATIONS):
                    encoded = enc_func(data)
                t_ser = ((time.perf_counter() - t0) / ITERATIONS) * 1000

                t0 = time.perf_counter()
                for _ in range(ITERATIONS):
                    _ = dec_func(encoded)
                t_des = ((time.perf_counter() - t0) / ITERATIONS) * 1000

                # Handle size for vectored list output
                if isinstance(encoded, list):
                    total_bytes = sum(len(c) for c in encoded)
                else:
                    total_bytes = len(encoded)

                size_str = (
                    f"{total_bytes / 1024 / 1024:.2f} MB"
                    if total_bytes > 1024**2
                    else f"{total_bytes / 1024:.2f} KB"
                )
                print(
                    f"{scen['name']:<15} | {name:<12} | {size_str:<10} | {t_ser:>7.3f} ms | {t_des:>7.3f} ms"
                )
            except Exception as e:
                print(f"{scen['name']:<15} | {name:<12} | FAILED ({e})")
        print("-" * 80)


def run_io():
    """Run disk I/O benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 2: DISK I/O (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    shape = (8192, 8192)
    data = np.random.rand(*shape).astype(np.float32)
    size_mb = data.nbytes / (1024 * 1024)
    print(f"Dataset: {size_mb:.0f} MB Matrix {shape}")

    print(f"{'FORMAT':<15} | {'WRITE (ms)':<10} | {'READ (ms)':<10}")
    print("-" * 60)

    # 1. Tenso
    with tempfile.NamedTemporaryFile(delete=False) as f:
        path = f.name
    try:
        t0 = time.perf_counter()
        with open(path, "wb") as f:
            tenso.dump(data, f, check_integrity=USE_INTEGRITY)
        t_write = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        # load automatically detects integrity footer from header
        with open(path, "rb") as f:
            tenso.load(f, mmap_mode=True)
        t_read = (time.perf_counter() - t0) * 1000
        print(f"{'Tenso':<15} | {t_write:>10.2f} | {t_read:>10.2f}")
    finally:
        if os.path.exists(path):
            os.remove(path)

    # 2. Numpy
    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as f:
        path = f.name
    try:
        t0 = time.perf_counter()
        np.save(path, data)
        t_write = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        np.load(path, mmap_mode="r")
        t_read = (time.perf_counter() - t0) * 1000
        print(f"{'Numpy .npy':<15} | {t_write:>10.2f} | {t_read:>10.2f}")
    finally:
        if os.path.exists(path):
            os.remove(path)

    # 3. Pickle
    with tempfile.NamedTemporaryFile(delete=False) as f:
        path = f.name
    try:
        t0 = time.perf_counter()
        with open(path, "wb") as f:
            pickle.dump(data, f)
        t_write = (time.perf_counter() - t0) * 1000

        with open(path, "rb") as f:
            t0 = time.perf_counter()
            pickle.load(f)
            t_read = (time.perf_counter() - t0) * 1000
        print(f"{'Pickle':<15} | {t_write:>10.2f} | {t_read:>10.2f}")
    finally:
        if os.path.exists(path):
            os.remove(path)


def run_stream_read():
    """Run stream read benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 3: STREAM READ (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    class FastStream(io.BytesIO):
        pass  # IOBase supports readinto

    data = np.random.rand(5000, 5000).astype(np.float32)
    packet = tenso.dumps(data, check_integrity=USE_INTEGRITY)
    size_mb = len(packet) / (1024 * 1024)

    print(f"Dataset: {size_mb:.0f} MB Packet")
    print(f"{'METHOD':<20} | {'TIME (ms)':<10} | {'THROUGHPUT':<15}")
    print("-" * 60)

    # Optimized
    stream = FastStream(packet)
    t0 = time.perf_counter()
    tenso.read_stream(stream)
    t_opt = (time.perf_counter() - t0) * 1000
    print(
        f"{'Tenso read_stream':<20} | {t_opt:>7.2f} ms | {size_mb / (t_opt / 1000):>7.2f} MB/s"
    )

    # Legacy Loop Simulation
    stream.seek(0)
    t0 = time.perf_counter()
    chunks = []
    while True:
        chunk = stream.read(65536)
        if not chunk:
            break
        chunks.append(chunk)
    full_data = b"".join(chunks)  # Standard Python way
    tenso.loads(full_data)
    t_old = (time.perf_counter() - t0) * 1000
    print(
        f"{'Naive Loop':<20} | {t_old:>7.2f} ms | {size_mb / (t_old / 1000):>7.2f} MB/s"
    )

    print("-" * 60)
    print(f"Speedup: {t_old / t_opt:.1f}x")


def run_stream_write():
    """Run network write benchmarks with improved synchronization."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 4: NETWORK WRITE (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    # Use a dynamic port to avoid conflicts with other processes
    s_temp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s_temp.bind(("", 0))
    PORT = s_temp.getsockname()[1]
    s_temp.close()

    def sink_server(port):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        try:
            s.bind(("localhost", port))
            s.listen(1)
            conn, _ = s.accept()
            while True:
                chunk = conn.recv(1024 * 1024)
                if not chunk:
                    break
            conn.close()
        except Exception:
            pass
        finally:
            s.close()

    server_thread = threading.Thread(target=sink_server, args=(PORT,), daemon=True)
    server_thread.start()

    # Ensure the server has time to start listening
    time.sleep(0.5)

    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        client.connect(("localhost", PORT))

        data = np.random.rand(16, 16).astype(np.float32)
        COUNT = 10000

        print(f"Sending {COUNT} tensors (1KB each) over localhost TCP...")

        t0 = time.perf_counter()
        for _ in range(COUNT):
            tenso.write_stream(data, client, check_integrity=USE_INTEGRITY)
        t_total = time.perf_counter() - t0

        client.close()

        print(f"Total Time: {t_total:.4f}s")
        print(f"Throughput: {COUNT / t_total:.0f} packets/sec")
        print(f"Latency:    {(t_total / COUNT) * 1_000_000:.1f} µs/packet")
    except Exception as e:
        print(f"Network write benchmark failed: {e}")


def run_memory_overhead():
    """Run memory overhead benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 5: MEMORY OVERHEAD (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    shapes = [(1000, 1000), (2000, 2000), (4000, 4000)]

    print(
        f"{'SIZE':<15} | {'FORMAT':<12} | {'RAW (MB)':<10} | {'SERIALIZED (MB)':<15} | {'OVERHEAD':<10}"
    )
    print("-" * 80)

    for shape in shapes:
        data = np.random.rand(*shape).astype(np.float32)
        raw_size = data.nbytes / 1024 / 1024

        results = {}

        # Tenso
        packet = tenso.dumps(data, check_integrity=USE_INTEGRITY)
        results["Tenso"] = len(packet) / 1024 / 1024

        # Pickle
        packet = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL)
        results["Pickle"] = len(packet) / 1024 / 1024

        # Arrow
        if "pa" in globals():
            arr = pa.array(data.flatten())
            batch = pa.RecordBatch.from_arrays([arr], names=["t"])
            sink = pa.BufferOutputStream()
            with pa.ipc.new_stream(sink, batch.schema) as writer:
                writer.write_batch(batch)
            results["Arrow"] = len(sink.getvalue()) / 1024 / 1024

        # SafeTensors
        if "st_save" in globals():
            packet = st_save({"t": data})
            results["Safetensors"] = len(packet) / 1024 / 1024

        for name, size in results.items():
            overhead = ((size - raw_size) / raw_size) * 100
            print(
                f"{str(shape):<15} | {name:<12} | {raw_size:>8.2f} | {size:>13.2f} | {overhead:>8.2f}%"
            )

        print("-" * 80)


def run_cpu_usage():
    """Run CPU usage benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 6: CPU USAGE (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    data = np.random.rand(4096, 4096).astype(np.float32)
    monitor = ResourceMonitor()

    print(f"{'FORMAT':<15} | {'SERIALIZE CPU%':<15} | {'DESERIALIZE CPU%':<18}")
    print("-" * 60)

    formats = {
        "Tenso": (lambda x: tenso.dumps(x, check_integrity=USE_INTEGRITY), tenso.loads),
        "Pickle": (
            lambda x: pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL),
            pickle.loads,
        ),
    }

    if "pa" in globals():

        def arrow_enc(x):
            arr = pa.array(x.flatten())
            batch = pa.RecordBatch.from_arrays([arr], names=["t"])
            sink = pa.BufferOutputStream()
            with pa.ipc.new_stream(sink, batch.schema) as writer:
                writer.write_batch(batch)
            return sink.getvalue()

        def arrow_dec(x):
            with pa.ipc.open_stream(x) as reader:
                batch = reader.read_next_batch()
                return (
                    batch.column(0).to_numpy(zero_copy_only=False).reshape(data.shape)
                )

        formats["Arrow"] = (arrow_enc, arrow_dec)

    if "st_save" in globals():
        formats["Safetensors"] = (
            lambda x: st_save({"t": x}),
            lambda x: st_load(x)["t"],
        )

    for name, (enc, dec) in formats.items():
        # Warmup
        packet = enc(data)
        _ = dec(packet)

        # Measure serialization CPU
        monitor.process.cpu_percent()  # Reset
        time.sleep(0.1)
        for _ in range(20):
            _ = enc(data)
        ser_cpu = monitor.process.cpu_percent()

        # Measure deserialization CPU
        monitor.process.cpu_percent()  # Reset
        time.sleep(0.1)
        for _ in range(20):
            _ = dec(packet)
        des_cpu = monitor.process.cpu_percent()

        print(f"{name:<15} | {ser_cpu:>13.1f}% | {des_cpu:>16.1f}%")


async def _bench_async_throughput(data, count=1000):
    """Measure async write/read throughput."""

    # 1. Async Write Benchmark
    # We write to a null sink to measure purely the serialization + overhead
    class NullWriter:
        def write(self, d):
            pass

        async def drain(self):
            pass

    writer = NullWriter()
    t0 = time.perf_counter()
    for _ in range(count):
        await awrite_stream(data, writer, check_integrity=USE_INTEGRITY)
    t_write = time.perf_counter() - t0

    return t_write


def run_async_benchmark():
    """Run async I/O benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK: ASYNC I/O (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    if not HAS_ASYNC:
        print("Async core not available.")
        return

    shape = (128, 128)  # Small-ish tensor
    data = np.random.rand(*shape).astype(np.float32)
    count = 5000
    total_mb = (data.nbytes * count) / (1024**2)

    print(
        f"Streaming {count} tensors of shape {shape} ({data.nbytes / 1024:.1f} KB each)..."
    )

    # Run Async Write
    t_write = asyncio.run(_bench_async_throughput(data, count))

    fps = count / t_write
    mbps = total_mb / t_write

    print(
        f"{'Async Write':<20} | {t_write:>8.4f}s | {fps:>10.0f} tensors/s | {mbps:>8.2f} MB/s"
    )
    print("-" * 80)


def run_dtype_coverage():
    """Run dtype coverage benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 7: DTYPE COVERAGE (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    dtypes = [
        np.float16,
        np.float32,
        np.float64,
        np.int8,
        np.int16,
        np.int32,
        np.int64,
        np.uint8,
        np.uint16,
        np.uint32,
        np.uint64,
        np.bool_,
        np.complex64,
        np.complex128,
    ]

    # Try adding bfloat16
    try:
        from ml_dtypes import bfloat16

        dtypes.append(bfloat16)
    except ImportError:
        pass

    shape = (100, 100)
    print(
        f"{'DTYPE':<25} | {'STATUS':<10} | {'SIZE (KB)':<12} | {'ROUNDTRIP (ms)':<15}"
    )
    print("-" * 75)

    for dtype in dtypes:
        try:
            if dtype == np.bool_:
                data = np.random.randint(0, 2, shape).astype(dtype)
            elif hasattr(dtype, "name") and dtype.name == "bfloat16":
                data = np.random.randn(*shape).astype(np.float32).astype(dtype)
            elif np.issubdtype(dtype, np.integer):
                data = np.random.randint(0, 100, shape).astype(dtype)
            elif np.issubdtype(dtype, np.complexfloating):
                data = (np.random.randn(*shape) + 1j * np.random.randn(*shape)).astype(
                    dtype
                )
            else:
                data = np.random.randn(*shape).astype(dtype)

            t0 = time.perf_counter()
            packet = tenso.dumps(data, check_integrity=USE_INTEGRITY)
            restored = tenso.loads(packet)
            roundtrip = (time.perf_counter() - t0) * 1000

            status = "✓ PASS" if np.array_equal(data, restored) else "✗ FAIL"

            print(
                f"{str(dtype):<25} | {status:<10} | {len(packet) / 1024:>10.2f} | {roundtrip:>13.3f}"
            )
        except Exception as e:
            print(f"{str(dtype):<25} | {'✗ ERROR':<10} | {'-':<12} | {str(e)[:20]}")


def run_arrow_comparison():
    """Run Arrow vs Tenso comparison benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 8: ARROW vs TENSO (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    if "pa" not in globals():
        print("Arrow not available. Skipping.")
        return

    sizes = [
        ("Small", (512, 512)),
        ("Medium", (2048, 2048)),
        ("Large", (4096, 4096)),
        ("XLarge", (8192, 8192)),
    ]

    print(
        f"{'SIZE':<10} | {'TENSO SER':<12} | {'ARROW SER':<12} | {'TENSO DES':<12} | {'ARROW DES':<12} | {'SPEEDUP':<10}"
    )
    print("-" * 95)

    for name, shape in sizes:
        data = np.random.rand(*shape).astype(np.float32)

        # Tenso
        t0 = time.perf_counter()
        tenso_packet = tenso.dumps(data, check_integrity=USE_INTEGRITY)
        tenso_ser = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        _ = tenso.loads(tenso_packet)
        tenso_des = (time.perf_counter() - t0) * 1000

        # Arrow
        arr = pa.array(data.flatten())
        batch = pa.RecordBatch.from_arrays([arr], names=["t"])

        t0 = time.perf_counter()
        sink = pa.BufferOutputStream()
        with pa.ipc.new_stream(sink, batch.schema) as writer:
            writer.write_batch(batch)
        arrow_packet = sink.getvalue()
        arrow_ser = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        with pa.ipc.open_stream(arrow_packet) as reader:
            batch = reader.read_next_batch()
            _ = batch.column(0).to_numpy(zero_copy_only=False).reshape(shape)
        arrow_des = (time.perf_counter() - t0) * 1000

        speedup = arrow_des / tenso_des

        print(
            f"{name:<10} | {tenso_ser:>10.3f}ms | {arrow_ser:>10.3f}ms | {tenso_des:>10.3f}ms | {arrow_des:>10.3f}ms | {speedup:>8.1f}x"
        )

    print("-" * 95)


def run_correctness():
    """Run correctness benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK 9: CORRECTNESS (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    tests = [
        ("Zeros", lambda: np.zeros((100, 100), dtype=np.float32)),
        ("Ones", lambda: np.ones((100, 100), dtype=np.float32)),
        ("Random", lambda: np.random.rand(100, 100).astype(np.float32)),
        ("Negative", lambda: -np.random.rand(100, 100).astype(np.float32)),
        ("Mixed", lambda: np.random.randn(100, 100).astype(np.float32)),
        ("Large Values", lambda: np.random.rand(100, 100).astype(np.float32) * 1e6),
        ("Small Values", lambda: np.random.rand(100, 100).astype(np.float32) * 1e-6),
    ]

    print(f"{'TEST':<20} | {'STATUS':<10} | {'MAX ERROR':<15}")
    print("-" * 55)

    for name, gen_func in tests:
        data = gen_func()
        packet = tenso.dumps(data, check_integrity=USE_INTEGRITY)
        restored = tenso.loads(packet)

        if np.allclose(data, restored, rtol=1e-7, atol=1e-7):
            max_error = np.abs(data - restored).max()
            print(f"{name:<20} | {'✓ PASS':<10} | {max_error:<15.2e}")
        else:
            print(f"{name:<20} | {'✗ FAIL':<10} | {'MISMATCH':<15}")

    print("-" * 55)

def run_gpu_benchmark():
    print("\n" + "=" * 80)
    print("BENCHMARK: GPU TRANSFER")
    print("=" * 80)

    try:
        import torch

        import tenso.gpu as t_gpu

        # Setup: Create a large tensor in RAM
        shape = (4096, 4096)  # ~64MB float32
        data = np.random.rand(*shape).astype(np.float32)
        stream = io.BytesIO(tenso.dumps(data))

        # Benchmark
        t0 = time.perf_counter()
        # This uses your pinned memory logic
        t_gpu.read_to_device(stream, device_id=0)
        torch.cuda.synchronize()  # Wait for GPU to finish
        t_total = (time.perf_counter() - t0) * 1000

        print(f"CPU -> GPU (Tenso): {t_total:.2f} ms")

        # Compare to Standard PyTorch Load
        stream.seek(0)
        t0 = time.perf_counter()
        # Standard way: Load to CPU -> Move to GPU
        cpu_arr = tenso.loads(stream.read())
        torch.tensor(cpu_arr).to("cuda")
        torch.cuda.synchronize()
        t_std = (time.perf_counter() - t0) * 1000

        print(f"CPU -> GPU (Standard): {t_std:.2f} ms")
        print(f"Speedup: {t_std / t_total:.1f}x")

    except ImportError:
        print("Skipping GPU benchmark (Torch/CUDA not found)")


def run_gpu_benchmark():
    print("\n" + "=" * 80)
    print("BENCHMARK: GPU TRANSFER")
    print("=" * 80)

    try:
        import torch

        import tenso.gpu as t_gpu

        # Setup: Create a large tensor in RAM
        shape = (4096, 4096)  # ~64MB float32
        data = np.random.rand(*shape).astype(np.float32)
        stream = io.BytesIO(tenso.dumps(data))

        # Benchmark
        t0 = time.perf_counter()
        # This uses your pinned memory logic
        t_gpu.read_to_device(stream, device_id=0)
        torch.cuda.synchronize()  # Wait for GPU to finish
        t_total = (time.perf_counter() - t0) * 1000

        print(f"CPU -> GPU (Tenso): {t_total:.2f} ms")

        # Compare to Standard PyTorch Load
        stream.seek(0)
        t0 = time.perf_counter()
        # Standard way: Load to CPU -> Move to GPU
        cpu_arr = tenso.loads(stream.read())
        torch.tensor(cpu_arr).to("cuda")
        torch.cuda.synchronize()
        t_std = (time.perf_counter() - t0) * 1000

        print(f"CPU -> GPU (Standard): {t_std:.2f} ms")
        print(f"Speedup: {t_std / t_total:.1f}x")

    except ImportError:
        print("Skipping GPU benchmark (Torch/CUDA not found)")


def run_shm_benchmark():
    """Run Shared Memory IPC benchmarks."""
    print("\n" + "=" * 80)
    print(f"BENCHMARK: SHARED MEMORY IPC (Integrity: {USE_INTEGRITY})")
    print("=" * 80)

    try:
        from tenso import TensoShm
    except ImportError:
        print("TensoShm not available. Skipping.")
        return

    # Use a fairly large tensor to demonstrate throughput (e.g. 64MB)
    shape = (4096, 4096)
    data = np.random.rand(*shape).astype(np.float32)
    size_mb = data.nbytes / (1024 * 1024)
    shm_name = "tenso_bench_shm_01"

    print(f"Dataset: {size_mb:.0f} MB Matrix {shape}")
    print(f"{'OPERATION':<15} | {'TIME (ms)':<10} | {'THROUGHPUT':<15}")
    print("-" * 60)

    shm = None
    try:
        # Measure Write (Creation + Copy into SHM)
        t0 = time.perf_counter()
        shm = TensoShm.create_from(shm_name, data, check_integrity=USE_INTEGRITY)
        t_write = (time.perf_counter() - t0) * 1000
        
        print(f"{'Write (Put)':<15} | {t_write:>7.3f} ms | {size_mb / (t_write / 1000):>7.2f} MB/s")

        # Measure Read (Zero-copy mapping)
        # We use a separate instance to simulate a reader process, although in same process
        # the OS overhead for opening existing SHM is what we want to measure + mapping time.
        t0 = time.perf_counter()
        with TensoShm(shm_name) as reader:
            res = reader.get()
            del res
        t_read = (time.perf_counter() - t0) * 1000
        
        print(f"{'Read (Get)':<15} | {t_read:>7.3f} ms | {size_mb / (t_read / 1000):>7.2f} MB/s")
        
    finally:
        if shm:
            shm.close()
            shm.unlink()


def print_summary():
    """Print system info and summary."""
    print("\n" + "=" * 80)
    print("SYSTEM INFORMATION")
    print("=" * 80)
    print(f"Python Version: {sys.version.split()[0]}")
    print(f"NumPy Version:  {np.__version__}")
    print(
        f"Tenso Version:  {tenso.__version__ if hasattr(tenso, '__version__') else 'unknown'}"
    )
    print(f"CPU Cores:      {psutil.cpu_count()}")
    print(f"Total Memory:   {psutil.virtual_memory().total / 1024**3:.1f} GB")
    print(f"Platform:       {sys.platform}")

    if "pa" in globals():
        print(f"Arrow Version:  {pa.__version__}")
    print("=" * 80)


# --- MAIN ---

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Tenso Comprehensive Benchmarks",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Benchmark Modes:
  all         - Run all benchmarks (default)
  ser         - Serialization speed
  io          - Disk I/O performance
  read        - Stream reading
  write       - Network writing
  memory      - Memory overhead
  cpu         - CPU usage
  dtypes      - Dtype coverage
  arrow       - Arrow comparison
  correctness - Data integrity
  quick       - Quick overview (ser + arrow)
        """,
    )

    parser.add_argument(
        "mode",
        nargs="?",
        choices=[
            "all",
            "ser",
            "io",
            "read",
            "write",
            "memory",
            "cpu",
            "dtypes",
            "arrow",
            "correctness",
            "shm",
            "quick",
        ],
        default="all",
        help="Benchmark mode to run",
    )

    parser.add_argument(
        "--no-summary", action="store_true", help="Skip system information summary"
    )

    parser.add_argument(
        "--integrity",
        action="store_true",
        help="Enable XXH3 integrity checks for Tenso benchmarks",
    )

    args = parser.parse_args()

    # Set the global flag
    USE_INTEGRITY = args.integrity

    if not args.no_summary:
        print_summary()

    if args.mode == "all":
        run_serialization()
        run_io()
        run_stream_read()
        run_stream_write()
        run_memory_overhead()
        run_cpu_usage()
        run_dtype_coverage()
        run_arrow_comparison()
        run_correctness()
        run_async_benchmark()
        run_shm_benchmark()
    elif args.mode == "shm":
        run_shm_benchmark()
    elif args.mode == "async":
        run_async_benchmark()
    elif args.mode == "dtypes":
        run_dtype_coverage()
    elif args.mode == "quick":
        run_serialization()
        run_arrow_comparison()
    elif args.mode == "ser":
        run_serialization()
    elif args.mode == "io":
        run_io()
    elif args.mode == "read":
        run_stream_read()
    elif args.mode == "write":
        run_stream_write()
    elif args.mode == "memory":
        run_memory_overhead()
    elif args.mode == "cpu":
        run_cpu_usage()
    elif args.mode == "arrow":
        run_arrow_comparison()
    elif args.mode == "correctness":
        run_correctness()

    print("\n" + "=" * 80)
    print("BENCHMARKS COMPLETE")
    print("=" * 80)
