import asyncio
import datetime
import json
from dataclasses import dataclass
from typing import NewType

import pytest
from pydantic import BaseModel

from kurier import BusType, Station
from kurier.core.basekurier import BaseKurier
from kurier.core.kurier_message import KurierMessage
from kurier.hermes import Hermes
from kurier.transports.inmemory.bus import InMemoryLine

Hermes.is_enabled = False


# Test data models
@dataclass
class TestDataClass:
    name: str
    age: int
    created_at: datetime.datetime


class TestPydanticModel(BaseModel):
    name: str
    age: int
    created_at: datetime.datetime


UserId = NewType("UserId", int)


class TestStationValidation:
    """Tests for Station handler validation."""

    def setup_method(self):
        """Clear the registry before each test."""
        Station.station_registry.clear()
        BaseKurier._registry.clear()
        InMemoryLine._registry.clear()

    def test_function_have_only_one_arg(self):
        """
        Tests that a function with exactly one argument can be successfully
        decorated as a handler using `Station.add_handler`.
        """

        @Station.add_handler(BusType.IN_MEMORY)
        async def f(data: str, *context) -> str:
            return data

        # Verify handler was registered
        assert BusType.IN_MEMORY in Station.station_registry
        assert "str" in Station.station_registry[BusType.IN_MEMORY]

    def test_function_have_more_than_one_arg_raises_error(self):
        """
        Tests that `Station.add_handler` raises a ValueError when attempting
        to decorate a function that accepts more than one argument.
        """
        with pytest.raises(ValueError):

            @Station.add_handler(BusType.IN_MEMORY)
            async def f(data: str, data2: str) -> tuple[str, str]:
                return data, data2

    def test_function_have_no_arg_raises_error(self):
        """
        Tests that `Station.add_handler` raises a ValueError when attempting
        to decorate a function that accepts no arguments.
        """
        with pytest.raises(ValueError):

            @Station.add_handler(BusType.IN_MEMORY)
            async def f():
                pass

    def test_sync_function_registration(self):
        """Test that synchronous functions can be registered."""

        @Station.add_handler(BusType.IN_MEMORY)
        def sync_handler(data: int, *context):
            return data * 2

        handlers = Station.station_registry[BusType.IN_MEMORY]["int"]
        assert len(handlers) == 1
        assert not handlers[0].is_async

    def test_async_function_registration(self):
        """Test that asynchronous functions can be registered."""

        @Station.add_handler(BusType.IN_MEMORY)
        async def async_handler(data: int, *context):
            return data * 2

        handlers = Station.station_registry[BusType.IN_MEMORY]["int"]
        assert len(handlers) == 1
        assert handlers[0].is_async

    def test_multiple_handlers_same_type(self):
        """Test registering multiple handlers for the same message type."""

        @Station.add_handler(BusType.IN_MEMORY)
        def handler1(data: str, *context):
            pass

        @Station.add_handler(BusType.IN_MEMORY)
        def handler2(data: str, *context):
            pass

        handlers = Station.station_registry[BusType.IN_MEMORY]["str"]
        assert len(handlers) == 2

    def test_complex_type_registration(self):
        """Test registering handlers for complex types."""

        @Station.add_handler(BusType.IN_MEMORY)
        def dataclass_handler(data: TestDataClass, *context):
            pass

        @Station.add_handler(BusType.IN_MEMORY)
        def pydantic_handler(data: TestPydanticModel, *context):
            pass

        assert "TestDataClass" in Station.station_registry[BusType.IN_MEMORY]
        assert "TestPydanticModel" in Station.station_registry[BusType.IN_MEMORY]


class TestStationDispatch:
    """Tests for Station passenger dispatching."""

    def setup_method(self):
        """Clear the registry before each test."""
        Station.station_registry.clear()
        BaseKurier._registry.clear()
        InMemoryLine._registry.clear()

    @pytest.fixture
    def line(self):
        """Create a fresh InMemoryLine instance for dispatch tests."""
        line = InMemoryLine()
        yield line
        asyncio.run(line.close())

    @pytest.mark.asyncio
    async def test_dispatch_to_sync_handler(self, line):
        """Test dispatching messages to synchronous handlers."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: str, *context):
            received_messages.append(data)

        envelope = KurierMessage(
            object_type_name="str",
            message_payload='"test_message"',
            metadata={},
        )

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_messages) == 1
        assert received_messages[0] == "test_message"

    @pytest.mark.asyncio
    async def test_dispatch_to_async_handler(self, line):
        """Test dispatching messages to asynchronous handlers."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        async def handler(data: str, *context):
            received_messages.append(data)

        envelope = KurierMessage(object_type_name="str", message_payload='"async_test"', metadata={})

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_messages) == 1
        assert received_messages[0] == "async_test"

    @pytest.mark.asyncio
    async def test_dispatch_dataclass_message(self, line):
        """Test dispatching dataclass messages."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: TestDataClass, *context):
            received_messages.append(data)

        envelope = KurierMessage(
            object_type_name="TestDataClass",
            message_payload=json.dumps(
                {
                    "name": "John",
                    "age": 30,
                    "created_at": "2023-01-01T10:00:00",
                }
            ),
            metadata={},
        )

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_messages) == 1
        assert received_messages[0].name == "John"
        assert received_messages[0].age == 30

    @pytest.mark.asyncio
    async def test_dispatch_pydantic_message(self, line):
        """Test dispatching Pydantic model messages."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: TestPydanticModel, *context):
            received_messages.append(data)

        envelope = KurierMessage(
            object_type_name="TestPydanticModel",
            message_payload='{"name": "Jane", "age": 25, "created_at": "2023-01-01T10:00:00"}',
            metadata={},
        )

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_messages) == 1
        assert received_messages[0].name == "Jane"
        assert received_messages[0].age == 25

    @pytest.mark.asyncio
    async def test_dispatch_multiple_handlers(self, line):
        """Test dispatching to multiple handlers for the same message type."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def handler1(data: str, *context):
            received_messages.append(f"handler1: {data}")

        @Station.add_handler(BusType.IN_MEMORY)
        def handler2(data: str, *context):
            received_messages.append(f"handler2: {data}")

        envelope = KurierMessage(object_type_name="str", message_payload='"multi_test"', metadata={})

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_messages) == 2
        assert "handler1: multi_test" in received_messages
        assert "handler2: multi_test" in received_messages

    @pytest.mark.asyncio
    async def test_dispatch_unknown_message_type(self, line):
        """Test dispatching unknown message types."""

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: str, *context):
            pass

        envelope = KurierMessage(
            object_type_name="UnknownType",
            message_payload='"test"',
            metadata={},
        )

        result = await Station.dispatch(line, envelope.model_dump_json())
        assert result is False

    @pytest.mark.asyncio
    async def test_dispatch_malformed_envelope(self, line):
        """Test dispatching with malformed envelope."""

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: str, *context):
            pass

        malformed_json = '{"invalid": "envelope"}'
        result = await Station.dispatch(line, malformed_json)
        assert result is False

    @pytest.mark.asyncio
    async def test_dispatch_handler_exception(self, line):
        """Test that exceptions in message handlers propagate to the caller."""

        @Station.add_handler(BusType.IN_MEMORY)
        def failing_handler(data: str, *context):
            raise Exception("Handler failed")

        envelope = KurierMessage(object_type_name="str", message_payload='"test"', metadata={})

        with pytest.raises(Exception, match="Handler failed"):
            await Station.dispatch(line, envelope.model_dump_json())

    @pytest.mark.asyncio
    async def test_dispatch_with_metadata(self, line):
        """Test that context metadata is passed to handlers."""
        received_context = []

        @Station.add_handler(BusType.IN_MEMORY)
        def handler(data: str, *context):
            received_context.extend(context)

        envelope = KurierMessage(
            object_type_name="str",
            message_payload='"test"',
            metadata={"source": "test", "priority": 1},
        )

        result = await Station.dispatch(line, envelope.model_dump_json())

        assert result is True
        assert len(received_context) == 1
        context_data = received_context[0]
        assert context_data["metadata"]["source"] == "test"
        assert context_data["metadata"]["priority"] == 1


async def _poll(predicate, timeout=2.0, interval=0.05):
    """Poll *predicate* until it returns True or *timeout* seconds elapse."""
    elapsed = 0.0
    while elapsed < timeout:
        if predicate():
            return
        await asyncio.sleep(interval)
        elapsed += interval
    raise AssertionError(f"Condition not met within {timeout}s")


class TestStationIntegration:
    """Integration tests for Station with InMemoryLine."""

    def setup_method(self):
        """Clear the registry before each test."""
        Station.station_registry.clear()
        BaseKurier._registry.clear()
        InMemoryLine._registry.clear()

    @pytest.mark.asyncio
    async def test_end_to_end_message_flow(self):
        """Test complete message flow from publish to handler."""
        received_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def message_handler(data: str, *context):
            received_messages.append(data)

        line = InMemoryLine.get("")

        await line.publish_async("integration_test")

        await _poll(lambda: len(received_messages) >= 1)

        assert len(received_messages) == 1
        assert received_messages[0] == "integration_test"

        await line.close()

    @pytest.mark.asyncio
    async def test_multiple_message_types_integration(self):
        """Test handling multiple different message types in integration."""
        string_messages = []
        int_messages = []
        dataclass_messages = []

        @Station.add_handler(BusType.IN_MEMORY)
        def string_handler(data: str, *context):
            string_messages.append(data)

        @Station.add_handler(BusType.IN_MEMORY)
        def int_handler(data: int, *context):
            int_messages.append(data)

        @Station.add_handler(BusType.IN_MEMORY)
        def dataclass_handler(data: TestDataClass, *context):
            dataclass_messages.append(data)

        line = InMemoryLine.get("")

        await line.publish_async("string message")
        await line.publish_async(123)
        await line.publish_async(TestDataClass(name="Test", age=25, created_at=datetime.datetime.now()))

        await _poll(lambda: len(string_messages) >= 1 and len(int_messages) >= 1 and len(dataclass_messages) >= 1)

        assert len(string_messages) == 1
        assert string_messages[0] == "string message"

        assert len(int_messages) == 1
        assert int_messages[0] == 123

        assert len(dataclass_messages) == 1
        assert dataclass_messages[0].name == "Test"
        assert dataclass_messages[0].age == 25

        await line.close()

    @pytest.mark.asyncio
    async def test_batch_processing_integration(self):
        """Test batch message processing."""
        received_messages = []
        correlation_ids = set()

        @Station.add_handler(BusType.IN_MEMORY)
        def batch_handler(data: str, context):
            received_messages.append(data)
            correlation_ids.add(context["correlation_id"])

        line = InMemoryLine.get("")

        messages = ["msg1", "msg2", "msg3", "msg4", "msg5"]
        await line.publish_batch_async(messages, {}, batch_size=2)

        await _poll(lambda: len(received_messages) >= len(messages))

        assert received_messages == messages
        assert len(correlation_ids) == 1

        await line.close()
