#!/usr/bin/env python3
"""
Graceful-shutdown regression tests for CeleryChannel.
"""

from __future__ import annotations

import threading
import socket
import sys
from pathlib import Path
from unittest.mock import Mock, patch

from test_support import install_sdk_dependency_stubs

install_sdk_dependency_stubs()

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

from sandai.operator.channel import CeleryChannel


class MockMessage:
    def __init__(self, task_id: str):
        self.task_id = task_id
        self.rejected = False
        self.acked = False
        self.properties = {"correlation_id": task_id}
        self.delivery_info = {"routing_key": "normal.test-operator.v1"}

    def reject(self) -> None:
        self.rejected = True

    def ack(self) -> None:
        self.acked = True


class FakeRedis:
    def __init__(self):
        self.values = {}
        self.expiries = {}

    def ping(self):
        return True

    def incr(self, key: str) -> int:
        value = int(self.values.get(key, 0)) + 1
        self.values[key] = value
        return value

    def expire(self, key: str, ttl: int) -> bool:
        self.expiries[key] = ttl
        return True

    def set(self, key: str, value, nx: bool = False, ex: int | None = None):
        if nx and key in self.values:
            return False
        self.values[key] = value
        if ex is not None:
            self.expiries[key] = ex
        return True


def _create_channel(visibility_timeout: int | None = None) -> CeleryChannel:
    with patch("sandai.operator.channel.Celery"), patch(
        "sandai.operator.env.get_celery_config"
    ) as mock_get_config, patch(
        "sandai.operator.env.get_default_celery_config"
    ) as mock_default_config:
        mock_config = Mock()
        mock_config.broker_url = "redis://localhost:6379/0"
        mock_config.result_backend = "redis://localhost:6379/0"
        mock_config.task_serializer = "json"
        mock_config.result_serializer = "json"
        mock_config.accept_content = ["json", "msgpack"]
        mock_get_config.return_value = mock_config
        mock_default_config.return_value = mock_config

        kwargs = {
            "operator_name": "test-operator",
            "operator_version": "v1",
            "debug": True,
        }
        if visibility_timeout is not None:
            kwargs["visibility_timeout"] = visibility_timeout
        channel = CeleryChannel(**kwargs)
        channel._stop_event.set()
        return channel


def _register_inflight(channel: CeleryChannel, message: MockMessage, priority: str = "normal") -> None:
    routing_key = f"{priority}.test-operator.v1"
    message.delivery_info = {"routing_key": routing_key}
    inflight = channel._build_inflight_task(message, routing_key, priority)
    channel._current_process_map[inflight.sdk_task_id] = inflight


def test_visibility_timeout_defaults_to_300_seconds() -> None:
    channel = _create_channel()
    assert channel.visibility_timeout == 300


def test_graceful_shutdown_on_empty_process_map_is_safe() -> None:
    channel = _create_channel()
    channel._current_process_map = {}
    channel.graceful_shutdown()

    assert channel._stop_event.is_set() is True
    assert channel._current_process_map == {}


def test_graceful_shutdown_rejects_pending_tasks() -> None:
    channel = _create_channel(visibility_timeout=300)
    messages = [MockMessage(f"task-{index}") for index in range(3)]
    for message in messages:
        _register_inflight(channel, message)

    channel.graceful_shutdown()

    assert all(message.rejected for message in messages)
    assert channel._current_process_map == {}
    assert channel._stop_event.is_set() is True


def test_graceful_shutdown_ignores_reject_errors_and_clears_state() -> None:
    class ErrorMessage(MockMessage):
        def reject(self) -> None:
            raise Exception("reject failed")

    channel = _create_channel()
    normal = MockMessage("normal-task")
    broken = ErrorMessage("broken-task")
    _register_inflight(channel, normal)
    _register_inflight(channel, broken)

    channel.graceful_shutdown()

    assert normal.rejected is True
    assert channel._current_process_map == {}


def test_close_calls_graceful_shutdown() -> None:
    channel = _create_channel()
    message = MockMessage("close-task")
    _register_inflight(channel, message)
    channel._stop_event.clear()

    channel.close()

    assert message.rejected is True
    assert channel._stop_event.is_set() is True


def test_multiple_graceful_shutdown_calls_do_not_reject_twice() -> None:
    channel = _create_channel()
    message = MockMessage("repeat-task")
    _register_inflight(channel, message)

    channel.graceful_shutdown()
    assert message.rejected is True

    message.rejected = False
    channel.graceful_shutdown()
    assert message.rejected is False


def test_concurrent_graceful_shutdown_clears_all_tasks_once() -> None:
    channel = _create_channel()
    messages = [MockMessage(f"concurrent-{index}") for index in range(10)]
    for message in messages:
        _register_inflight(channel, message)

    threads = [threading.Thread(target=channel.graceful_shutdown) for _ in range(3)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=1.0)

    assert sum(1 for message in messages if message.rejected) == 10
    assert channel._current_process_map == {}


def test_normalize_task_payload_preserves_client_identity_in_meta() -> None:
    channel = _create_channel()
    message = MockMessage("client-task-1")
    inflight = channel._build_inflight_task(message, "high.test-operator.v1", "high")

    payload = channel._normalize_task_payload(
        {"task_id": "user-supplied", "options": {}, "artifacts": []},
        inflight,
    )

    assert payload["task_id"] == inflight.sdk_task_id
    assert payload["priority"] == "high"
    assert payload["meta"]["celery"]["client_task_id"] == "client-task-1"
    assert payload["meta"]["celery"]["routing_key"] == "high.test-operator.v1"
    assert payload["meta"]["celery"]["sdk_task_id"] == inflight.sdk_task_id


def test_store_result_uses_client_task_id_instead_of_sdk_task_id() -> None:
    channel = _create_channel()
    message = MockMessage("client-task-2")
    inflight = channel._build_inflight_task(message, "normal.test-operator.v1", "normal")
    channel.celery_app.backend = Mock()

    channel._store_celery_result(
        inflight,
        {
            "task_id": inflight.sdk_task_id,
            "status": "success",
            "results": {"ok": True},
            "artifacts": [],
        },
    )

    args = channel.celery_app.backend.store_result.call_args.args
    assert args[0] == "client-task-2"
    assert args[2] == "SUCCESS"


def test_timeout_detection_accepts_socket_and_builtin_timeouts() -> None:
    channel = _create_channel()

    assert channel._is_timeout_error(TimeoutError("timeout")) is True
    assert channel._is_timeout_error(socket.timeout("timed out")) is True
    assert channel._is_timeout_error(Exception("timed out")) is True
    assert channel._is_timeout_error(Exception("connection reset")) is False


def test_handle_incoming_message_rejects_when_prefetch_is_stopped() -> None:
    channel = _create_channel()
    message = MockMessage("blocked-task")
    channel.max_prefetch_size = 0
    channel._stop_event.set()

    channel._handle_incoming_message({"options": {}, "artifacts": []}, message)

    assert message.rejected is True
    assert channel.source_queue.empty() is True


def test_build_consumer_queues_preserves_visibility_timeout() -> None:
    channel = _create_channel(visibility_timeout=123)

    queues = channel._build_consumer_queues()

    assert [queue.name for queue in queues] == channel.queue_names
    assert all(queue.queue_arguments["x-visibility-timeout"] == 123 for queue in queues)


def test_excessive_redelivery_is_acked_and_marked_failed() -> None:
    channel = _create_channel()
    channel.max_delivery_attempts = 2
    channel._state_redis_client = FakeRedis()
    channel.celery_app.backend = Mock()

    first = MockMessage("client-task-redelivery")
    channel._handle_incoming_message({"options": {}, "artifacts": []}, first)
    assert first.acked is False
    assert first.rejected is False
    assert channel.source_queue.get(timeout=0.1)["meta"]["celery"]["delivery_attempt"] == 1

    second = MockMessage("client-task-redelivery")
    channel._handle_incoming_message({"options": {}, "artifacts": []}, second)
    assert second.acked is False
    assert second.rejected is False
    assert channel.source_queue.get(timeout=0.1)["meta"]["celery"]["delivery_attempt"] == 2

    third = MockMessage("client-task-redelivery")
    channel._handle_incoming_message({"options": {}, "artifacts": []}, third)

    assert third.acked is True
    assert third.rejected is False
    assert channel.source_queue.empty() is True

    args = channel.celery_app.backend.store_result.call_args.args
    assert args[0] == "client-task-redelivery"
    assert args[1]["status"] == "error"
    assert args[1]["error"]["error_type"] == "TaskDeliveryLimitExceededError"
    assert args[1]["error"]["error_details"]["delivery_attempt"] == 3
    assert args[1]["error"]["error_details"]["max_delivery_attempts"] == 2
    assert args[2] == "SUCCESS"


def test_late_worker_result_does_not_override_existing_terminal_result() -> None:
    channel = _create_channel()
    channel.max_delivery_attempts = 2
    channel._state_redis_client = FakeRedis()
    channel.celery_app.backend = Mock()

    first_message = MockMessage("client-task-finalized")
    inflight = channel._build_inflight_task(first_message, "normal.test-operator.v1", "normal")
    channel._state_redis_client.set(
        channel._finalized_result_key("client-task-finalized"),
        "delivery-limit-exceeded",
        nx=True,
        ex=channel._delivery_tracking_ttl(),
    )

    channel._complete_inflight_task(
        inflight,
        {
            "task_id": inflight.sdk_task_id,
            "status": "success",
            "results": {"ok": True},
            "artifacts": [],
        },
    )

    assert first_message.acked is True
    channel.celery_app.backend.store_result.assert_not_called()


def main() -> int:
    tests = [
        ("default_visibility_timeout", test_visibility_timeout_defaults_to_300_seconds),
        ("empty_process_map", test_graceful_shutdown_on_empty_process_map_is_safe),
        ("reject_pending_tasks", test_graceful_shutdown_rejects_pending_tasks),
        ("reject_error_handling", test_graceful_shutdown_ignores_reject_errors_and_clears_state),
        ("close_invokes_shutdown", test_close_calls_graceful_shutdown),
        ("idempotent_shutdown", test_multiple_graceful_shutdown_calls_do_not_reject_twice),
        ("concurrent_shutdown", test_concurrent_graceful_shutdown_clears_all_tasks_once),
        ("payload_identity_meta", test_normalize_task_payload_preserves_client_identity_in_meta),
        ("result_uses_client_task_id", test_store_result_uses_client_task_id_instead_of_sdk_task_id),
        ("timeout_detection", test_timeout_detection_accepts_socket_and_builtin_timeouts),
        ("rejects_when_stopped", test_handle_incoming_message_rejects_when_prefetch_is_stopped),
        ("consumer_queue_visibility", test_build_consumer_queues_preserves_visibility_timeout),
        ("delivery_limit_failure", test_excessive_redelivery_is_acked_and_marked_failed),
        ("late_result_dedup", test_late_worker_result_does_not_override_existing_terminal_result),
    ]

    passed = 0
    failed = 0

    print("Graceful shutdown 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())
