#!/usr/bin/env python3
"""
Regression tests for the Celery watchdog feature.

All tests use mocks — no real Celery or Redis required.
"""

from __future__ import annotations

import asyncio
import os
import sys
import time
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, 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-sdk"))

from sandai.operator import BatchProcessor
from sandai.operator.env import get_celery_watchdog_timeout


# ---------------------------------------------------------------------------
# 1. env defaults
# ---------------------------------------------------------------------------

def test_watchdog_timeout_env_default() -> None:
    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("SANDAI_OPERATOR_CELERY_WATCHDOG_TIMEOUT", None)
        assert get_celery_watchdog_timeout() == 0


def test_watchdog_timeout_env_override() -> None:
    with patch.dict(os.environ, {"SANDAI_OPERATOR_CELERY_WATCHDOG_TIMEOUT": "42"}):
        assert get_celery_watchdog_timeout() == 42


# ---------------------------------------------------------------------------
# 2. _celery_watchdog_loop — queue empty → no restart
# ---------------------------------------------------------------------------

def test_watchdog_not_triggered_when_queue_empty() -> None:
    processor = BatchProcessor("test", "v1", ignore_nongil_check=True)
    processor._last_compute_consumed_at = time.monotonic() - 9999  # very old

    channel = Mock()
    channel.get_pending_task_count.return_value = 0  # queues empty

    exit_calls = []

    async def _run():
        processor._stop_event.clear()
        task = asyncio.create_task(
            processor._celery_watchdog_loop(channel, timeout_secs=1)
        )
        await asyncio.sleep(0.2)
        processor._stop_event.set()
        await asyncio.sleep(0.05)
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

    with patch("os._exit", side_effect=exit_calls.append):
        asyncio.run(_run())

    assert exit_calls == [], f"os._exit should NOT be called when queues are empty, got {exit_calls}"


# ---------------------------------------------------------------------------
# 3. _celery_watchdog_loop — elapsed >= timeout AND pending > 0 → restart
# ---------------------------------------------------------------------------

def test_watchdog_triggers_restart_when_stuck() -> None:
    processor = BatchProcessor("test", "v1", ignore_nongil_check=True)
    processor._last_compute_consumed_at = time.monotonic() - 9999  # very old

    channel = Mock()
    channel.get_pending_task_count.return_value = 5  # tasks pending

    exit_called_with = []

    # Save the real asyncio.sleep before patching.
    _real_sleep = asyncio.sleep

    async def instant_sleep(_secs):
        """Yield to the event loop without actually waiting."""
        await _real_sleep(0)

    def _os_exit_side_effect(code):
        exit_called_with.append(code)
        # Stop the watchdog loop so the task exits cleanly.
        processor._stop_event.set()

    async def _run_fast():
        processor._stop_event.clear()
        import sandai.operator.processor as _proc_mod
        original = _proc_mod.asyncio.sleep
        _proc_mod.asyncio.sleep = instant_sleep
        try:
            with patch("os._exit", side_effect=_os_exit_side_effect):
                task = asyncio.create_task(
                    processor._celery_watchdog_loop(channel, timeout_secs=1)
                )
                for _ in range(50):
                    await _real_sleep(0)
                    if exit_called_with:
                        break
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
        finally:
            _proc_mod.asyncio.sleep = original

    asyncio.run(_run_fast())
    assert exit_called_with == [1], f"Expected os._exit(1), got {exit_called_with}"


# ---------------------------------------------------------------------------
# 4. _execute_batch updates _last_compute_consumed_at
# ---------------------------------------------------------------------------

def test_watchdog_reset_on_task_consumption() -> None:
    processor = BatchProcessor("test", "v1", ignore_nongil_check=True)
    processor._last_compute_consumed_at = 0.0  # sentinel

    # Minimal PreparedTask stub
    class FakeTask:
        task_input = Mock()
        task_input.task_id = "t1"
        persistent_output = None

    called = []

    async def fake_process_batch(inputs):
        called.append(inputs)
        return [Mock() for _ in inputs]

    processor.process_batch = fake_process_batch

    before = time.monotonic()

    async def _run():
        result = await processor._execute_batch([FakeTask()])
        return result

    asyncio.run(_run())

    assert processor._last_compute_consumed_at >= before, (
        "_last_compute_consumed_at should be updated after _execute_batch"
    )
    assert called, "process_batch should have been called"


# ---------------------------------------------------------------------------
# 5. get_pending_task_count — mock Redis llen sums all queues
# ---------------------------------------------------------------------------

def test_get_pending_task_count_sums_all_queues() -> None:
    from sandai.operator.channel import CeleryChannel

    with patch("sandai.operator.channel.CeleryChannel._setup_celery_consumer"), \
         patch("sandai.operator.channel.CeleryChannel._setup_state_redis_client"), \
         patch("sandai.operator.channel.CeleryChannel.start_prepare_source"), \
         patch("sandai.operator.channel.CeleryChannel.start_prepare_sink"):

        ch = CeleryChannel.__new__(CeleryChannel)
        ch.operator_name = "op"
        ch.operator_version = "v1"
        ch.priorities = ["high", "normal"]
        ch.queue_names = ["high.op.v1", "normal.op.v1"]
        ch.debug = False
        ch._watchdog_redis_client = None
        ch.celery_app = Mock()
        ch.celery_app.conf = _ChannelConf(broker_url="redis://localhost:6379/0")

        mock_redis = Mock()
        mock_redis.ping.return_value = True
        mock_redis.llen.side_effect = lambda q: {"high.op.v1": 3, "normal.op.v1": 7}.get(q, 0)

        with patch("redis.Redis", return_value=mock_redis):
            count = ch.get_pending_task_count()

    assert count == 10, f"Expected 10, got {count}"


# ---------------------------------------------------------------------------
# 6. get_pending_task_count — Redis unavailable → returns 0, no exception
# ---------------------------------------------------------------------------

def test_get_pending_task_count_returns_zero_without_redis() -> None:
    from sandai.operator.channel import CeleryChannel

    with patch("sandai.operator.channel.CeleryChannel._setup_celery_consumer"), \
         patch("sandai.operator.channel.CeleryChannel._setup_state_redis_client"), \
         patch("sandai.operator.channel.CeleryChannel.start_prepare_source"), \
         patch("sandai.operator.channel.CeleryChannel.start_prepare_sink"):

        ch = CeleryChannel.__new__(CeleryChannel)
        ch.operator_name = "op"
        ch.operator_version = "v1"
        ch.priorities = ["high"]
        ch.queue_names = ["high.op.v1"]
        ch.debug = False
        ch._watchdog_redis_client = None
        ch.celery_app = Mock()
        ch.celery_app.conf = _ChannelConf(broker_url="redis://localhost:6379/0")

        failing_redis = Mock()
        failing_redis.ping.side_effect = ConnectionError("unreachable")

        with patch("redis.Redis", return_value=failing_redis):
            count = ch.get_pending_task_count()

    assert count == 0


# ---------------------------------------------------------------------------
# 7. BatchProcessor default — watchdog_timeout == 0
# ---------------------------------------------------------------------------

def test_watchdog_disabled_by_default_in_processor() -> None:
    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("SANDAI_OPERATOR_CELERY_WATCHDOG_TIMEOUT", None)
        processor = BatchProcessor("test", "v1", ignore_nongil_check=True)

    assert processor.config["channel.celery"]["watchdog_timeout"] == 0


# ---------------------------------------------------------------------------
# 8. @on_batch override
# ---------------------------------------------------------------------------

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

    @processor.on_batch(celery_watchdog_timeout=60)
    def process(batch_inputs, context):
        pass

    assert processor.config["channel.celery"]["watchdog_timeout"] == 60


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

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

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


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    tests = [
        test_watchdog_timeout_env_default,
        test_watchdog_timeout_env_override,
        test_watchdog_not_triggered_when_queue_empty,
        test_watchdog_triggers_restart_when_stuck,
        test_watchdog_reset_on_task_consumption,
        test_get_pending_task_count_sums_all_queues,
        test_get_pending_task_count_returns_zero_without_redis,
        test_watchdog_disabled_by_default_in_processor,
        test_watchdog_on_batch_override,
    ]

    passed = 0
    failed = 0
    for fn in tests:
        try:
            fn()
            print(f"  PASS  {fn.__name__}")
            passed += 1
        except Exception as exc:
            import traceback
            print(f"  FAIL  {fn.__name__}: {exc}")
            traceback.print_exc()
            failed += 1

    print(f"\n{'=' * 50}")
    print(f"Celery Watchdog tests: passed={passed} failed={failed} total={passed+failed}")
    if failed:
        sys.exit(1)
