import base64
import json
import logging

from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from fastapi import HTTPException, Request, Response
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint

from app.config import settings

logger = logging.getLogger("videogen")


def _get_aesgcm() -> AESGCM:
    """Get AESGCM instance from configured key."""
    key_bytes = base64.b64decode(settings.ENCRYPTION_KEY)
    if len(key_bytes) != 32:
        # Pad or truncate to 32 bytes
        key_bytes = (key_bytes + b"\x00" * 32)[:32]
    return AESGCM(key_bytes)


def encrypt_data(plaintext: bytes) -> bytes:
    """Encrypt data using AES-256-GCM. Returns nonce + ciphertext."""
    import os
    aesgcm = _get_aesgcm()
    nonce = os.urandom(12)
    ciphertext = aesgcm.encrypt(nonce, plaintext, None)
    return base64.b64encode(nonce + ciphertext)


def decrypt_data(data: bytes) -> bytes:
    """Decrypt AES-256-GCM encrypted data."""
    raw = base64.b64decode(data)
    nonce = raw[:12]
    ciphertext = raw[12:]
    aesgcm = _get_aesgcm()
    return aesgcm.decrypt(nonce, ciphertext, None)


class RequestEncryptMiddleware(BaseHTTPMiddleware):
    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        # 白名单：支付回调接口不需要加密/解密
        path = request.url.path
        if "/payments/alipay/callback" in path or "/payments/wechat/callback" in path:
            return await call_next(request)
        
        encrypted = request.headers.get("X-Encrypted", "").lower() == "true"
        if not encrypted:
            return await call_next(request)

        # GET requests have no body to decrypt — just pass through
        if request.method == "GET":
            response = await call_next(request)
            response_body = b""
            async for chunk in response.body_iterator:
                if isinstance(chunk, str):
                    response_body += chunk.encode()
                else:
                    response_body += chunk
            encrypted_response = encrypt_data(response_body)
            return Response(
                content=json.dumps({"data": encrypted_response.decode()}),
                status_code=response.status_code,
                headers={"X-Encrypted": "true", "Content-Type": "application/json"},
            )

        try:
            body = await request.body()
            # Frontend sends {"data": "<encrypted_base64>"}
            body_json = json.loads(body)
            encrypted_b64 = body_json.get("data", "")
            decrypted = decrypt_data(encrypted_b64.encode())

            # Replace request body with decrypted content
            async def receive():
                return {"type": "http.request", "body": decrypted, "more_body": False}

            request._receive = receive
            request._body = decrypted

            response = await call_next(request)

            # Encrypt response
            response_body = b""
            async for chunk in response.body_iterator:
                if isinstance(chunk, str):
                    response_body += chunk.encode()
                else:
                    response_body += chunk

            encrypted_response = encrypt_data(response_body)
            # Wrap in {"data": "..."} to match frontend's expected format
            return Response(
                content=json.dumps({"data": encrypted_response.decode()}),
                status_code=response.status_code,
                headers={"X-Encrypted": "true", "Content-Type": "application/json"},
            )
        except Exception:
            logger.exception("Request decryption failed")
            return Response(
                content=json.dumps({"detail": "解密失败"}, ensure_ascii=False),
                status_code=400,
                media_type="application/json",
            )