from __future__ import annotations

import asyncio
import logging
import uuid
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterator,
    Final,
    Iterable,
    Mapping,
    Optional,
    Set,
    Tuple,
)

import aiohttp_cors
import attrs
import sqlalchemy as sa
import trafaret as t
from aiohttp import web
from aiohttp_sse import sse_response
from aiotools import adefer
from sqlalchemy.orm import load_only

from ai.backend.common import validators as tx
from ai.backend.common.events.dispatcher import (
    EventDispatcher,
)
from ai.backend.common.events.event_types.kernel.broadcast import (
    KernelCancelledBroadcastEvent,
    KernelCreatingBroadcastEvent,
    KernelPreparingBroadcastEvent,
    KernelPullingBroadcastEvent,
    KernelStartedBroadcastEvent,
    KernelTerminatedBroadcastEvent,
    KernelTerminatingBroadcastEvent,
)
from ai.backend.common.events.event_types.session.broadcast import (
    SessionCancelledBroadcastEvent,
    SessionEnqueuedBroadcastEvent,
    SessionFailureBroadcastEvent,
    SessionScheduledBroadcastEvent,
    SessionStartedBroadcastEvent,
    SessionSuccessBroadcastEvent,
    SessionTerminatedBroadcastEvent,
    SessionTerminatingBroadcastEvent,
)
from ai.backend.common.events.hub.propagators.cache import WithCachePropagator
from ai.backend.common.events.types import EventCacheDomain, EventDomain
from ai.backend.common.json import dump_json_str
from ai.backend.common.types import AgentId
from ai.backend.logging import BraceStyleAdapter

from ..errors.exceptions import GenericForbidden, GroupNotFound
from ..models import UserRole, groups, kernels
from ..models.session import SessionRow
from ..models.utils import execute_with_retry
from ..types import Sentinel
from .auth import auth_required
from .manager import READ_ALLOWED, server_status_required
from .utils import check_api_params

if TYPE_CHECKING:
    from .context import RootContext
    from .types import CORSOptions, WebMiddleware

log = BraceStyleAdapter(logging.getLogger(__spec__.name))

sentinel: Final = Sentinel.token

SessionEventInfo = Tuple[str, dict, str, Optional[int]]


@server_status_required(READ_ALLOWED)
@auth_required
@check_api_params(
    t.Dict({
        tx.AliasedKey(["name", "sessionName"], default="*") >> "session_name": t.String,
        t.Key("ownerAccessKey", default=None) >> "owner_access_key": t.Null | t.String,
        t.Key("sessionId", default=None) >> "session_id": t.Null | tx.UUID,
        # NOTE: if set, sessionId overrides sessionName and ownerAccessKey parameters.
        tx.AliasedKey(["group", "groupName"], default="*") >> "group_name": t.String,
        t.Key("scope", default="*"): t.Enum("*", "session", "kernel"),
    })
)
@adefer
async def push_session_events(
    defer,
    request: web.Request,
    params: Mapping[str, Any],
) -> web.StreamResponse:
    root_ctx: RootContext = request.app["_root.context"]
    app_ctx: PrivateContext = request.app["events.context"]
    session_name = params["session_name"]
    session_id = params["session_id"]
    scope = params["scope"]
    user_role = request["user"]["role"]
    user_uuid = request["user"]["uuid"]
    access_key = params["owner_access_key"]
    if access_key is None:
        access_key = request["keypair"]["access_key"]
    if user_role == UserRole.USER:
        if access_key != request["keypair"]["access_key"]:
            raise GenericForbidden
    group_name = params["group_name"]
    my_queue: asyncio.Queue[Sentinel | SessionEventInfo] = asyncio.Queue()
    log.info("PUSH_SESSION_EVENTS (ak:{}, s:{}, g:{})", access_key, session_name, group_name)
    if group_name == "*":
        group_id = "*"
    else:
        async with root_ctx.db.begin_readonly() as conn:
            query = sa.select([groups.c.id]).select_from(groups).where(groups.c.name == group_name)
            result = await conn.execute(query)
            row = result.first()
            if row is None:
                raise GroupNotFound
            group_id = row["id"]
    app_ctx.session_event_queues.add(my_queue)
    defer(lambda: app_ctx.session_event_queues.remove(my_queue))
    async with sse_response(request) as resp:
        try:
            while True:
                evdata = await my_queue.get()
                try:
                    if evdata is sentinel:
                        break
                    event_name, row, reason, exit_code = evdata
                    if user_role in (UserRole.USER, UserRole.ADMIN):
                        if row["domain_name"] != request["user"]["domain_name"]:
                            continue
                    if user_role == UserRole.USER:
                        if row["user_uuid"] != user_uuid:
                            continue
                    if group_id != "*" and row["group_id"] != group_id:
                        continue
                    if scope == "session" and not event_name.startswith("session_"):
                        continue
                    if scope == "kernel" and not event_name.startswith("kernel_"):
                        continue
                    if session_id is not None:
                        if row["session_id"] != session_id:
                            continue
                    else:
                        if session_name != "*" and not (
                            (row["session_name"] == session_name)
                            and (row["access_key"] == access_key)
                        ):
                            continue
                    response_data = {
                        "reason": reason,
                        "sessionName": row["session_name"],
                        "ownerAccessKey": row["access_key"],
                        "sessionId": str(row["session_id"]),
                        "exitCode": exit_code,
                    }
                    if kernel_id := row.get("id"):
                        response_data["kernelId"] = str(kernel_id)
                    if cluster_role := row.get("cluster_role"):
                        response_data["clusterRole"] = cluster_role
                    if cluster_idx := row.get("cluster_idx"):
                        response_data["clusterIdx"] = cluster_idx
                    await resp.send(dump_json_str(response_data), event=event_name)
                finally:
                    my_queue.task_done()
        finally:
            return resp


@server_status_required(READ_ALLOWED)
@auth_required
@check_api_params(
    t.Dict({
        tx.AliasedKey(["task_id", "taskId"]): tx.UUID,
    })
)
async def push_background_task_events(
    request: web.Request,
    params: Mapping[str, Any],
) -> web.StreamResponse:
    root_ctx: RootContext = request.app["_root.context"]
    task_id: uuid.UUID = params["task_id"]
    access_key = request["keypair"]["access_key"]
    log.info("PUSH_BACKGROUND_TASK_EVENTS (ak:{}, t:{})", access_key, task_id)
    async with sse_response(request) as resp:
        propagator = WithCachePropagator(root_ctx.event_fetcher)
        root_ctx.event_hub.register_event_propagator(
            propagator, [(EventDomain.BGTASK, str(task_id))]
        )
        try:
            cache_id = EventCacheDomain.BGTASK.cache_id(task_id)
            async for event in propagator.receive(cache_id):
                user_event = event.user_event()
                if user_event is None:
                    log.warning(
                        "Received unsupported user event: {}",
                        event.event_name(),
                    )
                    continue
                await resp.send(
                    dump_json_str(user_event.user_event_mapping()),
                    event=user_event.event_name(),
                    retry=user_event.retry_count(),
                )
                if user_event.is_close_event():
                    log.debug(
                        "Received close event: {}",
                        user_event.event_name(),
                    )
                    break
            await resp.send(dump_json_str({}), event="server_close")
        finally:
            root_ctx.event_hub.unregister_event_propagator(propagator.id())
    return resp


async def enqueue_kernel_creation_status_update(
    app: web.Application,
    source: AgentId,
    event: KernelPreparingBroadcastEvent
    | KernelPullingBroadcastEvent
    | KernelCreatingBroadcastEvent
    | KernelStartedBroadcastEvent,
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]

    async def _fetch():
        async with root_ctx.db.begin_readonly() as conn:
            query = (
                sa.select([
                    kernels.c.id,
                    kernels.c.session_id,
                    kernels.c.session_name,
                    kernels.c.access_key,
                    kernels.c.cluster_role,
                    kernels.c.cluster_idx,
                    kernels.c.domain_name,
                    kernels.c.group_id,
                    kernels.c.user_uuid,
                ])
                .select_from(kernels)
                .where(
                    (kernels.c.id == event.kernel_id),
                )
            )
            result = await conn.execute(query)
            return result.first()

    row = await execute_with_retry(_fetch)
    if row is None:
        return
    for q in app_ctx.session_event_queues:
        q.put_nowait((event.event_name(), row._mapping, event.reason, None))


async def enqueue_kernel_termination_status_update(
    app: web.Application,
    agent_id: AgentId,
    event: KernelCancelledBroadcastEvent
    | KernelTerminatingBroadcastEvent
    | KernelTerminatedBroadcastEvent,
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]

    async def _fetch():
        async with root_ctx.db.begin_readonly() as conn:
            query = (
                sa.select([
                    kernels.c.id,
                    kernels.c.session_id,
                    kernels.c.session_name,
                    kernels.c.access_key,
                    kernels.c.cluster_role,
                    kernels.c.cluster_idx,
                    kernels.c.domain_name,
                    kernels.c.group_id,
                    kernels.c.user_uuid,
                ])
                .select_from(kernels)
                .where(
                    (kernels.c.id == event.kernel_id),
                )
            )
            result = await conn.execute(query)
            return result.first()

    row = await execute_with_retry(_fetch)
    if row is None:
        return
    for q in app_ctx.session_event_queues:
        exit_code = (
            event.exit_code
            if isinstance(event, (KernelTerminatingBroadcastEvent, KernelTerminatedBroadcastEvent))
            else None
        )
        q.put_nowait((event.event_name(), row._mapping, event.reason, exit_code))


async def enqueue_session_creation_status_update(
    app: web.Application,
    source: AgentId,
    event: (
        SessionEnqueuedBroadcastEvent
        | SessionScheduledBroadcastEvent
        | SessionStartedBroadcastEvent
        | SessionCancelledBroadcastEvent
    ),
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]

    async def _fetch() -> SessionRow | None:
        async with root_ctx.db.begin_readonly_session() as db_session:
            query = (
                sa.select(SessionRow)
                .where(SessionRow.id == event.session_id)
                .options(
                    load_only(
                        SessionRow.id,
                        SessionRow.name,
                        SessionRow.access_key,
                        SessionRow.domain_name,
                        SessionRow.group_id,
                        SessionRow.user_uuid,
                    )
                )
            )
            return await db_session.scalar(query)

    row = await execute_with_retry(_fetch)
    if row is None:
        return
    row_map = {
        "session_id": row.id,
        "session_name": row.name,
        "domain_name": row.domain_name,
        "user_uuid": row.user_uuid,
        "group_id": row.group_id,
        "access_key": row.access_key,
    }
    for q in app_ctx.session_event_queues:
        q.put_nowait((event.event_name(), row_map, event.reason, None))


async def enqueue_session_termination_status_update(
    app: web.Application,
    agent_id: AgentId,
    event: SessionTerminatingBroadcastEvent | SessionTerminatedBroadcastEvent,
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]

    async def _fetch() -> SessionRow | None:
        async with root_ctx.db.begin_readonly_session() as db_session:
            query = (
                sa.select(SessionRow)
                .where(SessionRow.id == event.session_id)
                .options(
                    load_only(
                        SessionRow.id,
                        SessionRow.name,
                        SessionRow.access_key,
                        SessionRow.domain_name,
                        SessionRow.group_id,
                        SessionRow.user_uuid,
                    )
                )
            )
            return await db_session.scalar(query)

    row = await execute_with_retry(_fetch)
    if row is None:
        return
    row_map = {
        "session_id": row.id,
        "session_name": row.name,
        "domain_name": row.domain_name,
        "user_uuid": row.user_uuid,
        "group_id": row.group_id,
        "access_key": row.access_key,
    }
    for q in app_ctx.session_event_queues:
        q.put_nowait((event.event_name(), row_map, event.reason, None))


async def enqueue_batch_task_result_update(
    app: web.Application,
    agent_id: AgentId,
    event: SessionSuccessBroadcastEvent | SessionFailureBroadcastEvent,
) -> None:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]

    async def _fetch() -> SessionRow | None:
        async with root_ctx.db.begin_readonly_session() as db_session:
            query = (
                sa.select(SessionRow)
                .where(SessionRow.id == event.session_id)
                .options(
                    load_only(
                        SessionRow.id,
                        SessionRow.name,
                        SessionRow.access_key,
                        SessionRow.domain_name,
                        SessionRow.group_id,
                        SessionRow.user_uuid,
                    )
                )
            )
            return await db_session.scalar(query)

    row = await execute_with_retry(_fetch)
    if row is None:
        return
    row_map = {
        "session_id": row.id,
        "session_name": row.name,
        "domain_name": row.domain_name,
        "user_uuid": row.user_uuid,
        "group_id": row.group_id,
        "access_key": row.access_key,
    }
    for q in app_ctx.session_event_queues:
        q.put_nowait((event.event_name(), row_map, event.reason, event.exit_code))


@attrs.define(slots=True, auto_attribs=True, init=False)
class PrivateContext:
    session_event_queues: Set[asyncio.Queue[Sentinel | SessionEventInfo]]


async def events_app_ctx(app: web.Application) -> AsyncIterator[None]:
    root_ctx: RootContext = app["_root.context"]
    app_ctx: PrivateContext = app["events.context"]
    app_ctx.session_event_queues = set()
    event_dispatcher: EventDispatcher = root_ctx.event_dispatcher
    event_dispatcher.subscribe(
        SessionEnqueuedBroadcastEvent, app, enqueue_session_creation_status_update
    )
    event_dispatcher.subscribe(
        SessionScheduledBroadcastEvent, app, enqueue_session_creation_status_update
    )
    event_dispatcher.subscribe(
        KernelPreparingBroadcastEvent, app, enqueue_kernel_creation_status_update
    )
    event_dispatcher.subscribe(
        KernelPullingBroadcastEvent, app, enqueue_kernel_creation_status_update
    )
    event_dispatcher.subscribe(
        KernelCreatingBroadcastEvent, app, enqueue_kernel_creation_status_update
    )
    event_dispatcher.subscribe(
        KernelStartedBroadcastEvent, app, enqueue_kernel_creation_status_update
    )
    event_dispatcher.subscribe(
        SessionStartedBroadcastEvent, app, enqueue_session_creation_status_update
    )
    event_dispatcher.subscribe(
        KernelTerminatingBroadcastEvent, app, enqueue_kernel_termination_status_update
    )
    event_dispatcher.subscribe(
        KernelTerminatedBroadcastEvent, app, enqueue_kernel_termination_status_update
    )
    event_dispatcher.subscribe(
        KernelCancelledBroadcastEvent, app, enqueue_kernel_termination_status_update
    )
    event_dispatcher.subscribe(
        SessionTerminatingBroadcastEvent, app, enqueue_session_termination_status_update
    )
    event_dispatcher.subscribe(
        SessionTerminatedBroadcastEvent, app, enqueue_session_termination_status_update
    )
    event_dispatcher.subscribe(
        SessionCancelledBroadcastEvent, app, enqueue_session_creation_status_update
    )
    event_dispatcher.subscribe(SessionSuccessBroadcastEvent, app, enqueue_batch_task_result_update)
    event_dispatcher.subscribe(SessionFailureBroadcastEvent, app, enqueue_batch_task_result_update)
    yield


async def events_shutdown(app: web.Application) -> None:
    # shutdown handler is called before waiting for closing active connections.
    # We need to put sentinels here to ensure delivery of them to active SSE connections.
    app_ctx: PrivateContext = app["events.context"]
    join_tasks = []
    for sq in app_ctx.session_event_queues:
        sq.put_nowait(sentinel)
        join_tasks.append(sq.join())
    await asyncio.gather(*join_tasks)


def create_app(
    default_cors_options: CORSOptions,
) -> Tuple[web.Application, Iterable[WebMiddleware]]:
    app = web.Application()
    app["prefix"] = "events"
    app["events.context"] = PrivateContext()
    app["api_versions"] = (3, 4)
    app.on_shutdown.append(events_shutdown)
    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
    add_route = app.router.add_route
    app.cleanup_ctx.append(events_app_ctx)
    cors.add(add_route("GET", r"/background-task", push_background_task_events))
    cors.add(add_route("GET", r"/session", push_session_events))
    return app, []
