from __future__ import annotations

import asyncio
from collections.abc import Callable
from dataclasses import dataclass

from aiogram import Bot, Dispatcher

from src.bootstrap.market_data_storage import (
    MarketDataStorageLifecycleProtocol,
)
from src.market_data.acquisition.runtime.trade_stream_production_runtime import (
    TradeStreamProductionRuntimeProtocol,
)


@dataclass(frozen=True, slots=True)
class ApplicationComposition:
    """Корневые компоненты одного запущенного экземпляра приложения."""

    bot: Bot
    dispatcher: Dispatcher
    trade_stream_runtime: (
        TradeStreamProductionRuntimeProtocol | None
    )
    market_data_storage_lifecycle: (
        MarketDataStorageLifecycleProtocol | None
    ) = None


async def run_application(
    application: ApplicationComposition,
) -> None:
    """
    Выполнять Telegram polling и опциональный Trade Stream как одно целое.

    Ошибка включённого Trade Stream считается фатальной для всего
    приложения. Остановка Telegram, Trade Stream и bot session
    выполняется одним владельцем и не оставляет фоновых root tasks.
    """
    polling_task: asyncio.Task[None] | None = None
    runtime_task: asyncio.Task[None] | None = None
    primary_error: BaseException | None = None

    try:
        storage_lifecycle = (
            application.market_data_storage_lifecycle
        )

        if storage_lifecycle is not None:
            await _run_blocking_lifecycle_operation(
                storage_lifecycle.start,
                task_name="market-data-storage-startup",
            )

        polling_task = asyncio.create_task(
            application.dispatcher.start_polling(
                application.bot,
                close_bot_session=False,
            ),
            name="telegram-polling",
        )
        runtime_task = (
            asyncio.create_task(
                application.trade_stream_runtime.run(),
                name="trade-stream-runtime",
            )
            if application.trade_stream_runtime is not None
            else None
        )

        await _wait_for_root_tasks(
            polling_task=polling_task,
            runtime_task=runtime_task,
        )
    except BaseException as error:
        primary_error = error

    cleanup_task = asyncio.create_task(
        _shutdown_application(
            application=application,
            polling_task=polling_task,
            runtime_task=runtime_task,
            primary_error=primary_error,
        ),
        name="application-shutdown",
    )
    cleanup_error = await _await_cleanup(
        cleanup_task,
        primary_error=primary_error,
    )

    if primary_error is not None:
        if cleanup_error is not None:
            primary_error.add_note(
                "Application cleanup also failed: "
                f"{type(cleanup_error).__name__}."
            )

        raise primary_error.with_traceback(
            primary_error.__traceback__,
        )

    if cleanup_error is not None:
        raise cleanup_error.with_traceback(
            cleanup_error.__traceback__,
        )


async def _wait_for_root_tasks(
    *,
    polling_task: asyncio.Task[None],
    runtime_task: asyncio.Task[None] | None,
) -> None:
    if runtime_task is None:
        await polling_task
        return

    done_tasks, _ = await asyncio.wait(
        (
            polling_task,
            runtime_task,
        ),
        return_when=asyncio.FIRST_COMPLETED,
    )

    if runtime_task in done_tasks:
        if runtime_task.cancelled():
            raise RuntimeError(
                "Trade Stream Runtime was cancelled unexpectedly."
            )

        runtime_error = runtime_task.exception()

        if runtime_error is not None:
            raise runtime_error.with_traceback(
                runtime_error.__traceback__,
            )

        raise RuntimeError(
            "Trade Stream Runtime terminated unexpectedly."
        )

    await polling_task


async def _shutdown_application(
    *,
    application: ApplicationComposition,
    polling_task: asyncio.Task[None] | None,
    runtime_task: asyncio.Task[None] | None,
    primary_error: BaseException | None,
) -> BaseException | None:
    cleanup_error: BaseException | None = None
    polling_cancelled = False
    runtime_cancelled = False

    if polling_task is not None and not polling_task.done():
        polling_cancelled = True
        polling_task.cancel()

    runtime = application.trade_stream_runtime

    if runtime is not None:
        try:
            await runtime.stop()
        except BaseException as error:
            cleanup_error = _merge_cleanup_error(
                cleanup_error,
                error,
            )

    if runtime_task is not None and not runtime_task.done():
        runtime_cancelled = True
        runtime_task.cancel()

    if polling_task is not None:
        cleanup_error = await _observe_root_task(
            polling_task,
            expected_cancellation=polling_cancelled,
            primary_error=primary_error,
            cleanup_error=cleanup_error,
        )

    if runtime_task is not None:
        cleanup_error = await _observe_root_task(
            runtime_task,
            expected_cancellation=runtime_cancelled,
            primary_error=primary_error,
            cleanup_error=cleanup_error,
        )

    storage_lifecycle = application.market_data_storage_lifecycle

    if storage_lifecycle is not None:
        try:
            await _run_blocking_lifecycle_operation(
                storage_lifecycle.stop,
                task_name="market-data-storage-shutdown",
            )
        except BaseException as error:
            cleanup_error = _merge_cleanup_error(
                cleanup_error,
                error,
            )

    try:
        await application.bot.session.close()
    except BaseException as error:
        cleanup_error = _merge_cleanup_error(
            cleanup_error,
            error,
        )

    return cleanup_error


async def _run_blocking_lifecycle_operation(
    operation: Callable[[], None],
    *,
    task_name: str,
) -> None:
    operation_task = asyncio.create_task(
        asyncio.to_thread(operation),
        name=task_name,
    )

    try:
        await asyncio.shield(operation_task)
    except asyncio.CancelledError as cancellation:
        try:
            await _wait_for_blocking_operation(operation_task)
        except BaseException as operation_error:
            cancellation.add_note(
                "Application blocking lifecycle operation also failed: "
                f"{type(operation_error).__name__}."
            )

        raise


async def _wait_for_blocking_operation(
    operation_task: asyncio.Task[None],
) -> None:
    while not operation_task.done():
        try:
            await asyncio.shield(operation_task)
        except asyncio.CancelledError:
            continue
        except BaseException:
            break

    operation_task.result()


async def _observe_root_task(
    task: asyncio.Task[None],
    *,
    expected_cancellation: bool,
    primary_error: BaseException | None,
    cleanup_error: BaseException | None,
) -> BaseException | None:
    try:
        await task
    except asyncio.CancelledError as error:
        if (
            not expected_cancellation
            and error is not primary_error
            and not isinstance(primary_error, asyncio.CancelledError)
        ):
            return _merge_cleanup_error(
                cleanup_error,
                error,
            )
    except BaseException as error:
        if error is not primary_error:
            return _merge_cleanup_error(
                cleanup_error,
                error,
            )

    return cleanup_error


async def _await_cleanup(
    cleanup_task: asyncio.Task[BaseException | None],
    *,
    primary_error: BaseException | None,
) -> BaseException | None:
    interrupted_error: asyncio.CancelledError | None = None

    while not cleanup_task.done():
        try:
            await asyncio.shield(cleanup_task)
        except asyncio.CancelledError as error:
            if primary_error is None:
                primary_error = error
                interrupted_error = error
            continue

    cleanup_error = cleanup_task.result()

    if interrupted_error is not None:
        if cleanup_error is not None:
            interrupted_error.add_note(
                "Application cleanup also failed: "
                f"{type(cleanup_error).__name__}."
            )

        raise interrupted_error

    return cleanup_error


def _merge_cleanup_error(
    current_error: BaseException | None,
    new_error: BaseException,
) -> BaseException:
    if current_error is None:
        return new_error

    current_error.add_note(
        "Additional application cleanup failure: "
        f"{type(new_error).__name__}."
    )
    return current_error