"""Concurrent throughput benchmark for core-product partition pipeline.

Mirrors production: FastAPI calls `asyncio.to_thread(partition, ...)` per request.
Uses pytest-async-benchmark for structured timing with rich output.

Usage (on benchmark VM):
    cd ~/core-product
    # Install from pedantic-mode branch
    uv pip install --python .venv/bin/python \
      'pytest-async-benchmark[asyncio] @ git+https://github.com/KRRT7/pytest-async-benchmark.git@feat/pedantic-mode'

    # Run all benchmarks (~5 min total)
    taskset -c 0 .venv/bin/python -m pytest bench/bench_throughput.py -v -s

    # Just concurrency scaling (fast strategy, ~20s)
    taskset -c 0 .venv/bin/python -m pytest bench/bench_throughput.py -v -s -k Concurrency

    # Just OCR pipeline (hi_res, ~5 min)
    taskset -c 0 .venv/bin/python -m pytest bench/bench_throughput.py -v -s -k OCR

Copy this file to ~/core-product/bench/ on the VM before running.
"""

from __future__ import annotations

import asyncio
from pathlib import Path

import pytest

REPO_ROOT = Path(__file__).resolve().parent.parent

# ---------------------------------------------------------------------------
# Fixtures -- chosen for fast iteration
#
# "fast" strategy fixtures (pdfminer, no OCR/detection, <0.5s each):
#   These test concurrency scaling without the 6s+ OCR overhead.
#
# "hi_res" fixture (full pipeline: render + detect + OCR + merge):
#   Single small doc to measure the hot path we're optimizing.
# ---------------------------------------------------------------------------
FAST_FIXTURES = [
    (
        "unstructured-api/sample-docs/embedded-images-tables.pdf",
        "img-tables-1p",
    ),
    ("unstructured_prop/tests/test_files/multi-column-2p.pdf", "multicol-2p"),
    (
        "unstructured-api/sample-docs/layout-parser-paper-with-table.pdf",
        "table-1p",
    ),
]

HIRES_FIXTURE = (
    "unstructured-api/sample-docs/embedded-images-tables.pdf",
    "img-tables-1p",
)


def resolve(rel: str) -> Path:
    p = REPO_ROOT / rel
    if not p.exists():
        pytest.skip(f"fixture not found: {p}")
    return p


# ---------------------------------------------------------------------------
# Partition wrappers
# ---------------------------------------------------------------------------
def partition_sync(filepath: Path, strategy: str = "auto") -> int:
    from unstructured.partition.auto import partition

    elements = partition(filename=str(filepath), strategy=strategy)
    return len(elements)


async def partition_async(filepath: Path, strategy: str = "auto") -> int:
    """Mirrors production: asyncio.to_thread(partition, ...)"""
    return await asyncio.to_thread(partition_sync, filepath, strategy)


async def batch_serial(filepaths: list[Path], strategy: str) -> list[int]:
    results = []
    for fp in filepaths:
        results.append(await partition_async(fp, strategy))
    return results


async def batch_concurrent(
    filepaths: list[Path], strategy: str, concurrency: int
) -> list[int]:
    sem = asyncio.Semaphore(concurrency)

    async def worker(fp: Path) -> int:
        async with sem:
            return await partition_async(fp, strategy)

    return list(await asyncio.gather(*[worker(fp) for fp in filepaths]))


# ===================================================================
# Concurrency scaling -- fast strategy (~0.2s/doc)
#
# Measures throughput scaling with concurrent requests.
# Uses fast strategy so each round completes in seconds, not minutes.
# The async patterns are identical to hi_res -- only the partition
# internals differ.
# ===================================================================
class TestConcurrencyScaling:
    """Concurrency scaling with fast strategy (practical iteration speed)."""

    @pytest.fixture
    def fast_paths(self) -> list[Path]:
        return [resolve(rel) for rel, _ in FAST_FIXTURES]

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=5, warmup_rounds=1)
    async def test_serial(self, async_benchmark, fast_paths):
        result = await async_benchmark(batch_serial, fast_paths, "fast")
        assert result["mean"] > 0

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=5, warmup_rounds=1)
    async def test_concurrent_2(self, async_benchmark, fast_paths):
        result = await async_benchmark(batch_concurrent, fast_paths, "fast", 2)
        assert result["mean"] > 0

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=5, warmup_rounds=1)
    async def test_concurrent_4(self, async_benchmark, fast_paths):
        result = await async_benchmark(batch_concurrent, fast_paths, "fast", 4)
        assert result["mean"] > 0


# ===================================================================
# OCR pipeline -- hi_res strategy (the hot path)
#
# Exercises: render PDF pages -> layout detection (ONNX) -> OCR
# (Tesseract) -> element merge.  This is what experimental-agent
# optimized and where native async will have the biggest impact.
#
# Baseline (main, taskset -c 0):
#   single 1p:           ~8s
#   3 docs serial:      ~46s
#   3 docs concurrent:  ~44s  (only 5% gain -- GIL-bound)
# ===================================================================
class TestOCRPipeline:
    """OCR pipeline latency (hi_res, single doc)."""

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=2, warmup_rounds=1)
    async def test_hires_single(self, async_benchmark):
        """Single 1-page PDF through full hi_res pipeline."""
        filepath = resolve(HIRES_FIXTURE[0])
        result = await async_benchmark(partition_async, filepath, "hi_res")
        assert result["mean"] > 0

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=2, warmup_rounds=1)
    async def test_hires_serial_3docs(self, async_benchmark):
        """3 docs through hi_res, serial baseline."""
        paths = [resolve(rel) for rel, _ in FAST_FIXTURES]
        result = await async_benchmark(batch_serial, paths, "hi_res")
        assert result["mean"] > 0

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=2, warmup_rounds=1)
    async def test_hires_concurrent_3docs(self, async_benchmark):
        """3 docs through hi_res, concurrent -- measures async gain."""
        paths = [resolve(rel) for rel, _ in FAST_FIXTURES]
        result = await async_benchmark(batch_concurrent, paths, "hi_res", 3)
        assert result["mean"] > 0


# ===================================================================
# Per-doc latency -- fast strategy, parametrized
# ===================================================================
@pytest.mark.parametrize(
    "fixture_rel,fixture_id",
    FAST_FIXTURES,
    ids=[fid for _, fid in FAST_FIXTURES],
)
class TestPerDocLatency:
    """Per-document partition latency (fast strategy)."""

    @pytest.mark.asyncio
    @pytest.mark.async_benchmark(rounds=5, warmup_rounds=1)
    async def test_fast(self, async_benchmark, fixture_rel, fixture_id):
        filepath = resolve(fixture_rel)
        result = await async_benchmark(partition_async, filepath, "fast")
        assert result["mean"] > 0
