#!/usr/bin/env python3
"""
Description-query regression tests.

These checks cover doc/schema preparation and direct `_handle_desc_query`
responses for typed and untyped processors.
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import Any, 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
from sandai.operator.processor import OperatorConfigurationError
from sandai.operator.types import TaskInput, TaskOutput


class DescOptions(BaseModel):
    name: str
    count: int = 10
    enabled: bool = True


class DescResults(BaseModel):
    message: str
    processed_count: int
    success: bool = True


def _build_typed_processor() -> BatchProcessor:
    processor = BatchProcessor("test-desc-processor", "v1.0", ignore_nongil_check=True)

    @processor.on_batch(max_batch_size=2)
    def test_process_function(
        batch_inputs: List[TaskInput[DescOptions]],
        config: Any,
        context: Any,
    ) -> Generator[TaskOutput[DescResults], None, None]:
        """
        Test processing function for description-query coverage.

        This function processes a batch of inputs and emits typed outputs so the
        SDK can derive both option and result schemas.

        Args:
            batch_inputs: Inputs using DescOptions.
            config: Operator configuration.
            context: Processing context.

        Yields:
            TaskOutput values using DescResults.
        """
        for task_input in batch_inputs:
            yield TaskOutput[DescResults](
                task_id=task_input.task_id,
                status="success",
                results=DescResults(
                    message=f"Processed {task_input.options.name}",
                    processed_count=task_input.options.count * 2,
                    success=True,
                ),
                artifacts=[],
            )

    return processor


def _build_untyped_processor() -> BatchProcessor:
    processor = BatchProcessor("no-types-processor", "v1.0", ignore_nongil_check=True)

    @processor.on_batch(max_batch_size=1)
    def simple_function(batch_inputs, config, context):
        """Simple function without type annotations."""
        for task_input in batch_inputs:
            yield {
                "task_id": task_input.task_id,
                "status": "success",
                "results": {"message": "processed"},
            }

    return processor


def test_function_description_is_prepared_during_decoration() -> None:
    processor = _build_typed_processor()
    desc = processor._function_description

    assert desc is not None
    assert "doc" in desc
    assert "options_schema" in desc
    assert "results_schema" in desc
    assert desc["doc"].strip() != ""
    assert desc["options_schema"] is not None
    assert desc["results_schema"] is not None


def test_docstring_extraction_contains_expected_sections() -> None:
    processor = _build_typed_processor()
    doc = processor._function_description["doc"]

    assert "Test processing function for description-query coverage." in doc
    assert "Args:" in doc
    assert "Yields:" in doc


def test_options_schema_matches_model_shape() -> None:
    processor = _build_typed_processor()
    schema = processor._function_description["options_schema"]
    properties = schema["properties"]

    assert schema["type"] == "object"
    assert properties["name"]["type"] == "string"
    assert "count" in properties
    assert "enabled" in properties
    assert "name" in schema.get("properties", {})


def test_results_schema_matches_model_shape() -> None:
    processor = _build_typed_processor()
    schema = processor._function_description["results_schema"]
    properties = schema["properties"]

    assert schema["type"] == "object"
    assert properties["message"]["type"] == "string"
    assert "processed_count" in properties
    assert "success" in properties


def test_handle_desc_query_returns_expected_payload() -> None:
    processor = _build_typed_processor()
    result = processor._handle_desc_query({"task_id": "desc-001", "task_type": "desc"})

    assert result["task_id"] == "desc-001"
    assert result["status"] == "success"
    assert "doc" in result
    assert "options_schema" in result
    assert "results_schema" in result
    assert result["__keep_raw_output__"] is True


def test_generated_schemas_accept_expected_field_types() -> None:
    processor = _build_typed_processor()
    schema = processor._function_description["options_schema"]
    properties = schema["properties"]
    valid_options = {"name": "test_task", "count": 5, "enabled": False}

    assert isinstance(valid_options["name"], str)
    assert isinstance(valid_options["count"], int)
    assert isinstance(valid_options["enabled"], bool)
    assert set(valid_options).issubset(properties)


def test_schema_payloads_are_structurally_complete() -> None:
    processor = _build_typed_processor()
    desc = processor._function_description

    for schema_name in ["options_schema", "results_schema"]:
        schema = desc[schema_name]
        assert "type" in schema
        assert "properties" in schema
        for prop_def in schema["properties"].values():
            assert "type" in prop_def


def test_desc_query_falls_back_when_description_is_missing() -> None:
    processor = _build_typed_processor()
    original_desc = processor._function_description
    processor._function_description = None
    processor.process_func = None

    result = processor._handle_desc_query({"task_id": "error-test", "task_type": "desc"})

    assert result["task_id"] == "error-test"
    assert result["status"] == "success"
    assert "error" in result

    processor._function_description = original_desc


def test_multiple_desc_queries_return_stable_schemas() -> None:
    processor = _build_typed_processor()
    results = [
        processor._handle_desc_query({"task_id": f"desc-{index}", "task_type": "desc"})
        for index in range(3)
    ]

    assert [result["task_id"] for result in results] == ["desc-0", "desc-1", "desc-2"]
    assert all(result["status"] == "success" for result in results)
    assert all(result["options_schema"] == results[0]["options_schema"] for result in results[1:])
    assert all(result["results_schema"] == results[0]["results_schema"] for result in results[1:])


def test_desc_query_without_types_still_returns_doc() -> None:
    processor = _build_untyped_processor()
    result = processor._handle_desc_query({"task_id": "no-types", "task_type": "desc"})

    assert result["status"] == "success"
    assert "doc" in result


def test_invalid_batch_input_annotation_is_rejected() -> None:
    processor = BatchProcessor("invalid-input-processor", "v1", ignore_nongil_check=True)

    try:
        @processor.on_batch(max_batch_size=1)
        def invalid_process(batch_inputs: DescOptions, config: Any, context: Any):
            yield TaskOutput(task_id="x", status="success")

        raise AssertionError("expected OperatorConfigurationError")
    except OperatorConfigurationError as exc:
        assert "batch_inputs" in str(exc)


def test_invalid_result_annotation_is_rejected() -> None:
    processor = BatchProcessor("invalid-result-processor", "v1", ignore_nongil_check=True)

    try:
        @processor.on_batch(max_batch_size=1)
        def invalid_process(
            batch_inputs: List[TaskInput[DescOptions]],
            config: Any,
            context: Any,
        ) -> Generator[dict, None, None]:
            yield {"task_id": "x", "status": "success"}

        raise AssertionError("expected OperatorConfigurationError")
    except OperatorConfigurationError as exc:
        assert "TaskOutput" in str(exc)


def test_runtime_check_is_compatible_noop() -> None:
    processor = BatchProcessor("runtime-check", "v1")
    assert processor.check_non_gil() is None


def main() -> int:
    tests = [
        ("prepare_description", test_function_description_is_prepared_during_decoration),
        ("docstring_extraction", test_docstring_extraction_contains_expected_sections),
        ("options_schema", test_options_schema_matches_model_shape),
        ("results_schema", test_results_schema_matches_model_shape),
        ("handle_desc_query", test_handle_desc_query_returns_expected_payload),
        ("schema_types", test_generated_schemas_accept_expected_field_types),
        ("schema_completeness", test_schema_payloads_are_structurally_complete),
        ("missing_description", test_desc_query_falls_back_when_description_is_missing),
        ("multiple_queries", test_multiple_desc_queries_return_stable_schemas),
        ("untyped_processor", test_desc_query_without_types_still_returns_doc),
        ("invalid_batch_inputs", test_invalid_batch_input_annotation_is_rejected),
        ("invalid_result_annotation", test_invalid_result_annotation_is_rejected),
        ("runtime_requirement", test_runtime_check_is_compatible_noop),
    ]

    passed = 0
    failed = 0

    print("Description-query 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())
