#!/usr/bin/env python3
"""
Focused regression tests for client timeout/cancel payloads and processor-side
skip logic.
"""

from __future__ import annotations

import sys
import time
from pathlib import Path
from unittest.mock import Mock, patch

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

TESTS_DIR = Path(__file__).parent
REPO_ROOT = TESTS_DIR.parent.parent
sys.path.insert(0, str(REPO_ROOT / "operator-client"))
sys.path.insert(0, str(REPO_ROOT / "operator-sdk"))

from sandai.operator import BatchProcessor
from sandai.operator.channel import CeleryChannel
from sandai.operator_client import OperatorClient, create_client
from sandai.operator_client.config import CeleryConfig, ClientConfig
from sandai.operator_client.exceptions import (
    TaskCancelledError,
    TaskDeadlineExceededError,
)


class _ChannelConf:
    def __init__(self, **values):
        self._values = values
        for key, value in values.items():
            setattr(self, key, value)

    def get(self, key, default=None):
        return self._values.get(key, default)


def _mock_client_runtime():
    config = ClientConfig(
        operator_name="test-operator",
        operator_version="v1",
        celery=CeleryConfig(
            broker_url="redis://localhost:6379/0",
            result_backend="redis://localhost:6379/0",
        ),
    )
    return patch(
        "sandai.operator_client.client.resolve_client_config",
        return_value=config,
    ), patch("sandai.operator_client.client.CeleryTransport")


def test_default_timeout_and_create_client() -> None:
    with _mock_client_runtime()[0], _mock_client_runtime()[1] as mock_transport:
        mock_transport.return_value = Mock(app=Mock())
        client = OperatorClient("test-operator", "v1")
        created = create_client("test-operator", "v1")

    assert client.timeout == 28800, client.timeout
    assert created.timeout == 28800, created.timeout


def test_async_call_sets_deadline_and_cancel_key() -> None:
    with _mock_client_runtime()[0], _mock_client_runtime()[1] as mock_transport_cls:
        mock_transport = Mock()
        mock_result = Mock()
        mock_result.id = "task-123"
        mock_transport.app = Mock()
        mock_transport.send_task.return_value = mock_result
        mock_transport_cls.return_value = mock_transport

        client = OperatorClient("test-operator", "v1", timeout=300, cancel_key="default-session")
        start = time.time()
        client.async_call(
            input_artifacts=[{"uri": "file://input.txt", "slot": "input"}],
            options={"mode": "fast"},
            timeout=120,
            cancel_key="override-session",
        )

        kwargs = mock_transport.send_task.call_args.kwargs["kwargs"]

    assert kwargs["cancel_key"] == "override-session", kwargs
    assert abs(kwargs["deadline"] - (start + 120)) < 1.0, kwargs["deadline"]
    assert kwargs["artifacts"][0]["slot"] == "input"
    assert kwargs["options"] == {"mode": "fast"}


def test_sync_raises_deadline_and_cancel_errors() -> None:
    deadline_result = {
        "status": "error",
        "error": {
            "error_type": "TaskDeadlineExceededError",
            "error_message": "deadline exceeded",
            "error_details": {
                "deadline": 100.0,
                "current_time": 101.5,
                "exceeded_by": 1.5,
            },
        },
    }
    cancel_result = {
        "status": "error",
        "error": {
            "error_type": "TaskCancelledError",
            "error_message": "cancelled",
            "error_details": {
                "cancel_key": "cancel:user-1",
            },
        },
    }

    with _mock_client_runtime()[0], _mock_client_runtime()[1] as mock_transport_cls:
        mock_transport = Mock()
        mock_transport.app = Mock()
        mock_transport_cls.return_value = mock_transport

        deadline_async_result = Mock()
        deadline_async_result.id = "deadline-task"

        cancel_async_result = Mock()
        cancel_async_result.id = "cancel-task"

        mock_transport.send_task.side_effect = [deadline_async_result, cancel_async_result]
        mock_transport.get_result.side_effect = [deadline_result, cancel_result]

        client = OperatorClient("test-operator", "v1")

        try:
            client.sync([], {}, timeout=60)
            raise AssertionError("expected TaskDeadlineExceededError")
        except TaskDeadlineExceededError as exc:
            assert exc.task_id == "deadline-task"
            assert exc.deadline == 100.0
            assert exc.current_time == 101.5

        try:
            client.sync([], {}, cancel_key="cancel:user-1")
            raise AssertionError("expected TaskCancelledError")
        except TaskCancelledError as exc:
            assert exc.task_id == "cancel-task"
            assert exc.cancel_key == "cancel:user-1"


def test_processor_deadline_skip_and_error_output() -> None:
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    now = time.time()

    expired = {
        "task_id": "expired-task",
        "artifacts": [],
        "options": {},
        "deadline": now - 10,
    }
    valid = {
        "task_id": "valid-task",
        "artifacts": [],
        "options": {},
        "deadline": now + 30,
    }

    assert processor._should_skip_due_to_deadline(expired) is True
    assert processor._should_skip_due_to_deadline(valid) is False

    error_output = processor._create_deadline_error_output(expired)
    assert error_output["task_id"] == "expired-task"
    assert error_output["error"]["error_type"] == "TaskDeadlineExceededError"
    assert error_output["error"]["error_details"]["skipped_before_processing"] is True
    assert error_output["__keep_raw_output__"] is True


def test_processor_cancel_skip_caches_redis_and_formats_error() -> None:
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    processor._redis_client = Mock()
    processor._redis_client.exists.side_effect = [1]

    cancelled = {
        "task_id": "cancelled-task",
        "artifacts": [],
        "options": {},
        "cancel_key": "cancel:user-42",
    }
    normal = {
        "task_id": "normal-task",
        "artifacts": [],
        "options": {},
    }

    assert processor._should_skip_due_to_cancel(cancelled) is True
    assert processor._should_skip_due_to_cancel(cancelled) is True
    assert processor._redis_client.exists.call_count == 1
    assert processor._should_skip_due_to_cancel(normal) is False

    error_output = processor._create_cancel_error_output(cancelled)
    assert error_output["task_id"] == "cancelled-task"
    assert error_output["error"]["error_type"] == "TaskCancelledError"
    assert error_output["error"]["error_details"]["cancel_key"] == "cancel:user-42"
    assert error_output["error"]["error_details"]["skipped_before_processing"] is True
    assert error_output["__keep_raw_output__"] is True


def test_setup_redis_client_prefers_result_backend_and_supports_rediss() -> None:
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    redis_client = Mock()

    with patch("sandai.operator.processor.redis.Redis", return_value=redis_client) as redis_cls:
        channel = Mock()
        channel.celery_app.conf = _ChannelConf(
            broker_url="amqp://guest:guest@mq:5672//",
            result_backend="rediss://user:secret@redis.example.com:6380/5",
        )

        processor._setup_redis_client(channel)

    redis_cls.assert_called_once_with(
        host="redis.example.com",
        port=6380,
        username="user",
        password="secret",
        db=5,
        decode_responses=True,
        ssl=True,
    )
    redis_client.ping.assert_called_once()
    assert processor._redis_client is redis_client


def test_setup_redis_client_falls_back_from_unavailable_backend_to_broker() -> None:
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)
    failing_client = Mock()
    failing_client.ping.side_effect = RuntimeError("backend down")
    broker_client = Mock()

    with patch(
        "sandai.operator.processor.redis.Redis",
        side_effect=[failing_client, broker_client],
    ) as redis_cls:
        channel = Mock()
        channel.celery_app.conf = _ChannelConf(
            broker_url="redis://broker-host:6379/2",
            result_backend="redis://backend-host:6379/1",
        )

        processor._setup_redis_client(channel)

    assert redis_cls.call_count == 2
    broker_client.ping.assert_called_once()
    assert processor._redis_client is broker_client


def test_setup_redis_client_disables_cancellation_checks_without_redis_urls() -> None:
    processor = BatchProcessor("test-operator", "v1", ignore_nongil_check=True)

    with patch("sandai.operator.processor.redis.Redis") as redis_cls:
        channel = Mock()
        channel.celery_app.conf = _ChannelConf(
            broker_url="amqp://guest:guest@mq:5672//",
            result_backend="rpc://",
        )

        processor._setup_redis_client(channel)

    redis_cls.assert_not_called()
    assert processor._redis_client is None


def test_celery_channel_stores_error_envelope_as_success_result() -> None:
    with patch.object(CeleryChannel, "_setup_celery_consumer", lambda self: None), patch.object(
        CeleryChannel, "start_prepare_source", lambda self: None
    ), patch.object(CeleryChannel, "start_prepare_sink", lambda self: None):
        channel = CeleryChannel(
            operator_name="demo",
            operator_version="v1",
            broker_url="redis://localhost:6379/0",
            result_backend="redis://localhost:6379/0",
        )

    store_result = Mock()
    channel.celery_app.backend = Mock(store_result=store_result)

    inflight_task = Mock(client_task_id="client-task-1")
    channel._store_celery_result(
        inflight_task,
        {
            "task_id": "sdk-task-1",
            "status": "error",
            "error": {
                "error_type": "TaskCancelledError",
                "error_message": "cancelled",
                "error_details": {"cancel_key": "cancel:1"},
            },
            "artifacts": [],
            "results": {},
        },
    )

    store_result.assert_called_once_with(
        "client-task-1",
        {
            "task_id": "sdk-task-1",
            "status": "error",
            "error": {
                "error_type": "TaskCancelledError",
                "error_message": "cancelled",
                "error_details": {"cancel_key": "cancel:1"},
            },
            "artifacts": [],
            "results": {},
        },
        "SUCCESS",
    )


def test_celery_channel_prefers_explicit_connection_settings() -> None:
    with patch.object(CeleryChannel, "start_prepare_source", lambda self: None), patch.object(
        CeleryChannel, "start_prepare_sink", lambda self: None
    ), patch.object(CeleryChannel, "_debug", lambda self, message: None):
        channel = CeleryChannel(
            operator_name="demo",
            operator_version="v1",
            broker_url="redis://explicit-broker:6379/7",
            result_backend="redis://explicit-backend:6379/8",
            task_serializer="json",
            result_serializer="json",
            accept_content=["json"],
        )

    assert channel.celery_app.conf["broker_url"] == "redis://explicit-broker:6379/7"
    assert channel.celery_app.conf["result_backend"] == "redis://explicit-backend:6379/8"


def test_celery_channel_accepts_json_and_pickle_by_default() -> None:
    with patch.object(CeleryChannel, "start_prepare_source", lambda self: None), patch.object(
        CeleryChannel, "start_prepare_sink", lambda self: None
    ), patch.object(CeleryChannel, "_debug", lambda self, message: None):
        channel = CeleryChannel(
            operator_name="demo",
            operator_version="v1",
            broker_url="redis://explicit-broker:6379/7",
            result_backend="redis://explicit-backend:6379/8",
        )

    assert channel.accept_content == ["json", "msgpack"]
    assert channel.celery_app.conf["accept_content"] == ["json", "msgpack"]


def test_batch_processor_passes_max_delivery_attempts_to_celery_channel() -> None:
    processor = BatchProcessor("demo", "v1", ignore_nongil_check=True)
    processor.config["channel.type"] = "celery"
    processor.config["channel.celery"] = {
        "broker_url": "redis://explicit-broker:6379/7",
        "result_backend": "redis://explicit-backend:6379/8",
        "max_delivery_attempts": 4,
    }

    @processor.on_batch(max_batch_size=1)
    def _process_batch(batch_inputs, operator_config, context):
        if False:
            yield None

    captured_kwargs = {}

    class _FakeChannel:
        def __init__(self, **kwargs):
            captured_kwargs.update(kwargs)
            self.celery_app = Mock(conf=_ChannelConf(
                broker_url=kwargs["broker_url"],
                result_backend=kwargs["result_backend"],
            ))

        def wait_ready(self, timeout=None):
            return True

        def sync(self):
            return None

        def graceful_shutdown(self):
            return None

        def close(self):
            return None

    with patch("sandai.operator.processor.CeleryChannel", _FakeChannel), patch.object(
        processor, "_setup_redis_client", lambda channel: None
    ), patch.object(processor, "_setup_temp_directory", lambda: None), patch.object(
        processor, "stop", lambda: None
    ), patch.object(
        processor, "_run_continuous_mode", side_effect=RuntimeError("stop after create")
    ):
        import asyncio

        try:
            asyncio.run(processor._async_run(job_mode=False))
        except RuntimeError as exc:
            assert str(exc) == "stop after create"
        except Exception:
            pass

    assert captured_kwargs["max_delivery_attempts"] == 4


def test_batch_processor_reads_max_delivery_attempts_from_env() -> None:
    with patch.dict(
        "os.environ",
        {"SANDAI_OPERATOR_CELERY_MAX_DELIVERY_ATTEMPTS": "5"},
        clear=False,
    ):
        processor = BatchProcessor("demo", "v1", ignore_nongil_check=True)

    assert processor.config["channel.celery"]["max_delivery_attempts"] == 5


def test_batch_processor_init_overrides_env_max_delivery_attempts() -> None:
    with patch.dict(
        "os.environ",
        {"SANDAI_OPERATOR_CELERY_MAX_DELIVERY_ATTEMPTS": "5"},
        clear=False,
    ):
        processor = BatchProcessor(
            "demo",
            "v1",
            ignore_nongil_check=True,
            max_delivery_attempts=2,
        )

    assert processor.config["channel.celery"]["max_delivery_attempts"] == 2


def test_on_batch_overrides_max_delivery_attempts() -> None:
    processor = BatchProcessor(
        "demo",
        "v1",
        ignore_nongil_check=True,
        max_delivery_attempts=2,
    )

    @processor.on_batch(max_delivery_attempts=7)
    def _process_batch(batch_inputs, operator_config, context):
        if False:
            yield None

    assert processor.config["channel.celery"]["max_delivery_attempts"] == 7


def main() -> int:
    tests = [
        ("client_timeout_defaults", test_default_timeout_and_create_client),
        ("client_payload", test_async_call_sets_deadline_and_cancel_key),
        ("client_error_mapping", test_sync_raises_deadline_and_cancel_errors),
        ("processor_deadline", test_processor_deadline_skip_and_error_output),
        ("processor_cancel", test_processor_cancel_skip_caches_redis_and_formats_error),
        ("redis_client_prefers_backend", test_setup_redis_client_prefers_result_backend_and_supports_rediss),
        ("redis_client_fallback", test_setup_redis_client_falls_back_from_unavailable_backend_to_broker),
        ("redis_client_disabled", test_setup_redis_client_disables_cancellation_checks_without_redis_urls),
        ("celery_channel_error_envelope", test_celery_channel_stores_error_envelope_as_success_result),
        ("celery_channel_explicit_config", test_celery_channel_prefers_explicit_connection_settings),
        ("celery_channel_accept_content", test_celery_channel_accepts_json_and_pickle_by_default),
        ("celery_channel_delivery_limit_config", test_batch_processor_passes_max_delivery_attempts_to_celery_channel),
        ("delivery_limit_env", test_batch_processor_reads_max_delivery_attempts_from_env),
        ("delivery_limit_init", test_batch_processor_init_overrides_env_max_delivery_attempts),
        ("delivery_limit_on_batch", test_on_batch_overrides_max_delivery_attempts),
    ]

    passed = 0
    failed = 0

    print("Task control 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())
