#!/usr/bin/env python3
"""
Processor contract regression tests.

These tests focus on stricter runtime validation introduced in BatchProcessor:
typed output enforcement, duplicate/unknown task ids, and richer input conversion
errors.
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path
from typing import Generator, List

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.processor import (
    BatchProcessor,
    OperatorConfigurationError,
    OutputEnvelope,
    PreparedTask,
)
from sandai.operator.types import TaskInput, TaskOutput


class ContractOptions(BaseModel):
    value: int = 1


class ContractResults(BaseModel):
    echoed: int


def _base_processor() -> BatchProcessor:
    processor = BatchProcessor("contract-processor", "v1", ignore_nongil_check=True)
    processor._setup_temp_directory()
    return processor


def test_convert_input_to_taskinput_rejects_missing_task_id() -> None:
    processor = _base_processor()

    async def run_case() -> None:
        try:
            await processor.convert_input_to_taskinput({"options": {}, "artifacts": []})
            raise AssertionError("expected missing task_id failure")
        except Exception as exc:
            assert "task_id" in str(exc)

    asyncio.run(run_case())


def test_convert_input_to_taskinput_rejects_invalid_artifact_shape() -> None:
    processor = _base_processor()

    async def run_case() -> None:
        try:
            await processor.convert_input_to_taskinput(
                {"task_id": "bad-artifact", "options": {}, "artifacts": ["oops"]}
            )
            raise AssertionError("expected invalid artifact failure")
        except Exception as exc:
            assert "artifact" in str(exc).lower()

    asyncio.run(run_case())


def test_safe_process_func_rejects_non_taskoutput_values() -> None:
    processor = _base_processor()

    @processor.on_batch(max_batch_size=1)
    def process_batch(
        batch_inputs: List[TaskInput[ContractOptions]],
        operator_config: dict,
        context,
    ) -> Generator[TaskOutput[ContractResults], None, None]:
        yield {"task_id": batch_inputs[0].task_id, "status": "success"}  # type: ignore[misc]

    task_input = TaskInput[ContractOptions](
        task_id="task-1",
        options=ContractOptions(value=3),
        artifacts=[],
    )
    outputs = processor._safe_process_func([task_input], {})

    assert len(outputs) == 1
    assert outputs[0].status == "error"
    assert outputs[0].error is not None
    assert outputs[0].error.error_type == "OutputValidationError"


def test_safe_process_func_rejects_duplicate_task_ids() -> None:
    processor = _base_processor()

    @processor.on_batch(max_batch_size=2)
    def process_batch(
        batch_inputs: List[TaskInput[ContractOptions]],
        operator_config: dict,
        context,
    ) -> Generator[TaskOutput[ContractResults], None, None]:
        first_task_id = batch_inputs[0].task_id
        yield TaskOutput(task_id=first_task_id, results=ContractResults(echoed=1))
        yield TaskOutput(task_id=first_task_id, results=ContractResults(echoed=2))

    task_inputs = [
        TaskInput[ContractOptions](task_id="task-a", options=ContractOptions(value=1), artifacts=[]),
        TaskInput[ContractOptions](task_id="task-b", options=ContractOptions(value=2), artifacts=[]),
    ]
    outputs = processor._safe_process_func(task_inputs, {})

    assert len(outputs) == 2
    assert outputs[0].status == "success"
    assert outputs[1].status == "error"
    assert outputs[1].error is not None
    assert outputs[1].error.error_type == "OutputValidationError"


def test_safe_process_func_rejects_unknown_task_ids() -> None:
    processor = _base_processor()

    @processor.on_batch(max_batch_size=1)
    def process_batch(
        batch_inputs: List[TaskInput[ContractOptions]],
        operator_config: dict,
        context,
    ) -> Generator[TaskOutput[ContractResults], None, None]:
        yield TaskOutput(task_id="foreign-task", results=ContractResults(echoed=99))

    task_input = TaskInput[ContractOptions](
        task_id="known-task",
        options=ContractOptions(value=4),
        artifacts=[],
    )
    outputs = processor._safe_process_func([task_input], {})

    assert len(outputs) == 1
    assert outputs[0].status == "error"
    assert outputs[0].error is not None
    assert outputs[0].error.error_type == "OutputValidationError"


def test_execute_batch_rejects_output_count_mismatch() -> None:
    processor = _base_processor()

    async def broken_process_batch(task_inputs):
        return [
            TaskOutput(
                task_id=task_inputs[0].task_id,
                results=ContractResults(echoed=task_inputs[0].options.value),
            )
        ]

    processor.process_batch = broken_process_batch  # type: ignore[method-assign]

    prepared_tasks = [
        PreparedTask(
            task_input=TaskInput[ContractOptions](
                task_id="task-a", options=ContractOptions(value=1), artifacts=[]
            )
        ),
        PreparedTask(
            task_input=TaskInput[ContractOptions](
                task_id="task-b", options=ContractOptions(value=2), artifacts=[]
            )
        ),
    ]

    output_batch = asyncio.run(processor._execute_batch(prepared_tasks))

    assert len(output_batch) == 2
    assert all(isinstance(envelope, OutputEnvelope) for envelope in output_batch)
    assert all(envelope.output.status == "error" for envelope in output_batch)
    assert output_batch[0].output.error is not None
    assert output_batch[0].output.error.error_type == "ProcessingError"
    assert "different number of outputs" in output_batch[0].output.error.error_message


def test_convert_taskoutput_to_output_serializes_nested_results_model() -> None:
    processor = _base_processor()
    processor._output_executor = None

    async def run_case() -> None:
        output = TaskOutput(
            task_id="task-nested",
            results=ContractResults(echoed=7),
            artifacts=[],
        )
        converted = await processor.convert_taskoutput_to_output(output)
        assert converted["task_id"] == "task-nested"
        assert converted["results"] == {"echoed": 7}

    asyncio.run(run_case())


def test_async_run_requires_registered_process_function() -> None:
    processor = BatchProcessor("missing-handler", "v1", ignore_nongil_check=True)

    async def run_case() -> None:
        try:
            await processor._async_run(job_mode=True)
            raise AssertionError("expected missing process function failure")
        except OperatorConfigurationError as exc:
            assert "No process function defined" in str(exc)

    asyncio.run(run_case())


def main() -> int:
    tests = [
        ("missing_task_id", test_convert_input_to_taskinput_rejects_missing_task_id),
        ("invalid_artifact", test_convert_input_to_taskinput_rejects_invalid_artifact_shape),
        ("non_taskoutput", test_safe_process_func_rejects_non_taskoutput_values),
        ("duplicate_task_id", test_safe_process_func_rejects_duplicate_task_ids),
        ("unknown_task_id", test_safe_process_func_rejects_unknown_task_ids),
        ("output_count_mismatch", test_execute_batch_rejects_output_count_mismatch),
        ("serialize_nested_results", test_convert_taskoutput_to_output_serializes_nested_results_model),
        ("missing_process_func", test_async_run_requires_registered_process_function),
    ]

    passed = 0
    failed = 0

    print("Processor contract 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())
