# -*- coding: utf-8 -*-
"""FastAPI 中间件：镜像 Flask after_request 的响应侧行为。

与 main.py `after_request` 对齐的关键点：
1. 音频直链保护——JSON 响应中的网易云音频地址替换为 HMAC 签名代理地址
   （`/audio_stream?token=...`），使直链不可被外站直接盗用；密钥与 Flask 同源，
   故两栈签发的 token 互相可校验。
2. 零成本安全响应头（nosniff / DENY / Referrer-Policy / CSP 基线）。
3. 抹去服务器版本指纹（Server 头统一为 NeteaseUrl）。
4. 429 统一带 Retry-After，供客户端退避。

实现注意：
- 用纯 ASGI 中间件逐段转发（不用 BaseHTTPMiddleware，避免它缓冲整段流式响应，
  破坏大 ZIP / 音频流式下载）。
- 仅对「JSON 且 200」的响应缓冲改写；此时必须延迟发送 response.start，待改写完成
  后**重算 Content-Length**——否则长度不一致会让客户端 IncompleteRead（实测踩过）。
"""
import json

from starlette.types import ASGIApp, Message, Receive, Scope, Send

from core import security

_SECURITY_HEADERS = (
    (b'x-content-type-options', b'nosniff'),
    (b'x-frame-options', b'DENY'),
    (b'referrer-policy', b'same-origin'),
    (b'content-security-policy',
     b"default-src 'self'; script-src 'self' 'unsafe-inline'; style-src 'self' 'unsafe-inline'; "
     b"img-src 'self' data: http: https:; media-src 'self' blob:; connect-src 'self'; "
     b"font-src 'self' data:; object-src 'none'; frame-ancestors 'none'; base-uri 'self'"),
)


def _augment_headers(headers: list[tuple[bytes, bytes]], status_code: int) -> list[tuple[bytes, bytes]]:
    existing = {k.lower() for k, _ in headers}
    out = list(headers)
    for hk, hv in _SECURITY_HEADERS:
        if hk not in existing:
            out.append((hk, hv))
    out.append((b'server', b'NeteaseUrl'))
    if status_code == 429 and b'retry-after' not in existing:
        out.append((b'retry-after', b'30'))
    return out


class ResponseHygieneMiddleware:
    """响应侧统一处理（安全头 + 音频直链保护 + 版本指纹抹除）。"""

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope['type'] != 'http':
            await self.app(scope, receive, send)
            return

        state = {'hold': False, 'start': None, 'chunks': []}

        async def send_wrapper(message: Message) -> None:
            mtype = message['type']
            if mtype == 'http.response.start':
                status = message['status']
                headers = _augment_headers(list(message.get('headers') or []), status)
                ctype = b''.join(v.lower() for k, v in headers if k.lower() == b'content-type')
                if status == 200 and b'application/json' in ctype:
                    # 需改写 body：延迟发送 start，待重算 Content-Length
                    state['hold'] = True
                    state['start'] = message
                    message['headers'] = headers
                else:
                    message['headers'] = headers
                    await send(message)
            elif mtype == 'http.response.body':
                if not state['hold']:
                    await send(message)
                    return
                state['chunks'].append(message.get('body', b''))
                if message.get('more_body'):
                    return
                await self._flush(state, send)

        await self.app(scope, receive, send_wrapper)

    async def _flush(self, state: dict, send: Send) -> None:
        raw = b''.join(state['chunks'])
        try:
            data = json.loads(raw)
            if isinstance(data, (dict, list)) and security.protect_audio_urls(data):
                raw = json.dumps(data, ensure_ascii=False).encode('utf-8')
        except Exception:
            pass
        headers = [(k, v) for k, v in state['start']['headers'] if k.lower() != b'content-length']
        headers.append((b'content-length', str(len(raw)).encode('ascii')))
        state['start']['headers'] = headers
        await send(state['start'])
        await send({'type': 'http.response.body', 'body': raw, 'more_body': False})