#!/usr/bin/env python3
"""
Pipeline behavior tests for prefetch, concurrency, and batch size.

These tests run as a standalone script and avoid optional third-party runtime
dependencies by installing lightweight stubs before importing the SDK.
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import tempfile
import threading
import time
from pathlib import Path

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

sys.path.insert(0, str(Path(__file__).parent.parent))

from pydantic import BaseModel

from sandai.operator import BatchProcessor, TaskInput, TaskOutput, OutputArtifact


class PerfOptions(BaseModel):
    payload: int = 1


class PerfResults(BaseModel):
    started_at: float
    finished_at: float
    batch_size: int


def _write_jsonl(path: str, rows: list[dict]) -> None:
    with open(path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")


def _read_jsonl(path: str) -> list[dict]:
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def test_batch_size_groups_tasks() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "batch_input.jsonl")
        output_path = os.path.join(tmpdir, "batch_output.jsonl")

        rows = [
            {"task_id": f"batch-{i}", "options": {"payload": i}, "artifacts": []}
            for i in range(5)
        ]
        _write_jsonl(input_path, rows)

        observed_batch_sizes = []
        processor = BatchProcessor(
            name="batch-size-check",
            version="1.0.0",
            max_concurrency=1,
            max_batch_size=2,
            max_prefetch_size=4,
            ignore_nongil_check=True,
        )

        @processor.on_batch(max_concurrency=1, max_batch_size=2, max_prefetch_size=4)
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            observed_batch_sizes.append(len(batch_inputs))
            now = time.time()
            for task_input in batch_inputs:
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results=PerfResults(
                        started_at=now,
                        finished_at=now,
                        batch_size=len(batch_inputs),
                    ),
                    artifacts=[],
                )

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }
        processor.run(job_mode=True)

        outputs = _read_jsonl(output_path)
        batch_sizes = sorted(result["results"]["batch_size"] for result in outputs)

        assert len(outputs) == 5, f"expected 5 outputs, got {len(outputs)}"
        assert observed_batch_sizes == [2, 2, 1], observed_batch_sizes
        assert batch_sizes == [1, 2, 2, 2, 2], batch_sizes
        return True


def test_concurrency_runs_batches_in_parallel() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "concurrency_input.jsonl")
        output_path = os.path.join(tmpdir, "concurrency_output.jsonl")

        rows = [
            {"task_id": f"concurrency-{i}", "options": {"payload": i}, "artifacts": []}
            for i in range(4)
        ]
        _write_jsonl(input_path, rows)

        processor = BatchProcessor(
            name="concurrency-check",
            version="1.0.0",
            max_concurrency=2,
            max_batch_size=1,
            max_prefetch_size=2,
            ignore_nongil_check=True,
        )

        batch_start_times = []
        lock = threading.Lock()

        @processor.on_batch(max_concurrency=2, max_batch_size=1, max_prefetch_size=2)
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            start = time.time()
            with lock:
                batch_start_times.append(start)
            time.sleep(0.25)
            end = time.time()
            for task_input in batch_inputs:
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results=PerfResults(
                        started_at=start,
                        finished_at=end,
                        batch_size=1,
                    ),
                    artifacts=[],
                )

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }

        started = time.time()
        processor.run(job_mode=True)
        elapsed = time.time() - started

        outputs = _read_jsonl(output_path)
        assert len(outputs) == 4, len(outputs)
        assert len(batch_start_times) == 4, batch_start_times
        sorted_starts = sorted(batch_start_times)
        overlap_gap = sorted_starts[1] - sorted_starts[0]
        compute_span = max(batch_start_times) - min(batch_start_times)

        assert overlap_gap < 0.20, overlap_gap
        assert compute_span < 0.60, compute_span
        assert elapsed < 1.80, elapsed
        return True


def test_prefetch_prepares_next_tasks_during_compute() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "prefetch_input.jsonl")
        output_path = os.path.join(tmpdir, "prefetch_output.jsonl")

        payloads = []
        for i in range(4):
            file_path = os.path.join(tmpdir, f"artifact_{i}.txt")
            Path(file_path).write_text(f"artifact-{i}", encoding="utf-8")
            payloads.append(
                {
                    "task_id": f"prefetch-{i}",
                    "options": {"payload": i},
                    "artifacts": [{"uri": file_path, "name": f"artifact-{i}"}],
                }
            )

        _write_jsonl(input_path, payloads)

        processor = BatchProcessor(
            name="prefetch-check",
            version="1.0.0",
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=3,
            ignore_nongil_check=True,
        )

        original_uri_to_path = processor._uri_to_path
        converted_task_ids = []
        conversion_times = {}
        lock = threading.Lock()

        def tracked_uri_to_path(uri: str, task_id: str, file_type: str = "inputs"):
            result = original_uri_to_path(uri, task_id, file_type)
            with lock:
                converted_task_ids.append(task_id)
                conversion_times.setdefault(task_id, time.time())
            time.sleep(0.05)
            return result

        processor._uri_to_path = tracked_uri_to_path

        compute_windows = []

        @processor.on_batch(max_concurrency=1, max_batch_size=1, max_prefetch_size=3)
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            start = time.time()
            time.sleep(0.30)
            end = time.time()
            task_id = batch_inputs[0].task_id
            compute_windows.append((task_id, start, end))
            artifact_path = os.path.join(context.get_task_workdir(task_id), f"{task_id}.txt")
            Path(artifact_path).write_text(task_id, encoding="utf-8")
            yield TaskOutput(
                task_id=task_id,
                results=PerfResults(
                    started_at=start,
                    finished_at=end,
                    batch_size=1,
                ),
                artifacts=[OutputArtifact(uri=artifact_path, name=f"{task_id}.txt")],
            )

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }
        processor.run(job_mode=True)

        outputs = _read_jsonl(output_path)
        assert len(outputs) == 4, len(outputs)
        assert converted_task_ids[:2] == ["prefetch-0", "prefetch-1"], converted_task_ids

        first_task_id, first_start, first_end = compute_windows[0]
        second_task_id = compute_windows[1][0]
        assert conversion_times[first_task_id] <= first_start
        assert conversion_times[second_task_id] < first_end, {
            "first_compute": (first_start, first_end),
            "second_conversion": conversion_times[second_task_id],
        }
        return True


def test_default_stage_worker_counts_match_max_concurrency() -> bool:
    processor = BatchProcessor(
        name="defaults-check",
        version="1.0.0",
        max_concurrency=3,
        max_batch_size=2,
        max_prefetch_size=4,
        ignore_nongil_check=True,
    )

    assert processor._prepare_worker_count() == 3
    assert processor._compute_worker_count() == 3
    assert processor._output_worker_count() == 3

    @processor.on_batch(max_concurrency=4, max_batch_size=2, max_prefetch_size=4)
    def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
        for task_input in batch_inputs:
            yield TaskOutput(task_id=task_input.task_id, results=None, artifacts=[])

    assert processor._prepare_worker_count() == 4
    assert processor._compute_worker_count() == 4
    assert processor._output_worker_count() == 4
    return True


def test_prepare_concurrency_can_exceed_compute_concurrency() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "stage_override_input.jsonl")
        output_path = os.path.join(tmpdir, "stage_override_output.jsonl")

        rows = [
            {"task_id": f"stage-{i}", "options": {"payload": i}, "artifacts": []}
            for i in range(4)
        ]
        _write_jsonl(input_path, rows)

        processor = BatchProcessor(
            name="stage-override-check",
            version="1.0.0",
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=4,
            prepare_concurrency=2,
            ignore_nongil_check=True,
        )

        original_convert = processor.convert_input_to_taskinput
        conversion_starts = []
        lock = threading.Lock()

        async def tracked_convert(raw_input):
            with lock:
                conversion_starts.append((raw_input["task_id"], time.time()))
            await asyncio.sleep(0.15)
            return await original_convert(raw_input)

        import asyncio

        processor.convert_input_to_taskinput = tracked_convert

        @processor.on_batch(
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=4,
            prepare_concurrency=2,
        )
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            time.sleep(0.25)
            for task_input in batch_inputs:
                yield TaskOutput(task_id=task_input.task_id, results=None, artifacts=[])

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }
        processor.run(job_mode=True)

        assert len(_read_jsonl(output_path)) == 4
        assert len(conversion_starts) >= 2
        sorted_starts = sorted(ts for _, ts in conversion_starts[:2])
        assert sorted_starts[1] - sorted_starts[0] < 0.10, conversion_starts
        return True


def test_output_concurrency_runs_uploads_in_parallel() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "output_concurrency_input.jsonl")
        output_path = os.path.join(tmpdir, "output_concurrency_output.jsonl")

        rows = [
            {"task_id": f"output-stage-{i}", "options": {"payload": i}, "artifacts": []}
            for i in range(4)
        ]
        _write_jsonl(input_path, rows)

        processor = BatchProcessor(
            name="output-stage-check",
            version="1.0.0",
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=4,
            output_concurrency=2,
            ignore_nongil_check=True,
        )

        upload_starts = []
        lock = threading.Lock()

        async def tracked_upload(local_file_path: str, ephemeral: bool = True):
            with lock:
                upload_starts.append((Path(local_file_path).name, time.time()))
            await asyncio.sleep(0.25)
            return f"fake://uploaded/{Path(local_file_path).name}", {"type": "fake"}

        processor._upload_to_storage = tracked_upload

        @processor.on_batch(
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=4,
            output_concurrency=2,
        )
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            for task_input in batch_inputs:
                artifact_path = os.path.join(
                    context.get_task_workdir(task_input.task_id),
                    f"{task_input.task_id}.txt",
                )
                Path(artifact_path).write_text(task_input.task_id, encoding="utf-8")
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results=None,
                    artifacts=[OutputArtifact(uri=artifact_path, name=f"{task_input.task_id}.txt")],
                )

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }

        started = time.time()
        processor.run(job_mode=True)
        elapsed = time.time() - started

        outputs = _read_jsonl(output_path)
        assert len(outputs) == 4, len(outputs)
        assert len(upload_starts) == 4, upload_starts

        sorted_starts = sorted(ts for _, ts in upload_starts)
        assert sorted_starts[1] - sorted_starts[0] < 0.10, upload_starts
        assert elapsed < 1.60, elapsed
        assert all(output["artifacts"][0]["uri"].startswith("fake://uploaded/") for output in outputs)
        return True


def test_job_mode_waits_for_output_flush() -> bool:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "job_mode_flush_input.jsonl")
        output_path = os.path.join(tmpdir, "job_mode_flush_output.jsonl")

        rows = [
            {"task_id": f"flush-{i}", "options": {"payload": i}, "artifacts": []}
            for i in range(2)
        ]
        _write_jsonl(input_path, rows)

        processor = BatchProcessor(
            name="job-mode-flush-check",
            version="1.0.0",
            max_concurrency=1,
            max_batch_size=1,
            max_prefetch_size=2,
            ignore_nongil_check=True,
        )

        async def slow_upload(local_file_path: str, ephemeral: bool = True):
            await asyncio.sleep(0.15)
            return f"fake://uploaded/{Path(local_file_path).name}", {"type": "fake"}

        processor._upload_to_storage = slow_upload

        @processor.on_batch(max_concurrency=1, max_batch_size=1, max_prefetch_size=2)
        def process_batch(batch_inputs: list[TaskInput[PerfOptions]], operator_config: dict, context):
            for task_input in batch_inputs:
                artifact_path = os.path.join(
                    context.get_task_workdir(task_input.task_id),
                    f"{task_input.task_id}.txt",
                )
                Path(artifact_path).write_text(task_input.task_id, encoding="utf-8")
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results=None,
                    artifacts=[OutputArtifact(uri=artifact_path, name=f"{task_input.task_id}.txt")],
                )

        processor.config = {
            "channel.file.input_path": input_path,
            "channel.file.output_path": output_path,
            "operator.config": {},
        }
        processor.run(job_mode=True)

        outputs = _read_jsonl(output_path)
        assert len(outputs) == 2, outputs
        assert all(output["artifacts"][0]["uri"].startswith("fake://uploaded/") for output in outputs)
        return True


def main() -> int:
    tests = [
        ("batch_size", test_batch_size_groups_tasks),
        ("concurrency", test_concurrency_runs_batches_in_parallel),
        ("prefetch", test_prefetch_prepares_next_tasks_during_compute),
        ("stage_default_counts", test_default_stage_worker_counts_match_max_concurrency),
        ("prepare_override", test_prepare_concurrency_can_exceed_compute_concurrency),
        ("output_concurrency", test_output_concurrency_runs_uploads_in_parallel),
        ("job_mode_flush", test_job_mode_waits_for_output_flush),
    ]

    passed = 0
    failed = 0

    print("Pipeline behavior tests")
    print("=" * 60)

    for name, test_func in tests:
        try:
            start = time.time()
            test_func()
            elapsed = time.time() - start
            print(f"PASS {name} ({elapsed:.2f}s)")
            passed += 1
        except Exception as exc:
            print(f"FAIL {name}: {exc}")
            import traceback

            traceback.print_exc()
            failed += 1

    print("=" * 60)
    print(f"passed={passed} failed={failed}")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())
