#!/usr/bin/env python3
"""
Job-mode regression tests.

These checks validate end-to-end file-channel job mode behavior for both normal
and empty input files.
"""

from __future__ import annotations

import json
import sys
import tempfile
import time
import types
from pathlib import Path

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

from pydantic import BaseModel

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

from sandai.operator import BatchProcessor, TaskInput, TaskOutput


class JobOptions(BaseModel):
    delay: float = 0.1
    message: str = "test"


class JobResults(BaseModel):
    processed_at: float
    batch_size: int
    message: str


def _write_jsonl(path: Path, 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: Path) -> list[dict]:
    if not path.exists():
        return []
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def test_job_mode_processes_all_tasks_and_keeps_batch_metadata() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "job_mode_input.jsonl"
        output_path = Path(tmpdir) / "job_mode_output.jsonl"

        rows = [
            {
                "task_id": f"job-{index:03d}",
                "options": {"delay": 0.01, "message": f"task-{index}"},
                "artifacts": [],
            }
            for index in range(7)
        ]
        _write_jsonl(input_path, rows)

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

        observed_batch_sizes = []

        @processor.on_batch(max_concurrency=2, max_batch_size=3, max_prefetch_size=10)
        def process_batch(batch_inputs: list[TaskInput[JobOptions]], operator_config: dict, context):
            observed_batch_sizes.append(len(batch_inputs))
            time.sleep(0.05)
            now = time.time()
            for task_input in batch_inputs:
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results=JobResults(
                        processed_at=now,
                        batch_size=len(batch_inputs),
                        message=f"Processed: {task_input.options.message}",
                    ),
                    artifacts=[],
                )

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

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

        outputs = _read_jsonl(output_path)
        assert len(outputs) == len(rows), len(outputs)
        assert sorted(observed_batch_sizes) == [1, 3, 3], observed_batch_sizes
        assert elapsed < 5.0, elapsed

        batch_sizes = sorted(item["results"]["batch_size"] for item in outputs)
        assert batch_sizes == [1, 3, 3, 3, 3, 3, 3], batch_sizes
        assert all(item["results"]["message"].startswith("Processed: task-") for item in outputs)


def test_job_mode_exits_cleanly_for_empty_input() -> None:
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = Path(tmpdir) / "empty_input.jsonl"
        output_path = Path(tmpdir) / "empty_output.jsonl"
        input_path.write_text("", encoding="utf-8")

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

        @processor.on_batch()
        def process_batch(batch_inputs, operator_config, context):
            for task_input in batch_inputs:
                yield TaskOutput(
                    task_id=task_input.task_id,
                    results={"processed": True},
                    artifacts=[],
                )

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

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

        assert elapsed < 3.0, elapsed
        outputs = _read_jsonl(output_path)
        assert outputs == [], outputs


def test_run_does_not_install_parent_guard() -> None:
    processor = BatchProcessor(
        name="guard-check",
        version="1.0.0",
        ignore_nongil_check=True,
    )

    calls: list[str] = []
    original_caoe = sys.modules.get("caoe")
    sys.modules["caoe"] = types.SimpleNamespace(install=lambda: calls.append("install"))

    async def _async_run(signal_handler=None, job_mode=False):
        return None

    original_async_run = processor._async_run
    processor._async_run = _async_run  # type: ignore[method-assign]
    try:
        processor.run()
    finally:
        processor._async_run = original_async_run  # type: ignore[method-assign]
        if original_caoe is None:
            sys.modules.pop("caoe", None)
        else:
            sys.modules["caoe"] = original_caoe

    assert calls == []


def main() -> int:
    tests = [
        ("job_mode_end_to_end", test_job_mode_processes_all_tasks_and_keeps_batch_metadata),
        ("empty_input", test_job_mode_exits_cleanly_for_empty_input),
        ("parent_guard_removed", test_run_does_not_install_parent_guard),
    ]

    passed = 0
    failed = 0

    print("Job mode regression tests")
    print("=" * 60)
    for name, test_func in tests:
        try:
            test_func()
            print(f"PASS {name}")
            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())
