"""Task adapters and serializers for fair task queue.

This module includes:
- Task adapters for different task execution frameworks (TaskIQ, etc.)
- Task serializers for converting tasks to/from bytes
- Invocation serializers for serializing task arguments
"""

import json
import pickle
from collections.abc import Callable
from dataclasses import asdict
from typing import TYPE_CHECKING, Any, NoReturn, Protocol, TypeVar
from uuid import uuid4

if TYPE_CHECKING:
    import msgpack  # type: ignore[reportMissingImports,import-untyped,import-not-found]
else:
    try:
        import msgpack
    except ImportError:
        msgpack = None  # type: ignore[assignment]

from pydantic import BaseModel
from taskiq.decor import AsyncTaskiqDecoratedTask
from taskiq.kicker import AsyncKicker
from taskiq.message import BrokerMessage, TaskiqMessage

from haolib.background.fair.abstract import (
    InvocationSerializer,
    SerializationStrategy,
    TaskInvocation,
    TaskRegistration,
)
from haolib.background.fair.exceptions import FairQueueSerializationError

if TYPE_CHECKING:
    from taskiq import AsyncBroker

T_Task = TypeVar("T_Task")

# ============================================================================
# Task Serializers
# ============================================================================


class AbstractTaskSerializer(Protocol[T_Task]):
    """Abstract task serializer protocol.

    Serializers are responsible for converting task objects to bytes for storage
    and back to task objects. This allows for different serialization strategies
    (pickle, JSON, msgpack, etc.) to be used based on requirements.
    """

    def serialize(self, task: T_Task) -> bytes:
        """Serialize a task to bytes.

        Args:
            task: Task to serialize

        Returns:
            Serialized task data as bytes

        """
        ...

    def deserialize(self, data: bytes) -> T_Task:
        """Deserialize bytes to a task.

        Args:
            data: Serialized task data

        Returns:
            Deserialized task object

        """
        ...


class JSONTaskSerializer:
    """JSON-based task serializer.

    Uses JSON for serialization. Safe for untrusted data but requires tasks
    to be JSON-serializable (Pydantic models work well).

    Note:
        This serializer works best with Pydantic models (like TaskiqMessage)
        that support model_dump_json() and model_validate_json().

    """

    def serialize(self, task: Any) -> bytes:
        """Serialize a task using JSON.

        Args:
            task: Task to serialize (must be JSON-serializable or Pydantic model)

        Returns:
            Serialized task data as bytes

        """

        # If it's a Pydantic model, use model_dump_json
        if isinstance(task, BaseModel):
            return task.model_dump_json().encode("utf-8")

        return json.dumps(task, default=str).encode("utf-8")

    def deserialize(self, data: bytes) -> Any:
        """Deserialize bytes to a task using JSON.

        Args:
            data: Serialized task data

        Returns:
            Deserialized task object (as dict, or Pydantic model if task_class provided)

        """

        json_str = data.decode("utf-8")
        task_dict = json.loads(json_str)

        # Try to reconstruct as TaskiqMessage if possible
        try:
            return TaskiqMessage.model_validate(task_dict)
        except Exception:
            # Return as dict if validation fails
            return task_dict


class PickleTaskSerializer:
    """Pickle-based task serializer.

    Uses Python's pickle module for serialization. Fast and supports complex
    Python objects, but has security implications if used with untrusted data.

    Warning:
        Only use with trusted data sources. Pickle can execute arbitrary code
        during deserialization.

    """

    def serialize(self, task: Any) -> bytes:
        """Serialize a task using pickle.

        Args:
            task: Task to serialize

        Returns:
            Serialized task data as bytes

        """
        return pickle.dumps(task)

    def deserialize(self, data: bytes) -> Any:
        """Deserialize bytes to a task using pickle.

        Args:
            data: Serialized task data

        Returns:
            Deserialized task object

        """
        return pickle.loads(data)  # noqa: S301


# ============================================================================
# Invocation Serializers
# ============================================================================


class DefaultInvocationSerializer:
    """Default implementation of InvocationSerializer.

    Supports JSON, PICKLE, and MSGPACK serialization strategies.
    """

    def _raise_msgpack_not_installed(self, strategy: SerializationStrategy) -> NoReturn:
        """Raise error for missing msgpack installation.

        Args:
            strategy: Serialization strategy that requires msgpack

        Raises:
            FairQueueSerializationError: Always raises with msgpack not installed message

        """
        raise FairQueueSerializationError(
            detail="msgpack not installed. Install with: pip install msgpack",
            serialization_strategy=strategy.value,
        ) from None

    def _raise_unsupported_strategy(self, strategy: SerializationStrategy) -> NoReturn:
        """Raise error for unsupported serialization strategy.

        Args:
            strategy: Unsupported serialization strategy

        Raises:
            FairQueueSerializationError: Always raises with unsupported strategy message

        """
        raise FairQueueSerializationError(
            detail=f"Unsupported serialization strategy: {strategy}",
            serialization_strategy=strategy.value,
        )

    def serialize_args(self, args: tuple[Any, ...], strategy: SerializationStrategy) -> bytes:
        """Serialize args using specified strategy.

        Args:
            args: Arguments to serialize
            strategy: Serialization strategy to use

        Returns:
            Serialized bytes

        Raises:
            FairQueueSerializationError: If serialization fails or strategy is unsupported

        """
        try:
            if strategy == SerializationStrategy.JSON:
                return json.dumps(args, default=str).encode("utf-8")
            if strategy == SerializationStrategy.PICKLE:
                return pickle.dumps(args)
            if strategy == SerializationStrategy.MSGPACK:
                if msgpack is None:
                    self._raise_msgpack_not_installed(strategy)
                return msgpack.packb(args, default=str)
            self._raise_unsupported_strategy(strategy)
        except Exception as e:
            raise FairQueueSerializationError(
                detail=f"Failed to serialize args: {e}",
                serialization_strategy=strategy.value,
                serialization_error=e,
            ) from e

    def deserialize_args(self, data: bytes, strategy: SerializationStrategy) -> tuple[Any, ...]:
        """Deserialize args using specified strategy.

        Args:
            data: Serialized bytes
            strategy: Serialization strategy used

        Returns:
            Deserialized arguments

        Raises:
            FairQueueSerializationError: If deserialization fails or strategy is unsupported

        """
        try:
            if strategy == SerializationStrategy.JSON:
                return tuple(json.loads(data.decode("utf-8")))
            if strategy == SerializationStrategy.PICKLE:
                return pickle.loads(data)  # noqa: S301
            if strategy == SerializationStrategy.MSGPACK:
                if msgpack is None:
                    self._raise_msgpack_not_installed(strategy)
                return tuple(msgpack.unpackb(data, raw=False))
            self._raise_unsupported_strategy(strategy)
        except Exception as e:
            raise FairQueueSerializationError(
                detail=f"Failed to deserialize args: {e}",
                serialization_strategy=strategy.value,
                serialization_error=e,
            ) from e

    def serialize_kwargs(self, kwargs: dict[str, Any], strategy: SerializationStrategy) -> bytes:
        """Serialize kwargs using specified strategy.

        Args:
            kwargs: Keyword arguments to serialize
            strategy: Serialization strategy to use

        Returns:
            Serialized bytes

        Raises:
            FairQueueSerializationError: If serialization fails or strategy is unsupported

        """
        try:
            if strategy == SerializationStrategy.JSON:
                return json.dumps(kwargs, default=str).encode("utf-8")
            if strategy == SerializationStrategy.PICKLE:
                return pickle.dumps(kwargs)
            if strategy == SerializationStrategy.MSGPACK:
                try:
                    import msgpack  # noqa: PLC0415  # pyright: ignore[reportMissingImports]

                    return msgpack.packb(kwargs, default=str)
                except ImportError:
                    self._raise_msgpack_not_installed(strategy)
            self._raise_unsupported_strategy(strategy)
        except Exception as e:
            raise FairQueueSerializationError(
                detail=f"Failed to serialize kwargs: {e}",
                serialization_strategy=strategy.value,
                serialization_error=e,
            ) from e

    def deserialize_kwargs(self, data: bytes, strategy: SerializationStrategy) -> dict[str, Any]:
        """Deserialize kwargs using specified strategy.

        Args:
            data: Serialized bytes
            strategy: Serialization strategy used

        Returns:
            Deserialized keyword arguments

        Raises:
            FairQueueSerializationError: If deserialization fails or strategy is unsupported

        """
        try:
            if strategy == SerializationStrategy.JSON:
                return json.loads(data.decode("utf-8"))
            if strategy == SerializationStrategy.PICKLE:
                return pickle.loads(data)  # noqa: S301
            if strategy == SerializationStrategy.MSGPACK:
                try:
                    import msgpack  # noqa: PLC0415  # pyright: ignore[reportMissingImports]

                    return msgpack.unpackb(data, raw=False)
                except ImportError:
                    self._raise_msgpack_not_installed(strategy)
            else:
                self._raise_unsupported_strategy(strategy)
        except Exception as e:
            raise FairQueueSerializationError(
                detail=f"Failed to deserialize kwargs: {e}",
                serialization_strategy=strategy.value,
                serialization_error=e,
            ) from e


# ============================================================================
# Task Adapters
# ============================================================================


class ResultBackend(Protocol):
    """Protocol for TaskIQ result backends."""

    async def get_result(self, message_id: str) -> Any:
        """Get result from result backend.

        Args:
            message_id: Message ID (task ID)

        Returns:
            Result object or None if not available

        """
        ...


class TaskiqResult(Protocol):
    """Protocol for TaskIQ result objects."""

    @property
    def return_value(self) -> Any:
        """Get return value from result."""
        ...

    @property
    def result(self) -> Any:
        """Get result value."""
        ...


class TaskIQAdapter:
    """TaskIQ adapter implementing AbstractTaskAdapter protocol.

    This adapter handles TaskIQ-specific task registration and dispatch
    while presenting a framework-agnostic interface to the fair queue.
    """

    def __init__(
        self,
        broker: AsyncBroker,  # type: ignore[valid-type]
        invocation_serializer: InvocationSerializer | None = None,
        serializer: AbstractTaskSerializer[Any] | None = None,
    ) -> None:
        """Initialize TaskIQ adapter.

        Args:
            broker: TaskIQ async broker instance
            invocation_serializer: Invocation serializer (defaults to DefaultInvocationSerializer)
            serializer: Legacy task serializer (for backward compatibility)

        """
        self.broker = broker
        self._invocation_serializer = (
            invocation_serializer if invocation_serializer is not None else DefaultInvocationSerializer()
        )
        self._serializer = serializer if serializer is not None else JSONTaskSerializer()
        # Cache of registered tasks (func -> decorated_task)
        self._registered_tasks: dict[Callable[..., Any], AsyncTaskiqDecoratedTask[Any, Any]] = {}
        # Cache of registrations (func, task_name) -> TaskRegistration
        self._registrations: dict[tuple[Callable[..., Any], str | None], TaskRegistration] = {}

    async def register_task(
        self,
        func: Callable[..., Any],
        task_name: str | None = None,
        default_serialization_strategy: SerializationStrategy = SerializationStrategy.JSON,
        **registration_kwargs: Any,
    ) -> TaskRegistration:
        """Set up a TaskIQ task (register with broker).

        This method registers the function with the TaskIQ broker and returns
        a TaskRegistration that can be used later to dispatch the task.

        Args:
            func: Task function to register
            task_name: Optional custom task name
            default_serialization_strategy: Default serialization for this task
            **registration_kwargs: Additional TaskIQ-specific options (labels, etc.)

        Returns:
            TaskRegistration with task metadata

        Example:
            ```python
            async def my_task(arg1: str) -> dict:
                return {"result": arg1}

            adapter = TaskIQAdapter(broker)
            descriptor = await adapter.register_task(my_task, task_name="my_custom_task")
            # descriptor.task_name == "my_custom_task"
            ```

        """
        cache_key = (func, task_name)

        # Check if already registered
        if cache_key in self._registrations:
            return self._registrations[cache_key]

        # Register with TaskIQ broker
        # broker.task() is a decorator that can be called with a function
        # When task_name is provided, we need to use it as a decorator parameter
        # The signature is: task(task_name: str | None = None, **labels: Any)
        # But when called with a function: broker.task(func) -> decorated_task
        # We can't pass task_name directly, so we register without it and use the auto-generated name
        # Or we can use the decorator pattern, but for programmatic registration, we call it directly
        if task_name is not None:
            # Try to register with task_name using decorator pattern
            # broker.task() returns a decorator when called without arguments
            # We need to call it as: broker.task(task_name=task_name)(func)
            decorated_task = self.broker.task(task_name=task_name, **registration_kwargs)(func)  # type: ignore[call-arg,operator]
        else:
            # Register without custom task_name
            decorated_task = self.broker.task(**registration_kwargs)(func)  # type: ignore[call-arg,operator]
        self._registered_tasks[func] = decorated_task

        registration = TaskRegistration(
            task_name=decorated_task.task_name,
            registration_metadata={
                "decorated_task_name": decorated_task.task_name,
                "broker_id": id(self.broker),
            },
            default_serialization_strategy=default_serialization_strategy,
        )

        self._registrations[cache_key] = registration
        return registration

    def create_invocation(
        self,
        registration: TaskRegistration,
        *args: Any,
        task_id: str | None = None,
        serialization_strategy: SerializationStrategy | None = None,
        **kwargs: Any,
    ) -> TaskInvocation:
        """Create TaskIQ invocation descriptor.

        Creates an immutable TaskInvocation from a TaskRegistration and arguments.

        Args:
            registration: Task registration
            *args: Task function arguments
            task_id: Optional task ID (generated if not provided)
            serialization_strategy: Override default serialization strategy
            **kwargs: Task function keyword arguments

        Returns:
            Immutable TaskInvocation

        """
        if task_id is None:
            task_id = uuid4().hex

        strategy = serialization_strategy or registration.default_serialization_strategy or SerializationStrategy.JSON

        return TaskInvocation(
            task_name=registration.task_name,
            args=args,
            kwargs=kwargs,
            task_id=task_id,
            serialization_strategy=strategy,
            invocation_metadata={},
        )

    async def dispatch(self, invocation: TaskInvocation) -> None:
        """Dispatch a TaskIQ task for execution.

        Creates a TaskiqMessage from the invocation and dispatches it
        via the broker without immediate execution.

        Args:
            invocation: Task invocation descriptor

        Raises:
            ValueError: If task is not registered with broker

        """
        # Get the decorated task from broker registry
        decorated_task = self.broker.find_task(invocation.task_name)
        if decorated_task is None:
            error_msg = (
                f"Task '{invocation.task_name}' not found in broker registry. "
                f"Make sure it was registered with register_task() first."
            )
            raise ValueError(error_msg)

        # Create message using kicker's _prepare_message (no execution)
        kicker: AsyncKicker = decorated_task.kicker()
        message: TaskiqMessage = kicker._prepare_message(  # noqa: SLF001
            *invocation.args,
            **invocation.kwargs,
        )

        # Update task_id if provided in invocation
        if invocation.task_id:
            message.task_id = invocation.task_id

        # Convert to BrokerMessage and dispatch
        message_bytes = self._serializer.serialize(message)
        broker_message = BrokerMessage(
            task_id=message.task_id,
            task_name=message.task_name,
            message=message_bytes,
            labels=message.labels,
        )
        await self.broker.kick(broker_message)

    async def get_result(self, task_id: str) -> Any | None:
        """Retrieve task result from TaskIQ result backend.

        Args:
            task_id: Task ID to retrieve result for

        Returns:
            Task result object (preserves error information) or None if not available yet
            The result object may have is_err, is_error, return_value, error properties

        """
        result_backend: ResultBackend | None = None
        try:
            result_backend = self.broker.result_backend  # type: ignore[attr-defined]
        except AttributeError:
            return None
        if result_backend is None:
            return None

        try:
            result = await result_backend.get_result(task_id)
        except Exception:
            return None
        else:
            # Return the full result object to preserve error information
            # The scheduler will check for is_err/is_error before extracting return_value
            return result

    def serialize_invocation(self, invocation: TaskInvocation) -> bytes:
        """Serialize task invocation for storage.

        Args:
            invocation: Task invocation

        Returns:
            Serialized bytes

        """
        # Use Pydantic's serialization for TaskInvocation
        # Note: We need to convert the dataclass to a dict first
        invocation_dict = asdict(invocation)
        # Convert SerializationStrategy enum to string
        invocation_dict["serialization_strategy"] = invocation.serialization_strategy.value
        return json.dumps(invocation_dict, default=str).encode("utf-8")

    def deserialize_invocation(self, data: bytes) -> TaskInvocation:
        """Deserialize task invocation from storage.

        Args:
            data: Serialized bytes

        Returns:
            Deserialized TaskInvocation

        """
        invocation_dict = json.loads(data.decode("utf-8"))
        # Convert serialization_strategy string back to enum
        if isinstance(invocation_dict.get("serialization_strategy"), str):
            invocation_dict["serialization_strategy"] = SerializationStrategy(invocation_dict["serialization_strategy"])
        # Convert args from list to tuple
        if "args" in invocation_dict and isinstance(invocation_dict["args"], list):
            invocation_dict["args"] = tuple(invocation_dict["args"])
        return TaskInvocation(**invocation_dict)

    def message_to_invocation(self, message: TaskiqMessage) -> TaskInvocation:
        """Convert TaskiqMessage to TaskInvocation.

        This method allows using TaskIQ's declarative pattern (task.kiq())
        with the fair queue by converting the message to an invocation.

        Args:
            message: TaskiqMessage from task.kiq()

        Returns:
            TaskInvocation for use with fair queue

        Example:
            ```python
            # Declarative pattern
            @broker.task
            async def my_task(arg1: str):
                return result

            # Create message using kiq()
            message = await my_task.kiq(arg1="value")

            # Convert to invocation for fair queue
            adapter = TaskIQAdapter(broker)
            invocation = adapter.message_to_invocation(message)

            # Enqueue in fair queue
            await queue.enqueue(invocation, group_id="user_123")
            ```

        """
        return TaskInvocation(
            task_name=message.task_name,
            args=tuple(message.args),
            kwargs=message.kwargs,
            task_id=message.task_id,
            serialization_strategy=SerializationStrategy.JSON,  # Default for TaskIQ messages
            invocation_metadata={},
        )
