"""Taskiq broker & RSTREAM producer for the job manager."""

import asyncio
import json
import logging
import os
import pickle
from typing import Any

from rstream import Producer
from rstream.exceptions import PreconditionFailed
from taskiq import Context, TaskiqDepends, TaskiqMessage
from taskiq.abc.formatter import TaskiqFormatter
from taskiq.compat import model_validate
from taskiq.message import BrokerMessage
from taskiq_aio_pika import AioPikaBroker

from digitalkin.core.common import ModuleFactory
from digitalkin.core.job_manager.base_job_manager import BaseJobManager
from digitalkin.core.task_manager.task_executor import TaskExecutor
from digitalkin.core.task_manager.task_session import TaskSession
from digitalkin.logger import logger
from digitalkin.models.module.module import ModuleCodeModel
from digitalkin.models.module.module_types import DataModel
from digitalkin.models.module.utility import EndOfStreamOutput
from digitalkin.modules._base_module import BaseModule
from digitalkin.services.services_config import ServicesConfig
from digitalkin.services.services_models import ServicesMode

logging.getLogger("taskiq").setLevel(logging.INFO)
logging.getLogger("aiormq").setLevel(logging.INFO)
logging.getLogger("aio_pika").setLevel(logging.INFO)
logging.getLogger("rstream").setLevel(logging.INFO)


class PickleFormatter(TaskiqFormatter):
    """Formatter that pickles the JSON-dumped TaskiqMessage.

    This lets you send arbitrary Python objects (classes, functions, etc.)
    by first converting to JSON-safe primitives, then pickling that string.
    """

    def dumps(self, message: TaskiqMessage) -> BrokerMessage:  # Required by TaskiqFormatter interface # noqa: PLR6301
        """Dumps message from python complex object to JSON.

        Args:
            message: TaskIQ message

        Returns:
            BrokerMessage with mandatory information for TaskIQ
        """
        payload: bytes = pickle.dumps(message)

        return BrokerMessage(
            task_id=message.task_id,
            task_name=message.task_name,
            message=payload,
            labels=message.labels,
        )

    def loads(self, message: bytes) -> TaskiqMessage:  # Required by TaskiqFormatter interface # noqa: PLR6301
        """Recreate Python object from bytes.

        Args:
            message: Broker message from bytes.

        Returns:
            message with TaskIQ format
        """
        json_str = pickle.loads(  # noqa: S301
            message
        )  # Pickle: required for Taskiq deserialization (internal broker messages only)
        return model_validate(TaskiqMessage, json_str)


def define_producer() -> Producer:
    """Get from the env the connection parameter to RabbitMQ.

    Returns:
        Producer
    """
    host: str = os.environ.get("RABBITMQ_RSTREAM_HOST", "localhost")
    port: str = os.environ.get("RABBITMQ_RSTREAM_PORT", "5552")
    username: str = os.environ.get("RABBITMQ_RSTREAM_USERNAME", "guest")
    password: str = os.environ.get("RABBITMQ_RSTREAM_PASSWORD", "guest")

    logger.info("Connection to RabbitMQ: %s:%s.", host, port)
    return Producer(host=host, port=int(port), username=username, password=password)


async def init_rstream() -> None:
    """Init a stream for every tasks."""
    try:
        await RSTREAM_PRODUCER.create_stream(
            STREAM,
            exists_ok=True,
            arguments={"max-length-bytes": STREAM_RETENTION},
        )
    except PreconditionFailed:
        logger.warning("stream already exist")


def define_broker() -> AioPikaBroker:
    """Define broker with from env paramter.

    Returns:
        Broker: connected to RabbitMQ and with custom formatter.
    """
    host: str = os.environ.get("RABBITMQ_BROKER_HOST", "localhost")
    port: str = os.environ.get("RABBITMQ_BROKER_PORT", "5672")
    username: str = os.environ.get("RABBITMQ_BROKER_USERNAME", "guest")
    password: str = os.environ.get("RABBITMQ_BROKER_PASSWORD", "guest")

    broker = AioPikaBroker(
        f"amqp://{username}:{password}@{host}:{port}",
        startup=[init_rstream],
    )
    broker.formatter = PickleFormatter()
    return broker


STREAM = "taskiq_data"
STREAM_RETENTION = 200_000
RSTREAM_PRODUCER = define_producer()
TASKIQ_BROKER = define_broker()


async def cleanup_global_resources() -> None:
    """Clean up global resources (producer and broker connections).

    This should be called during shutdown to prevent connection leaks.
    """
    try:
        await RSTREAM_PRODUCER.close()
        logger.info("RStream producer closed successfully")
    except Exception as e:
        logger.warning("Failed to close RStream producer: %s", e)

    try:
        await TASKIQ_BROKER.shutdown()
        logger.info("Taskiq broker shut down successfully")
    except Exception as e:
        logger.warning("Failed to shutdown Taskiq broker: %s", e)


async def send_message_to_stream(job_id: str, output_data: DataModel | ModuleCodeModel) -> None:
    """Callback define to add a message frame to the Rstream.

    Args:
        job_id: id of the job that sent the message
        output_data: message body as a OutputModelT or error / stream_code
    """
    body = json.dumps({"job_id": job_id, "output_data": output_data.model_dump(mode="json")}).encode("utf-8")
    await RSTREAM_PRODUCER.send(stream=STREAM, message=body)


@TASKIQ_BROKER.task
async def run_start_module(
    mission_id: str,
    setup_id: str,
    setup_version_id: str,
    module_class: type[BaseModule],
    services_mode: ServicesMode,
    input_data: dict,
    setup_data: dict,
    request_metadata: dict[str, str] | None = None,
    context: Context = TaskiqDepends(),
) -> None:
    """TaskIQ task allowing a module to compute in the background asynchronously.

    Args:
        mission_id: str,
        setup_id: The setup ID associated with the module.
        setup_version_id: The setup ID associated with the module.
        module_class: type[BaseModule],
        services_mode: ServicesMode,
        input_data: dict,
        setup_data: dict,
        request_metadata: gRPC request metadata (headers) to forward to the module.
        context: Allow TaskIQ context access
    """
    logger.info("Starting module with services_mode: %s", services_mode)
    services_config = ServicesConfig(
        services_config_strategies=module_class.services_config_strategies,
        services_config_params=module_class.services_config_params,
        mode=services_mode,
    )
    module_class.services_config = services_config
    logger.debug("Services config: %s | Module config: %s", services_config, module_class.services_config)
    module_class.discover()

    job_id = context.message.task_id
    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
    module = ModuleFactory.create_module_instance(
        module_class, job_id, mission_id, setup_id, setup_version_id, request_metadata=request_metadata
    )

    try:
        # Create TaskExecutor and supporting components for worker execution
        executor = TaskExecutor()
        session = TaskSession(job_id, mission_id, module)

        # Execute the task using TaskExecutor
        async def send_end_of_stream(_: Any) -> None:
            try:
                await callback(DataModel(root=EndOfStreamOutput()))
            except Exception as e:
                logger.error("Error sending end of stream: %s", e, exc_info=True)

        # Reconstruct Pydantic models from dicts for type safety
        try:
            input_model = module_class.create_input_model(input_data)
            setup_model = await module_class.create_setup_model(setup_data)
        except Exception as e:
            logger.error("Failed to reconstruct models for job %s: %s", job_id, e, exc_info=True)
            raise

        supervisor_task = await executor.execute_task(
            task_id=job_id,
            mission_id=mission_id,
            coro=module.start(
                input_model,
                setup_model,
                callback,
                done_callback=lambda result: asyncio.ensure_future(send_end_of_stream(result)),
            ),
            session=session,
        )

        # Wait for the supervisor task to complete
        await supervisor_task
        logger.info("Module task %s completed", job_id)
    except Exception:
        logger.exception("Error running module %s", job_id)
        raise
    finally:
        # Cleanup via module context
        try:
            await module.context.cleanup()
        except Exception:
            logger.exception("Error cleaning up module context for job %s", job_id)


@TASKIQ_BROKER.task
async def run_config_module(
    mission_id: str,
    setup_id: str,
    setup_version_id: str,
    module_class: type[BaseModule],
    services_mode: ServicesMode,
    config_setup_data: dict,
    request_metadata: dict[str, str] | None = None,
    context: Context = TaskiqDepends(),
) -> None:
    """TaskIQ task allowing a module to compute in the background asynchronously.

    Args:
        mission_id: str,
        setup_id: The setup ID associated with the module.
        setup_version_id: The setup ID associated with the module.
        module_class: type[BaseModule],
        services_mode: ServicesMode,
        config_setup_data: dict,
        request_metadata: gRPC request metadata (headers) to forward to the module.
        context: Allow TaskIQ context access
    """
    logger.info("Starting config module with services_mode: %s", services_mode)
    services_config = ServicesConfig(
        services_config_strategies=module_class.services_config_strategies,
        services_config_params=module_class.services_config_params,
        mode=services_mode,
    )
    module_class.services_config = services_config
    logger.debug("Services config: %s | Module config: %s", services_config, module_class.services_config)

    job_id = context.message.task_id
    callback = await BaseJobManager.job_specific_callback(  # type: ignore[type-var]
        send_message_to_stream, job_id
    )
    module = ModuleFactory.create_module_instance(
        module_class, job_id, mission_id, setup_id, setup_version_id, request_metadata=request_metadata
    )

    try:
        # Create TaskExecutor and supporting components for worker execution
        executor = TaskExecutor()
        session = TaskSession(job_id, mission_id, module)

        # Create and run the config setup task with TaskExecutor
        setup_model = module_class.create_config_setup_model(config_setup_data)

        supervisor_task = await executor.execute_task(
            task_id=job_id,
            mission_id=mission_id,
            coro=module.start_config_setup(setup_model, callback),
            session=session,
        )

        # Wait for the supervisor task to complete
        await supervisor_task
        logger.info("Config module task %s completed", job_id)
    except Exception:
        logger.exception("Error running config module %s", job_id)
        raise
    finally:
        # Cleanup via module context
        try:
            await module.context.cleanup()
        except Exception:
            logger.exception("Error cleaning up module context for job %s", job_id)
