from __future__ import annotations

import base64
import hashlib
import json
import os
import uuid
from dataclasses import dataclass
from typing import Any, Mapping, Sequence
from urllib.parse import quote, urlencode, urlsplit, urlunsplit

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


@dataclass
class UgreenEncryptedRequest:
    url: str
    headers: dict[str, str]
    params: dict[str, str]
    json: dict[str, Any] | None
    aes_key: str
    plain_query: str


class UgreenCrypto:
    """
    绿联接口请求加解密工具。
    """

    def __init__(
        self,
        public_key: str,
        token: str | None = None,
        client_id: str | None = None,
        client_version: str | None = "76363",
        ug_agent: str | None = "PC/WEB",
        language: str = "zh-CN",
    ) -> None:
        self.public_key_pem = self.normalize_public_key(public_key)
        self.public_key = serialization.load_pem_public_key(
            self.public_key_pem.encode("utf-8")
        )
        self.token = token
        self.client_id = client_id
        self.client_version = client_version
        self.ug_agent = ug_agent
        self.language = language

    @staticmethod
    def normalize_public_key(public_key: str) -> str:
        key = (public_key or "").strip().strip('"').replace("\\n", "\n")
        if "BEGIN" in key:
            return key if key.endswith("\n") else f"{key}\n"
        return (
            "-----BEGIN RSA PUBLIC KEY-----\n"
            f"{key}\n"
            "-----END RSA PUBLIC KEY-----\n"
        )

    @staticmethod
    def generate_aes_key() -> str:
        return uuid.uuid4().hex

    @staticmethod
    def _flatten_query(prefix: str, value: Any) -> list[tuple[str, str]]:
        pairs: list[tuple[str, str]] = []
        if isinstance(value, Mapping):
            for key, item in value.items():
                next_prefix = f"{prefix}[{key}]" if prefix else str(key)
                pairs.extend(UgreenCrypto._flatten_query(next_prefix, item))
            return pairs
        if isinstance(value, Sequence) and not isinstance(
            value, (str, bytes, bytearray)
        ):
            for item in value:
                next_prefix = f"{prefix}[]"
                pairs.extend(UgreenCrypto._flatten_query(next_prefix, item))
            return pairs
        if isinstance(value, bool):
            pairs.append((prefix, "true" if value else "false"))
            return pairs
        if value is None:
            pairs.append((prefix, ""))
            return pairs
        pairs.append((prefix, str(value)))
        return pairs

    @classmethod
    def encode_query(cls, params: Mapping[str, Any] | None) -> str:
        if not params:
            return ""
        pairs: list[tuple[str, str]] = []
        for key, value in params.items():
            pairs.extend(cls._flatten_query(str(key), value))
        return urlencode(pairs, doseq=False, quote_via=quote, safe="")

    def rsa_encrypt_long(self, plaintext: str) -> str:
        if not plaintext:
            return ""
        key_size = self.public_key.key_size // 8
        max_chunk = key_size - 11
        encrypted_chunks: list[bytes] = []
        raw = plaintext.encode("utf-8")
        for start in range(0, len(raw), max_chunk):
            chunk = raw[start : start + max_chunk]
            encrypted_chunks.append(
                self.public_key.encrypt(chunk, padding.PKCS1v15())
            )
        return base64.b64encode(b"".join(encrypted_chunks)).decode("utf-8")

    @staticmethod
    def aes_gcm_encrypt(plaintext: str, aes_key: str) -> str:
        iv = os.urandom(12)
        cipher = AESGCM(aes_key.encode("utf-8"))
        encrypted = cipher.encrypt(iv, plaintext.encode("utf-8"), None)
        # encrypt 返回 ciphertext + tag
        return base64.b64encode(iv + encrypted).decode("utf-8")

    @staticmethod
    def aes_gcm_decrypt(payload_b64: str, aes_key: str) -> str:
        raw = base64.b64decode(payload_b64)
        iv = raw[:12]
        encrypted = raw[12:]
        cipher = AESGCM(aes_key.encode("utf-8"))
        plain = cipher.decrypt(iv, encrypted, None)
        return plain.decode("utf-8")

    @staticmethod
    def build_security_key(token: str) -> str:
        return hashlib.md5(token.encode("utf-8")).hexdigest()

    @staticmethod
    def _normalize_body(data: Any) -> str:
        if isinstance(data, str):
            return data
        if isinstance(data, (bytes, bytearray)):
            return bytes(data).decode("utf-8")
        return json.dumps(data, ensure_ascii=False, separators=(",", ":"))

    def encrypt_body(self, data: Any, aes_key: str) -> dict[str, str]:
        plain = self._normalize_body(data)
        return {
            "encrypt_req_body": self.aes_gcm_encrypt(plain, aes_key),
            "req_body_sha256": hashlib.sha256(plain.encode("utf-8")).hexdigest(),
        }

    def build_headers(
        self,
        aes_key: str,
        token: str | None = None,
        extra_headers: Mapping[str, str] | None = None,
        encrypt_token: bool = True,
    ) -> dict[str, str]:
        token_value = token if token is not None else self.token
        headers: dict[str, str] = dict(extra_headers or {})

        if self.client_id:
            headers.setdefault("Client-Id", self.client_id)
        if self.client_version:
            headers.setdefault("Client-Version", self.client_version)
        if self.ug_agent:
            headers.setdefault("UG-Agent", self.ug_agent)
        headers.setdefault("X-Specify-Language", self.language)
        headers.setdefault("Accept", "application/json, text/plain, */*")

        if token_value:
            headers["X-Ugreen-Security-Key"] = self.build_security_key(token_value)
            headers["X-Ugreen-Security-Code"] = self.rsa_encrypt_long(aes_key)
            headers["X-Ugreen-Token"] = (
                self.rsa_encrypt_long(token_value) if encrypt_token else token_value
            )
        return headers

    def build_encrypted_request(
        self,
        url: str,
        method: str = "GET",
        params: Mapping[str, Any] | None = None,
        data: Any | None = None,
        extra_headers: Mapping[str, str] | None = None,
        token: str | None = None,
        encrypt_token: bool = True,
        encrypt_body: bool = True,
    ) -> UgreenEncryptedRequest:
        """
        构建绿联加密请求。

        关键点：
        - 传入的是明文 `params`；
        - 方法内部会将其序列化并加密成 `encrypt_query`；
        - 业务侧不需要、也不应该手工拼接 `encrypt_query`。
        """
        parsed = urlsplit(url)
        clean_url = urlunsplit(
            (parsed.scheme, parsed.netloc, parsed.path, "", parsed.fragment)
        )

        url_query_plain = parsed.query
        input_query_plain = self.encode_query(params)
        plain_query = "&".join(filter(None, [url_query_plain, input_query_plain]))

        aes_key = self.generate_aes_key()
        encrypted_query = self.aes_gcm_encrypt(plain_query, aes_key)

        req_json = None
        if data is not None:
            req_json = self.encrypt_body(data, aes_key) if encrypt_body else data

        headers = self.build_headers(
            aes_key=aes_key,
            token=token,
            extra_headers=extra_headers,
            encrypt_token=encrypt_token,
        )
        if req_json is not None:
            headers.setdefault("Content-Type", "application/json")

        _ = method  # 保留参数，便于上层统一调用

        return UgreenEncryptedRequest(
            url=clean_url,
            headers=headers,
            # 绿联接口约定：查询参数统一透传为 encrypt_query
            params={"encrypt_query": encrypted_query},
            json=req_json,
            aes_key=aes_key,
            plain_query=plain_query,
        )

    def decrypt_response(self, response_json: Any, aes_key: str) -> Any:
        if not isinstance(response_json, Mapping):
            return response_json
        encrypted = response_json.get("encrypt_resp_body")
        if not encrypted:
            return response_json
        plain = self.aes_gcm_decrypt(str(encrypted), aes_key)
        try:
            return json.loads(plain)
        except json.JSONDecodeError:
            return plain
