import base64
import os
from contextvars import ContextVar
from pathlib import Path
from typing import Dict

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCMSIV
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt

from loveace.config.manager import config_manager

console = Console()

rsa_context: ContextVar[Dict[str, "RSAUtils"]] = ContextVar("rsa_context")


class RSAUtils:
    """RSA 工具类，支持 AES-GCM-SIV 加密的密钥保护"""

    private_key_path: str
    private_key: RSAPrivateKey
    public_key: RSAPublicKey

    def __init__(self, private_key_path: str | None = None):
        """初始化 RSAUtils 类

        Args:
            private_key_path (str): 私钥文件路径
        """
        settings = config_manager.get_settings()
        self.private_key_path = str(
            Path(settings.app.rsa_protect_key_path).joinpath(
                Path(
                    private_key_path
                    or config_manager.get_settings().app.rsa_private_key_path
                ).name
            )
        )
        # 转换路径扩展名为 .hex
        self.private_key_path = str(self.private_key_path).replace(".pem", ".hex")
        self.load_keys()

    def _derive_key_from_password(
        self, password: str, salt: bytes | None = None
    ) -> tuple[bytes, bytes]:
        """从密码派生 AES 密钥

        Args:
            password (str): 用户输入的密码
            salt (bytes): 盐值，如果为 None 则生成新的

        Returns:
            tuple[bytes, bytes]: (派生密钥, 盐值)
        """
        if salt is None:
            salt = os.urandom(16)

        kdf_obj = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=16,  # AES-128 需要 16 字节密钥
            salt=salt,
            iterations=100000,
        )
        key = kdf_obj.derive(password.encode("utf-8"))
        return key, salt

    def load_keys(self):
        """加载密钥对（从加密的 AES 文件中）"""
        path = Path(self.private_key_path)
        console.print(
            Panel(
                f"[bold cyan]正在操作密钥文件[/bold cyan]\n"
                f"[cyan]文件路径：{self.private_key_path}[/cyan]",
                expand=False,
            )
        )
        if not path.exists():
            console.print(
                Panel(
                    "[bold yellow]RSA 密钥对不存在，将为您生成新的密钥对[/bold yellow]",
                    title="[bold blue]密钥生成[/bold blue]",
                    expand=False,
                )
            )
            self.generate_keys()
        else:
            self._load_encrypted_key()

    def _load_encrypted_key(self):
        """从加密的 .hex 文件加载密钥"""
        console.print(
            Panel(
                f"[bold cyan]检测到本地 RSA 私钥文件[/bold cyan]\n"
                f"[cyan]文件路径：{self.private_key_path}[/cyan]",
                title="[bold blue]密钥加载[/bold blue]",
                expand=False,
            )
        )

        console.print(
            "[bold yellow]该密钥文件受密码保护，需要您输入密码来解密[/bold yellow]"
        )
        password = Prompt.ask(
            "[bold]请输入 RSA 私钥密码[/bold]", password=True, console=console
        )

        with open(self.private_key_path, "rb") as key_file:
            encrypted_data = key_file.read()

        # 解析加密数据：salt(16) + nonce(12) + ciphertext
        salt = encrypted_data[:16]
        nonce = encrypted_data[16:28]
        ciphertext = encrypted_data[28:]

        # 派生密钥
        key, _ = self._derive_key_from_password(password, salt)

        # 使用 AES-GCM-SIV 解密
        try:
            aesgcmsiv = AESGCMSIV(key)
            plaintext = aesgcmsiv.decrypt(nonce, ciphertext, None)
            console.print("[bold green]✓ 私钥密码验证成功[/bold green]")
        except Exception as e:
            console.print(
                Panel(
                    "[bold red]✗ 私钥密码错误或密钥文件已损坏[/bold red]",
                    title="[bold red]错误[/bold red]",
                    expand=False,
                )
            )
            raise ValueError("Invalid password or corrupted key file") from e

        # 加载 PEM 格式的私钥
        try:
            pk = serialization.load_pem_private_key(
                plaintext, password=None, backend=default_backend()
            )
            if isinstance(pk, RSAPrivateKey):
                self.private_key = pk
            else:
                raise ValueError("Loaded key is not an RSA private key")
        except Exception:
            console.print(
                Panel(
                    "[bold red]✗ 密钥格式错误[/bold red]",
                    title="[bold red]错误[/bold red]",
                    expand=False,
                )
            )
            raise

        self.public_key = self.private_key.public_key()

    def generate_keys(self, key_size: int = 2048):
        """生成 RSA 密钥对并使用 AES 加密保存到文件

        Args:
            key_size (int): 密钥大小，默认2048位
        """
        path = Path(self.private_key_path)
        path.parent.mkdir(parents=True, exist_ok=True)

        # 提示用户设置密码
        console.print(
            Panel(
                "[bold cyan]请设置 RSA 私钥密码（用于保护密钥文件）[/bold cyan]",
                title="[bold blue]密钥保护[/bold blue]",
                expand=False,
            )
        )
        password = Prompt.ask("[bold]请输入密码[/bold]", password=True, console=console)
        password_confirm = Prompt.ask(
            "[bold]请确认密码[/bold]", password=True, console=console
        )

        if password != password_confirm:
            console.print(
                Panel(
                    "[bold red]✗ 两次输入的密码不一致[/bold red]",
                    title="[bold red]错误[/bold red]",
                    expand=False,
                )
            )
            raise ValueError("Passwords do not match")

        # 生成 RSA 密钥对
        console.print("[bold cyan]正在生成 RSA 密钥对...[/bold cyan]")
        private_key = rsa.generate_private_key(
            public_exponent=65537, key_size=key_size, backend=default_backend()
        )
        public_key = private_key.public_key()

        # 将私钥序列化为 PEM 格式
        pem_private = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )

        pem_public = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        # 使用 AES-GCM-SIV 加密私钥
        key, salt = self._derive_key_from_password(password)
        aesgcmsiv = AESGCMSIV(key)
        nonce = os.urandom(12)
        ciphertext = aesgcmsiv.encrypt(nonce, pem_private, None)

        # 保存加密的私钥：salt + nonce + ciphertext
        with open(self.private_key_path, "wb") as private_file:
            private_file.write(salt + nonce + ciphertext)

        # 保存公钥（不加密）
        public_key_path = self.private_key_path.replace(".hex", "_public.pem")
        with open(public_key_path, "wb") as public_file:
            public_file.write(pem_public)

        self.private_key = private_key
        self.public_key = public_key

        console.print(
            Panel(
                f"[bold green]✓ RSA 密钥对生成成功[/bold green]\n"
                f"[cyan]私钥路径：[/cyan]{self.private_key_path}\n"
                f"[cyan]公钥路径：[/cyan]{public_key_path}",
                title="[bold blue]完成[/bold blue]",
                expand=False,
            )
        )

    def encrypt(self, plaintext: str) -> str:
        """使用公钥加密数据

        Args:
            plaintext (str): 明文字符串

        Returns:
            str: Base64 编码的密文字符串
        """
        ciphertext = self.public_key.encrypt(
            plaintext.encode("utf-8"),
            padding.PKCS1v15(),
        )
        return base64.b64encode(ciphertext).decode("utf-8")

    def decrypt(self, b64_ciphertext: str) -> str:
        """使用私钥解密数据

        Args:
            b64_ciphertext (str): Base64 编码的密文字符串

        Returns:
            str: 解密后的明文字符串
        """
        ciphertext = base64.b64decode(b64_ciphertext)
        plaintext = self.private_key.decrypt(
            ciphertext,
            padding.PKCS1v15(),
        )
        return plaintext.decode("utf-8")

    @staticmethod
    def encrypt_file_with_aes(
        plaintext: bytes, password: str | None = None
    ) -> tuple[bytes, str]:
        """使用 AES-GCM-SIV 和密码加密数据

        Args:
            plaintext (bytes): 明文数据
            password (str): 密码，如果为 None 则生成随机密钥

        Returns:
            tuple[bytes, str]: (加密数据, 密钥的十六进制字符串)
        """
        if password is None:
            # 生成随机密钥
            key = AESGCMSIV.generate_key(bit_length=128)
            aesgcmsiv = AESGCMSIV(key)
            nonce = os.urandom(12)
            ciphertext = aesgcmsiv.encrypt(nonce, plaintext, None)
            encrypted_data = key + nonce + ciphertext
        else:
            # 从密码派生密钥
            salt = os.urandom(16)
            kdf_obj = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=16,
                salt=salt,
                iterations=100000,
            )
            key = kdf_obj.derive(password.encode("utf-8"))
            aesgcmsiv = AESGCMSIV(key)
            nonce = os.urandom(12)
            ciphertext = aesgcmsiv.encrypt(nonce, plaintext, None)
            encrypted_data = salt + nonce + ciphertext

        key_hex = key.hex()
        return encrypted_data, key_hex

    @staticmethod
    def get_or_create_rsa_utils(private_key_path: str | None = None) -> "RSAUtils":
        """
        获取或创建 RSAUtils 实例
        Args:
            private_key_path (str | None): 私钥文件路径，如果为 None 则使用配置中的默认路径
        """
        private_key_path = (
            private_key_path or config_manager.get_settings().app.rsa_private_key_path
        )
        try:
            rsa_utils_dict = rsa_context.get()
            if private_key_path in rsa_utils_dict:
                return rsa_utils_dict[private_key_path]
            else:
                rsa_utils = RSAUtils(private_key_path)
                rsa_utils_dict[private_key_path] = rsa_utils
                rsa_context.set(rsa_utils_dict)
                return rsa_utils
        except LookupError:
            rsa_utils = RSAUtils(private_key_path)
            rsa_utils_dict = {private_key_path: rsa_utils}
            rsa_context.set(rsa_utils_dict)
            return rsa_utils