#!/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_client import OperatorClient, create_client
from sandai.operator_client.config import CeleryConfig, ClientConfig
from sandai.operator_client.exceptions import (
    TaskCancelledError,
    TaskDeadlineExceededError,
)


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 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),
    ]

    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())
