"""Performance benchmark for SDK-567: JSON parsing approaches.

Compares three approaches for parsing API responses:
1. Current: orjson.loads → dict → adapter transform → msgspec.Struct
2. Msgspec direct: msgspec.json.decode → raw struct → transform → final struct
3. Msgspec optimized: msgspec.json.decode with custom decode hooks (if API matches)

This benchmark uses realistic API response payloads to assess real-world impact.
"""

import sys
import time

from typing import Any

import msgspec
import orjson
import pytest

from pinecone._internal.adapters.inference_adapter import InferenceAdapter
from pinecone.models.inference import (
    Embedding,
    EmbedResponse,
    RerankResponse,
    RerankResult,
    Usage,
)


pytestmark = pytest.mark.unit


# Raw response structs that match API format exactly
class _RawEmbeddingItem(msgspec.Struct, kw_only=True):
    """Raw embedding item from API."""

    values: list[float]


class _RawEmbedResponse(msgspec.Struct, kw_only=True):
    """Raw embed response matching API format exactly."""

    data: list[_RawEmbeddingItem]
    model: str
    vector_type: str
    usage: dict[str, int]


class _RawRerankItem(msgspec.Struct, kw_only=True):
    """Raw rerank item from API."""

    index: int
    score: float
    document: str | None = None


class _RawRerankResponse(msgspec.Struct, kw_only=True):
    """Raw rerank response matching API format exactly."""

    data: list[_RawRerankItem]
    model: str
    usage: dict[str, int]


def create_realistic_embed_response_json(num_embeddings: int = 100, dimensions: int = 768) -> bytes:
    """Create realistic JSON response for embed API."""
    response = {
        "data": [
            {"values": [float(i + j * 0.01) for j in range(dimensions)]}
            for i in range(num_embeddings)
        ],
        "model": "multilingual-e5-large",
        "vector_type": "dense",
        "usage": {"total_tokens": num_embeddings * 50},
    }
    return orjson.dumps(response)


def create_realistic_rerank_response_json(num_results: int = 100) -> bytes:
    """Create realistic JSON response for rerank API."""
    response = {
        "data": [
            {
                "index": i,
                "score": 0.95 - i * 0.001,
                "document": f"This is document number {i} with some realistic text content that would appear in a reranking response.",
            }
            for i in range(num_results)
        ],
        "model": "bge-reranker-v2-m3",
        "usage": {"rerank_units": num_results * 10},
    }
    return orjson.dumps(response)


# Approach 1: Current implementation (orjson → dict → adapter)
def parse_embed_current(json_bytes: bytes) -> EmbedResponse:
    """Current approach: orjson.loads → dict → adapter."""
    data: dict[str, Any] = orjson.loads(json_bytes)
    adapter = InferenceAdapter()
    return adapter.to_embed_response(data)


# Approach 2: msgspec direct decode → transform
def parse_embed_msgspec_raw(json_bytes: bytes) -> EmbedResponse:
    """Msgspec approach: decode to raw struct → transform."""
    raw = msgspec.json.decode(json_bytes, type=_RawEmbedResponse)

    # Transform raw response to final model (same logic as adapter)
    embeddings = [Embedding(values=item.values, index=i) for i, item in enumerate(raw.data)]

    return EmbedResponse(
        data=embeddings,
        model=raw.model,
        vector_type=raw.vector_type,
        usage=Usage(total_tokens=raw.usage["total_tokens"]),
    )


# Approach 1: Current implementation for rerank
def parse_rerank_current(json_bytes: bytes) -> RerankResponse:
    """Current approach: orjson.loads → dict → adapter."""
    data: dict[str, Any] = orjson.loads(json_bytes)
    adapter = InferenceAdapter()
    return adapter.to_rerank_response(data)


# Approach 2: msgspec direct decode → transform for rerank
def parse_rerank_msgspec_raw(json_bytes: bytes) -> RerankResponse:
    """Msgspec approach: decode to raw struct → transform."""
    raw = msgspec.json.decode(json_bytes, type=_RawRerankResponse)

    # Transform raw response to final model (same logic as adapter)
    results = [
        RerankResult(
            index=item.index,
            score=item.score,
            document=item.document,
        )
        for item in raw.data
    ]

    # Handle rerank_units vs total_tokens
    usage_dict = raw.usage
    total_tokens = usage_dict.get("total_tokens") or usage_dict.get("rerank_units", 0)

    return RerankResponse(
        data=results,
        model=raw.model,
        usage=Usage(total_tokens=total_tokens),
    )


def benchmark_parsing_approach(
    name: str,
    json_bytes: bytes,
    parse_func: Any,
    iterations: int = 1000,
) -> float:
    """Benchmark a parsing approach."""
    # Warmup
    for _ in range(10):
        _ = parse_func(json_bytes)

    # Actual benchmark
    start = time.perf_counter()
    for _ in range(iterations):
        _ = parse_func(json_bytes)
    return time.perf_counter() - start


def test_embed_response_parsing_performance() -> None:
    """Benchmark embed response parsing approaches."""
    # Test with realistic payload size (100 embeddings x 768 dimensions)
    json_bytes = create_realistic_embed_response_json(num_embeddings=100, dimensions=768)
    json_size_kb = len(json_bytes) / 1024
    iterations = 1000

    print(f"\n{'=' * 80}")
    print("SDK-567: Embed Response JSON Parsing Benchmark")
    print(f"{'=' * 80}")
    print(f"Payload: 100 embeddings x 768 dimensions ({json_size_kb:.1f} KB JSON)")
    print(f"Iterations: {iterations}")
    print()

    # Benchmark current approach
    current_time = benchmark_parsing_approach(
        "Current (orjson → dict → adapter)",
        json_bytes,
        parse_embed_current,
        iterations,
    )

    # Benchmark msgspec approach
    msgspec_time = benchmark_parsing_approach(
        "Proposed (msgspec → raw struct → transform)",
        json_bytes,
        parse_embed_msgspec_raw,
        iterations,
    )

    # Calculate speedup/slowdown
    ratio = current_time / msgspec_time

    print(
        f"Current approach (orjson):     {current_time:.6f}s ({current_time / iterations * 1000:.3f}ms per call)"
    )
    print(
        f"Proposed approach (msgspec):   {msgspec_time:.6f}s ({msgspec_time / iterations * 1000:.3f}ms per call)"
    )
    print()

    if ratio > 1:
        print(f"Result: Proposed approach is {ratio:.2f}x FASTER ✓")
    elif ratio < 1:
        print(f"Result: Proposed approach is {1 / ratio:.2f}x SLOWER ✗")
    else:
        print("Result: No significant difference")

    print(f"{'=' * 80}\n")

    # Verify correctness
    result_current = parse_embed_current(json_bytes)
    result_msgspec = parse_embed_msgspec_raw(json_bytes)

    assert len(result_current.data) == len(result_msgspec.data)
    assert result_current.model == result_msgspec.model
    assert result_current.usage.total_tokens == result_msgspec.usage.total_tokens

    # Performance assertion - only fail if significantly slower (>20%)
    if ratio < 0.8:
        print(f"WARNING: Proposed approach is significantly slower ({1 / ratio:.2f}x)")


def test_rerank_response_parsing_performance() -> None:
    """Benchmark rerank response parsing approaches."""
    # Test with realistic payload size (100 rerank results with documents)
    json_bytes = create_realistic_rerank_response_json(num_results=100)
    json_size_kb = len(json_bytes) / 1024
    iterations = 1000

    print(f"\n{'=' * 80}")
    print("SDK-567: Rerank Response JSON Parsing Benchmark")
    print(f"{'=' * 80}")
    print(f"Payload: 100 rerank results with documents ({json_size_kb:.1f} KB JSON)")
    print(f"Iterations: {iterations}")
    print()

    # Benchmark current approach
    current_time = benchmark_parsing_approach(
        "Current (orjson → dict → adapter)",
        json_bytes,
        parse_rerank_current,
        iterations,
    )

    # Benchmark msgspec approach
    msgspec_time = benchmark_parsing_approach(
        "Proposed (msgspec → raw struct → transform)",
        json_bytes,
        parse_rerank_msgspec_raw,
        iterations,
    )

    # Calculate speedup/slowdown
    ratio = current_time / msgspec_time

    print(
        f"Current approach (orjson):     {current_time:.6f}s ({current_time / iterations * 1000:.3f}ms per call)"
    )
    print(
        f"Proposed approach (msgspec):   {msgspec_time:.6f}s ({msgspec_time / iterations * 1000:.3f}ms per call)"
    )
    print()

    if ratio > 1:
        print(f"Result: Proposed approach is {ratio:.2f}x FASTER ✓")
    elif ratio < 1:
        print(f"Result: Proposed approach is {1 / ratio:.2f}x SLOWER ✗")
    else:
        print("Result: No significant difference")

    print(f"{'=' * 80}\n")

    # Verify correctness
    result_current = parse_rerank_current(json_bytes)
    result_msgspec = parse_rerank_msgspec_raw(json_bytes)

    assert len(result_current.data) == len(result_msgspec.data)
    assert result_current.model == result_msgspec.model
    assert result_current.usage.total_tokens == result_msgspec.usage.total_tokens

    # Performance assertion - only fail if significantly slower (>20%)
    if ratio < 0.8:
        print(f"WARNING: Proposed approach is significantly slower ({1 / ratio:.2f}x)")


def test_varying_payload_sizes() -> None:
    """Benchmark different payload sizes to assess scaling behavior."""
    print(f"\n{'=' * 80}")
    print("SDK-567: Scaling Analysis - Varying Payload Sizes")
    print(f"{'=' * 80}")
    print()

    test_sizes = [
        (10, 384, "Small"),  # 10 embeddings x 384 dims
        (50, 768, "Medium"),  # 50 embeddings x 768 dims
        (200, 1024, "Large"),  # 200 embeddings x 1024 dims
    ]

    print(
        f"{'Size':<10} {'Embeddings':<12} {'Dims':<6} {'JSON KB':<10} {'Current':<12} {'Msgspec':<12} {'Speedup'}"
    )
    print("-" * 80)

    for num_embeddings, dimensions, size_label in test_sizes:
        json_bytes = create_realistic_embed_response_json(num_embeddings, dimensions)
        json_size_kb = len(json_bytes) / 1024

        current_time = benchmark_parsing_approach(
            "current", json_bytes, parse_embed_current, iterations=500
        )
        msgspec_time = benchmark_parsing_approach(
            "msgspec", json_bytes, parse_embed_msgspec_raw, iterations=500
        )

        ratio = current_time / msgspec_time
        speedup_str = f"{ratio:.2f}x"

        print(
            f"{size_label:<10} {num_embeddings:<12} {dimensions:<6} "
            f"{json_size_kb:<10.1f} {current_time:<12.6f} {msgspec_time:<12.6f} {speedup_str}"
        )

    print(f"{'=' * 80}\n")


def test_memory_characteristics() -> None:
    """Analyze memory allocation patterns of both approaches."""
    print(f"\n{'=' * 80}")
    print("SDK-567: Memory Allocation Analysis")
    print(f"{'=' * 80}")
    print()

    json_bytes = create_realistic_embed_response_json(num_embeddings=100, dimensions=768)

    # Current approach creates intermediate dict
    result_current = parse_embed_current(json_bytes)
    current_size = sys.getsizeof(result_current)

    # Msgspec approach avoids intermediate dict
    result_msgspec = parse_embed_msgspec_raw(json_bytes)
    msgspec_size = sys.getsizeof(result_msgspec)

    print(f"Final object size (current):  {current_size} bytes")
    print(f"Final object size (msgspec):  {msgspec_size} bytes")
    print()
    print("Note: Current approach allocates intermediate dict + final struct")
    print("      Msgspec approach allocates raw struct + final struct")
    print("      Both approaches still require transform step due to API schema differences")
    print(f"{'=' * 80}\n")
