# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# OpenOLC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         `http://license.coscl.org.cn/MulanPSL2`
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""OlcFastAPIAdapter —— 纯 ASGI 形态的 OLC FastAPI 适配器。

直接以标准 ASGI 三元函数形式实现（`__call__(scope, receive, send)`），
绕开 starlette BaseHTTPMiddleware 内部 task + memory stream 调度，将
适配器自身开销压到接近原生 ASGI 中间件水平。

对外用法保持不变：
    from olc.adapters.fastapi import OlcFastAPIAdapter, OlcAdapterConfig
    app.add_middleware(OlcFastAPIAdapter, config=adapter_config)
"""

import asyncio
import inspect
import logging
from typing import Optional

from fastapi.responses import JSONResponse
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp, Message, Receive, Scope, Send

from olc.adapters.fastapi.adapter_config import OlcAdapterConfig
from olc.bean.dimension import Method, URL
from olc.bean.olc_control_request import OlcControlRequest
from olc.config.properties_loader import OlcConfigManager
from olc.control.context.context import OlcContext
from olc.olc import OLC

logger = logging.getLogger(__name__)


class OlcFastAPIAdapter:
    """OLC 的 FastAPI / Starlette 适配器（纯 ASGI 中间件）。

    主要行为：
      - SDK 总开关 OFF 时短路透传
      - 标签提取：默认从 scope 直读 path/method；自定义 extractor 才构造 Request
      - block 分支：复用 starlette Response 对象，通过 response(scope, receive, send) 发送
      - 放行分支：包装 send，在响应 body 结束时 asyncio.create_task 触发 _cleanup
      - 异常兜底：发生异常时透传给下游，并在 finally 中确保 _cleanup 至少触发一次
    """

    def __init__(
        self,
        app: ASGIApp,
        config: Optional[OlcAdapterConfig] = None,
    ) -> None:
        self.app = app
        self.config = config or OlcAdapterConfig()

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        # 非 http 流量（websocket / lifespan 等）直接透传
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        # SDK 总开关 OFF：透传
        if not OlcConfigManager.get_instance().get_sdk_config().switch:
            await self.app(scope, receive, send)
            return

        context: Optional[OlcContext] = None
        try:
            # 标签提取
            tags = await self._extract_tags(scope, receive)
            olc_request = OlcControlRequest(
                tags=tags,
                token_number=self.config.token_number,
            )
            context = await OLC.async_process(olc_request)

            # 命中拦截：复用 starlette Response 对象发送，并在发送完成后立即清理
            if context.result.block:
                response = await self._build_block_response(scope, receive, context)
                await response(scope, receive, send)
                await self._cleanup(context)
                return

        except Exception as e:
            logger.error(f"OLC process error: {e}", exc_info=True)
            # OLC 自身处理失败时不阻断业务，按透传处理
            await self.app(scope, receive, send)
            return

        # 放行分支：包装 send，捕获响应 body 末尾消息后异步清理
        cleanup_scheduled = False

        async def send_wrapper(message: Message) -> None:
            nonlocal cleanup_scheduled
            await send(message)
            if (
                not cleanup_scheduled
                and message["type"] == "http.response.body"
                and not message.get("more_body", False)
            ):
                cleanup_scheduled = True
                asyncio.create_task(self._cleanup(context))

        try:
            await self.app(scope, receive, send_wrapper)
        finally:
            # 兜底：下游早 close 或异常导致 body 末尾消息未发出时，确保清理
            if not cleanup_scheduled and context is not None:
                asyncio.create_task(self._cleanup(context))

    async def _cleanup(self, context: OlcContext) -> None:
        """异步清理 OLC 上下文，对外吞掉异常避免污染请求生命周期。"""
        try:
            await OLC.async_clean(context)
        except Exception as e:
            logger.error(f"OLC cleanup error: {e}", exc_info=True)

    async def _extract_tags(self, scope: Scope, receive: Receive) -> dict:
        """提取请求标签。

        - 默认 extractor：直接从 scope 取 path / method，避免构造 Request 对象
        - 自定义 extractor：构造一次 Request(scope, receive=receive) 后交给用户回调
        """
        if self.config.tag_extractor is None:
            return self._default_tag_extractor_from_scope(scope)

        request = Request(scope, receive=receive)
        result = self.config.tag_extractor(request)
        if inspect.isawaitable(result):
            return await result
        return result

    @staticmethod
    def _default_tag_extractor_from_scope(scope: Scope) -> dict:
        """默认标签提取器：直接从 ASGI scope 中读取 path / method。"""
        return {
            URL: scope.get("path", ""),
            Method: scope.get("method", ""),
        }

    async def _build_block_response(
        self, scope: Scope, receive: Receive, context: OlcContext
    ) -> Response:
        """构造拦截响应。

        与原适配器一致：优先使用 block_response_builder；否则返回 429 JSONResponse。
        builder 入参仍是 starlette Request，保持对外接口不变。
        """
        block_msg = context.result.block_msg or self.config.block_msg
        if self.config.block_response_builder is not None:
            request = Request(scope, receive=receive)
            result = self.config.block_response_builder(request, block_msg)
            if inspect.isawaitable(result):
                return await result
            return result
        return JSONResponse(
            status_code=429,
            content={
                "error": "Too Many Requests",
                "message": block_msg,
                "path": scope.get("path", ""),
            },
        )