#!/usr/bin/env python3

# Libervia plugin for OpenPGP for XMPP
# Copyright (C) 2025-2026 Jérôme Poisosn (goffi@goffi.org)
# Copyright (C) 2022-2026 Tim Henkes (me@syndace.dev)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import tempfile
from typing import Any

import gnupg

from .exceptions import (
    DecryptionFailed,
    GPGProviderError,
    InvalidPacket,
    UnknownKey,
    VerificationFailed,
)
from .models import GPGProviderBase, GPGPublicKeyBase, GPGSecretKeyBase


class GPGPublicKey(GPGPublicKeyBase):
    """
    GPG public key implementation based on python-gnupg.
    """

    def __init__(self, key_obj: dict[str, Any]) -> None:
        """
        @param key_obj: The python-gnupg key dictionary.
        """

        self._key_obj = key_obj

    @property
    def fingerprint(self) -> str:
        return self._key_obj["fingerprint"]

    @property
    def key_obj(self) -> dict[str, Any]:
        return self._key_obj


class GPGSecretKey(GPGSecretKeyBase):
    """
    GPG secret key implementation based on python-gnupg.
    """

    def __init__(self, public_key: GPGPublicKey) -> None:
        """
        @param public_key: The public key corresponding to this secret key.
        """

        self._public_key = public_key

    @property
    def public_key(self) -> GPGPublicKey:
        return self._public_key


class GPGProvider(GPGProviderBase):
    """
    GPG provider implementation based on python-gnupg.
    """

    ALGORITHM = "ed25519"

    def __init__(self, home_dir: str | None = None) -> None:
        """
        @param home_dir: Optional GPG home directory path to use for all operations.
        """

        self._home_dir = home_dir
        self._gpg = gnupg.GPG(gnupghome=self._home_dir)

    def export_public_key(self, public_key: GPGPublicKeyBase) -> bytes:
        assert isinstance(public_key, GPGPublicKey)

        try:
            result = self._gpg.export_keys(public_key.fingerprint, armor=False)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result:
            raise UnknownKey(f"Public key {public_key.fingerprint} not found.")

        return result

    def import_public_key(self, packet: bytes) -> GPGPublicKeyBase:
        try:
            result = self._gpg.import_keys(packet)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "fingerprints") or not result.fingerprints:
            raise InvalidPacket("Data not considered for public key import")

        if len(result.fingerprints) != 1:
            raise InvalidPacket(
                "Public key packet does not contain exactly one public key (not"
                " counting subkeys)."
            )

        fingerprint = result.fingerprints[0]
        key_list = self._gpg.list_keys(keys=fingerprint)
        if not key_list:
            raise GPGProviderError("Newly imported public key not found")

        return GPGPublicKey(key_list[0])

    def backup_secret_key(self, secret_key: GPGSecretKeyBase) -> bytes:
        assert isinstance(secret_key, GPGSecretKey)

        try:
            result = self._gpg.export_keys(
                secret_key.public_key.fingerprint, secret=True, armor=False
            )
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result:
            raise UnknownKey(f"Secret key {secret_key.public_key.fingerprint} not found.")

        return result

    def restore_secret_keys(self, data: bytes) -> set[GPGSecretKeyBase]:
        try:
            result = self._gpg.import_keys(data)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "fingerprints") or not result.fingerprints:
            raise InvalidPacket("Data not considered for secret key import")

        secret_keys = set()
        for fingerprint in result.fingerprints:
            key_list = self._gpg.list_keys(keys=fingerprint, secret=True)
            if not key_list:
                raise GPGProviderError("Newly imported secret key not found")

            secret_keys.add(GPGSecretKey(GPGPublicKey(key_list[0])))

        return secret_keys

    def encrypt_symmetrically(self, plaintext: bytes, password: str) -> bytes:
        try:
            result = self._gpg.encrypt(
                plaintext, None, passphrase=password, symmetric=True, armor=False
            )
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data") or not result.data:
            raise DecryptionFailed("Encryption failed")

        return result.data

    def decrypt_symmetrically(self, ciphertext: bytes, password: str) -> bytes:
        try:
            result = self._gpg.decrypt(ciphertext, passphrase=password)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data") or not result.data:
            raise DecryptionFailed("Decryption failed")

        return result.data

    def sign(self, data: bytes, secret_keys: set[GPGSecretKeyBase]) -> bytes:
        if not secret_keys:
            raise GPGProviderError("No secret keys provided for signing")

        secret_key = next(iter(secret_keys))
        assert isinstance(secret_key, GPGSecretKey)

        try:
            result = self._gpg.sign(
                data,
                keyid=secret_key.public_key.fingerprint,
                detach=False,
                binary=True,
            )
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data") or not result.data:
            raise GPGProviderError("Signing failed")

        return result.data

    def sign_detached(self, data: bytes, secret_keys: set[GPGSecretKeyBase]) -> bytes:
        if not secret_keys:
            raise GPGProviderError("No secret keys provided for signing")

        secret_key = next(iter(secret_keys))
        assert isinstance(secret_key, GPGSecretKey)

        try:
            result = self._gpg.sign(
                data,
                keyid=secret_key.public_key.fingerprint,
                detach=True,
                binary=True,
            )
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data") or not result.data:
            raise GPGProviderError("Detached signing failed")

        return result.data

    def verify(self, signed_data: bytes, public_keys: set[GPGPublicKeyBase]) -> bytes:
        try:
            result = self._gpg.verify(signed_data)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        for public_key in public_keys:
            assert isinstance(public_key, GPGPublicKey)
            if result.pubkey_fingerprint == public_key.fingerprint:
                break
        else:
            raise VerificationFailed("Data not signed by one of the expected public keys")

        return result.data

    def verify_detached(
        self, data: bytes, signature: bytes, public_keys: set[GPGPublicKeyBase]
    ) -> None:
        try:
            # FIXME: There is no method in gnupg to have signature and data in memory, we
            #   have to write signature to a file. This is probably due to the limitation
            #   of the gpg CLI.
            with tempfile.NamedTemporaryFile(delete=False) as signature_file:
                signature_file.write(signature)
                signature_path = signature_file.name
            try:
                result = self._gpg.verify_data(signature_path, data)
            finally:
                os.unlink(signature_path)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        for public_key in public_keys:
            assert isinstance(public_key, GPGPublicKey)
            if result.pubkey_fingerprint == public_key.fingerprint:
                break
        else:
            raise VerificationFailed("Data not signed by one of the expected public keys")

    def encrypt(
        self,
        plaintext: bytes,
        public_keys: set[GPGPublicKeyBase],
        signing_keys: set[GPGSecretKeyBase] | None = None,
    ) -> bytes:
        recipients = []
        for public_key in public_keys:
            assert isinstance(public_key, GPGPublicKey)
            recipients.append(public_key.fingerprint)

        kwargs: dict[str, Any] = {"armor": False}
        if signing_keys:
            secret_key = next(iter(signing_keys))
            assert isinstance(secret_key, GPGSecretKey)
            kwargs["sign"] = secret_key.public_key.fingerprint

        try:
            result = self._gpg.encrypt(plaintext, recipients, **kwargs)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data") or not result.data:
            raise GPGProviderError("Encryption failed")

        return result.data

    def decrypt(
        self,
        ciphertext: bytes,
        secret_keys: set[GPGSecretKeyBase],
        public_keys: set[GPGPublicKeyBase] | None = None,
    ) -> bytes:
        verify = public_keys is not None

        try:
            result = self._gpg.decrypt(ciphertext)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "data"):
            raise DecryptionFailed("Decryption failed")

        if verify:
            valid_signature_found = False
            for public_key in public_keys:
                assert isinstance(public_key, GPGPublicKey)

                if result.pubkey_fingerprint == public_key.fingerprint:
                    valid_signature_found = True
                    break

            if not valid_signature_found:
                raise VerificationFailed(
                    "Data not signed by one of the expected public keys"
                )

        return result.data

    def list_public_keys(self, user_id: str) -> set[GPGPublicKeyBase]:
        try:
            key_list = self._gpg.list_keys(keys=user_id)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        return {GPGPublicKey(key) for key in key_list if key.get("type") == "pub"}

    def list_secret_keys(self, user_id: str) -> set[GPGSecretKeyBase]:
        try:
            key_list = self._gpg.list_keys(keys=user_id, secret=True)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        return {
            GPGSecretKey(GPGPublicKey(key))
            for key in key_list
            if key.get("type") == "sec"
        }

    def can_sign(self, public_key: GPGPublicKeyBase) -> bool:
        assert isinstance(public_key, GPGPublicKey)

        key_obj = public_key.key_obj
        return "s" in key_obj.get("cap", "").lower()

    def can_encrypt(self, public_key: GPGPublicKeyBase) -> bool:
        assert isinstance(public_key, GPGPublicKey)

        key_obj = public_key.key_obj
        return "c" in key_obj.get("cap", "").lower()

    def create_key(self, user_id: str) -> GPGSecretKeyBase:
        key_input = self._gpg.gen_key_input(
            key_type=self.ALGORITHM,
            name_real=user_id,
            no_protection=True,
            expire_data=0,
        )
        try:
            result = self._gpg.gen_key(key_input)
        except Exception as e:
            raise GPGProviderError("Internal python-gnupg error") from e

        if not result or not hasattr(result, "fingerprint") or not result.fingerprint:
            raise GPGProviderError("Key creation failed")

        key_list = self._gpg.list_keys(keys=result.fingerprint, secret=True)
        if not key_list:
            raise GPGProviderError("Newly created key not found")

        return GPGSecretKey(GPGPublicKey(key_list[0]))
