import asyncio
from datetime import timedelta
from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch

import bittensor as bt
import pytest
from django.utils import timezone

import abstract_block_dumper._internal.dal.django_dal as abd_dal
import abstract_block_dumper._internal.services.utils as abd_utils
from abstract_block_dumper._internal.dal.memory_registry import task_registry
from abstract_block_dumper._internal.providers.bittensor_client import BittensorConnectionClient
from abstract_block_dumper._internal.services.block_processor import block_processor_factory
from abstract_block_dumper._internal.services.executor import CeleryExecutor
from abstract_block_dumper.models import TaskAttempt
from abstract_block_dumper.v1.decorators import block_task
from tests.fatories import TaskAttemptFactory


def failing_task(block_number: int):
    raise ValueError("Task failed")


# Needs its own name: Celery's task registry is global and keyed by name, so re-registering
# `failing_task` with different celery_kwargs would return the first registration's options.
def prioritised_failing_task(block_number: int):
    raise ValueError("Task failed")


def successful_task(block_number: int):
    return {"result": "success", "block": block_number}


def flaky_function(block_number: int):
    executable_path = abd_utils.get_executable_path(flaky_function)
    task_attempt = TaskAttempt.objects.get(
        block_number=block_number,
        executable_path=executable_path,
    )

    if task_attempt.attempt_count == 0:
        raise ValueError("Flaky failure")

    return f"Success on retry for block {block_number}"


@pytest.mark.django_db
def test_task_failure_triggers_retry():
    """Test that a failing task is marked as FAILED with retry info set."""
    block_number = 100
    executable_path = abd_utils.get_executable_path(failing_task)
    task_attempt, _ = abd_dal.task_create_or_get_pending(block_number, executable_path)

    block_task(failing_task)

    registry_item = task_registry.get_by_executable_path(executable_path)

    assert registry_item is not None
    assert callable(registry_item.function)

    # Note: CELERY_TASK_ALWAYS_EAGER allows to execute it directly
    # Execute the task - it will fail but not raise (failure is recorded in DB)
    registry_item.function(block_number)
    # This task will be rescheduled more times and will fail due to max retry attempts reached

    # Reload from DB to see post-failure state
    task_attempt.refresh_from_db()
    assert task_attempt.status == TaskAttempt.Status.FAILED
    assert task_attempt.attempt_count == abd_utils.get_max_attempt_limit()
    assert abd_dal.task_can_retry(task_attempt) is False
    assert task_attempt.next_retry_at is None


@pytest.mark.django_db
def test_successful_retry_completes_task() -> None:
    """Test that retry is scheduled with correct ETA."""
    current_block = 1000
    executable_path = abd_utils.get_executable_path(flaky_function)
    task_attempt, _ = abd_dal.task_create_or_get_pending(current_block, executable_path)

    block_task(flaky_function)

    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None
    assert callable(registry_item.function)

    # Note: CELERY_TASK_ALWAYS_EAGER allows to execute it directly
    registry_item.function(current_block)

    task_attempt.refresh_from_db()

    assert task_attempt.status == TaskAttempt.Status.SUCCESS
    assert task_attempt.attempt_count == 1


@pytest.mark.django_db
def test_restry_schedules_celery_task_with_eta():
    """Test that a flaky task fails first, then succeeds on retry."""
    current_block = 100
    executable_path = abd_utils.get_executable_path(failing_task)
    task_attempt, _ = abd_dal.task_create_or_get_pending(current_block, executable_path)

    block_task(failing_task)

    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None
    assert callable(registry_item.function)

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        # Note: CELERY_TASK_ALWAYS_EAGER allows to execute it directly
        registry_item.function(current_block)

        task_attempt.refresh_from_db()
        assert task_attempt.status == TaskAttempt.Status.PENDING
        assert task_attempt.attempt_count == 1
        assert task_attempt.next_retry_at is not None

        # Verify retry was scheduled with correct parameters
        mock_apply_async.assert_called_once()
        mock_apply_async.assert_called_once_with(
            kwargs={
                "block_number": current_block,
            },
            eta=task_attempt.next_retry_at,
        )


def _register_failing_task_submitted_to(queue: str | None, block_number: int):
    """Register `failing_task` and submit `block_number` to `queue`, as a scheduler would."""
    executable_path = abd_utils.get_executable_path(failing_task)
    block_task(failing_task)

    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None

    with patch.object(registry_item.function, "apply_async"):
        CeleryExecutor().execute(registry_item, block_number, {}, queue=queue)

    return registry_item, executable_path


@pytest.mark.django_db
def test_backfill_retry_stays_on_backfill_queue():
    current_block = 100
    registry_item, executable_path = _register_failing_task_submitted_to("block-backfill", current_block)

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        registry_item.function(current_block)

        task_attempt = TaskAttempt.objects.get(
            block_number=current_block,
            executable_path=executable_path,
        )
        mock_apply_async.assert_called_once_with(
            kwargs={"block_number": current_block},
            eta=task_attempt.next_retry_at,
            queue="block-backfill",
        )


@pytest.mark.django_db
def test_recovered_backfill_retry_stays_on_backfill_queue():
    """The queue survives a scheduler restart: recovery rebuilds it from the database."""
    current_block = 100
    registry_item, executable_path = _register_failing_task_submitted_to("block-backfill", current_block)

    task_attempt = TaskAttempt.objects.get(block_number=current_block, executable_path=executable_path)
    task_attempt.status = TaskAttempt.Status.FAILED
    task_attempt.attempt_count = 1
    task_attempt.next_retry_at = timezone.now() - timedelta(minutes=5)
    task_attempt.save()

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        block_processor_factory().recover_failed_retries(poll_interval=0)

    task_attempt.refresh_from_db()
    mock_apply_async.assert_called_once_with(
        kwargs={"block_number": current_block, "_use_archive_network": False},
        eta=task_attempt.next_retry_at,
        queue="block-backfill",
    )
    assert task_attempt.celery_queue_override == "block-backfill"


@pytest.mark.django_db
def test_task_level_queue_is_not_recorded_and_survives_retries():
    """`celery_kwargs={"queue": ...}` is a Celery task-level default, not a recorded override."""
    current_block = 100
    executable_path = abd_utils.get_executable_path(prioritised_failing_task)
    block_task(celery_kwargs={"queue": "high-priority"})(prioritised_failing_task)

    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None
    # Celery applies this whenever apply_async does not pass an explicit queue.
    assert registry_item.function.queue == "high-priority"

    with patch.object(registry_item.function, "apply_async"):
        CeleryExecutor().execute(registry_item, current_block, {})

    task_attempt = TaskAttempt.objects.get(block_number=current_block, executable_path=executable_path)
    assert task_attempt.celery_queue_override is None

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        registry_item.function(current_block)

        task_attempt.refresh_from_db()
        # No explicit queue, so the retry falls through to the task-level "high-priority".
        mock_apply_async.assert_called_once_with(
            kwargs={"block_number": current_block},
            eta=task_attempt.next_retry_at,
        )
        assert task_attempt.celery_queue_override is None


@pytest.mark.django_db
def test_task_backfill_queue_overrides_live_queue_and_survives_retries():
    current_block = 100
    executable_path = abd_utils.get_executable_path(prioritised_failing_task)
    block_task(
        celery_kwargs={"queue": "high-priority"},
        backfill_queue="high-priority-backfill",
    )(prioritised_failing_task)

    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        CeleryExecutor().execute(registry_item, current_block, {}, queue=registry_item.backfill_queue)
        assert mock_apply_async.call_args.kwargs["queue"] == "high-priority-backfill"

    task_attempt = TaskAttempt.objects.get(block_number=current_block, executable_path=executable_path)
    assert task_attempt.celery_queue_override == "high-priority-backfill"

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        registry_item.function(current_block)

        task_attempt.refresh_from_db()
        mock_apply_async.assert_called_once_with(
            kwargs={"block_number": current_block},
            eta=task_attempt.next_retry_at,
            queue="high-priority-backfill",
        )
        assert task_attempt.celery_queue_override == "high-priority-backfill"


@pytest.mark.django_db
def test_live_retry_is_unaffected_by_a_task_backfill_queue():
    """A live submission records no queue, so it must not be pulled onto the backfill queue."""
    current_block = 100
    executable_path = abd_utils.get_executable_path(failing_task)
    block_task(backfill_queue="block-backfill")(failing_task)
    registry_item = task_registry.get_by_executable_path(executable_path)
    assert registry_item is not None

    with patch.object(registry_item.function, "apply_async"):
        CeleryExecutor().execute(registry_item, current_block, {})

    with patch.object(registry_item.function, "apply_async") as mock_apply_async:
        registry_item.function(current_block)

        task_attempt = TaskAttempt.objects.get(
            block_number=current_block,
            executable_path=executable_path,
        )
        mock_apply_async.assert_called_once_with(
            kwargs={"block_number": current_block},
            eta=task_attempt.next_retry_at,
        )
        assert task_attempt.celery_queue_override is None


@pytest.mark.django_db
def test_retry_recover_mechanism():
    """Test that scheduler recovers orphaned failed tasks ready for retry."""
    batch_size = 30
    executable_path = abd_utils.get_executable_path(successful_task)

    past_time = timezone.now() - timedelta(days=1)

    pending_attempts = TaskAttemptFactory.create_batch(
        size=batch_size, is_pending=True, executable_path=executable_path, next_retry_at=past_time, attempt_count=1
    )
    pending_attempt_ids = [attempt.id for attempt in pending_attempts]

    failed_attempts = TaskAttemptFactory.create_batch(
        size=batch_size, is_pending=True, executable_path=executable_path, next_retry_at=past_time, attempt_count=1
    )
    failed_attempt_ids = [attempt.id for attempt in failed_attempts]

    TaskAttemptFactory.create_batch(
        size=batch_size, is_success=True, executable_path=executable_path, next_retry_at=past_time, attempt_count=1
    )

    block_task(successful_task)

    task_registry.get_by_executable_path(executable_path)

    block_processor = block_processor_factory()
    block_processor.recover_failed_retries(poll_interval=0)  # No delay for testing

    recover_ids = pending_attempt_ids + failed_attempt_ids

    qs = TaskAttempt.objects.filter(
        id__in=recover_ids,
        status=TaskAttempt.Status.SUCCESS,
    )
    assert qs.count() == len(recover_ids)


def test_bittensor_client_uses_archive_for_old_blocks(
    mock_subtensor,
    mock_archive_subtensor,
) -> None:
    """Test that BittensorConnectionClient uses archive_subtensor for blocks older than 300 from current head."""
    current_block = 1000

    client = BittensorConnectionClient(network="testnet")
    client._subtensor = mock_subtensor
    client._archive_subtensor = mock_archive_subtensor
    client._current_block_cache = current_block

    # For old block (400 behind), should return archive_subtensor
    old_block = 600  # 400 blocks behind (> 300 threshold)
    result = client.get_subtensor_for_block(old_block)
    assert result is mock_archive_subtensor

    # For recent block (within 300), should return regular subtensor
    recent_block = 800  # 200 blocks behind (< 300 threshold)
    result = client.get_subtensor_for_block(recent_block)
    assert result is mock_subtensor

    # For block exactly at threshold (300 behind), should return regular subtensor
    threshold_block = 700  # exactly 300 blocks behind
    result = client.get_subtensor_for_block(threshold_block)
    assert result is mock_subtensor

    # For block just over threshold (301 behind), should return archive_subtensor
    just_over_threshold_block = 699  # 301 blocks behind
    result = client.get_subtensor_for_block(just_over_threshold_block)
    assert result is mock_archive_subtensor


def _mocked_rpc_substrate(
    finalized_heads: list | Exception,
    block_numbers: list | None = None,
) -> MagicMock:
    """A stand-in for ``bt.RpcSubstrate`` whose raw chain reads return the given results in order."""
    substrate = MagicMock()
    substrate.connect = AsyncMock()
    substrate.close = AsyncMock()
    substrate.raw.get_chain_finalised_head = AsyncMock(side_effect=finalized_heads)
    substrate.raw.get_block_number = AsyncMock(side_effect=block_numbers or [])
    return substrate


def _client_with_endpoint() -> BittensorConnectionClient:
    client = BittensorConnectionClient(network="test")
    client._subtensor = MagicMock(endpoint="ws://localhost:9944")
    return client


def test_bittensor_client_reads_the_latest_block_from_the_chain_head() -> None:
    subtensor = MagicMock()
    type(subtensor).block = PropertyMock(side_effect=[1000, 1001])
    client = BittensorConnectionClient(network="test")
    client._subtensor = subtensor

    assert client.get_block() == 1000
    assert client.get_block(finalized=False) == 1001


def test_bittensor_client_reads_the_finalized_block_over_rpc() -> None:
    substrate = _mocked_rpc_substrate(["0xfinalized", "0xfinalized"], [998, 999])
    client = _client_with_endpoint()

    with patch.object(bt, "RpcSubstrate", return_value=substrate) as rpc_substrate_cls:
        assert client.get_block(finalized=True) == 998
        assert client.get_block(finalized=True) == 999
        client.close()

    # One connection, reused across reads.
    rpc_substrate_cls.assert_called_once_with("ws://localhost:9944")
    substrate.connect.assert_awaited_once()
    assert substrate.raw.get_chain_finalised_head.await_count == 2
    assert substrate.raw.get_block_number.await_args_list == [
        call(block_hash="0xfinalized"),
        call(block_hash="0xfinalized"),
    ]


def test_bittensor_client_reconnects_after_a_failed_finalized_read() -> None:
    failing = _mocked_rpc_substrate(ConnectionError("socket closed"))
    healthy = _mocked_rpc_substrate(["0xfinalized"], [998])
    client = _client_with_endpoint()

    with patch.object(bt, "RpcSubstrate", side_effect=[failing, healthy]):
        with pytest.raises(ConnectionError):
            client.get_block(finalized=True)

        # The dead connection is dropped, so the next read opens a fresh one.
        assert client.get_block(finalized=True) == 998
        client.close()

    failing.close.assert_awaited_once()
    healthy.connect.assert_awaited_once()


def test_bittensor_client_times_out_a_finalized_read_that_never_answers() -> None:
    """A socket that accepts the request and goes quiet must not park the caller forever."""

    async def never_answers(*args, **kwargs):
        await asyncio.Event().wait()

    stalled = _mocked_rpc_substrate([])
    stalled.raw.get_chain_finalised_head = AsyncMock(side_effect=never_answers)
    healthy = _mocked_rpc_substrate(["0xfinalized"], [998])
    client = _client_with_endpoint()
    client.rpc_timeout = 0.05

    with patch.object(bt, "RpcSubstrate", side_effect=[stalled, healthy]):
        with pytest.raises(TimeoutError):
            client.get_block(finalized=True)

        # The stalled connection is dropped, so the next read is served by a fresh one.
        assert client.get_block(finalized=True) == 998
        client.close()

    stalled.close.assert_awaited_once()
    healthy.connect.assert_awaited_once()
