"""
xtendr_signing.py

Core logic for signing XtendR plugins and maintaining a whitelist.
Deliberately kept free of any GTK/Adw imports so it can be unit tested
and reused by both the signer GUI and (later) XtendRSystem's verification
hook.

Design:
- Ed25519 keypair: private key stays with the signer, public key ships
  inside XtendR / PyPac to verify at attach() time.
- Each whitelist entry covers the plugin's manifest (<name>.json) and its
  module file (<module>.py), hashed together with SHA-256. The signature
  covers the whole entry (hash + name + module + class), so an attacker
  can't splice a valid hash onto a different plugin identity.
- The whitelist file itself can optionally be encrypted at rest with a
  passphrase (Fernet / AES-128-CBC+HMAC via PBKDF2-derived key). This is
  for confidentiality only -- the signature is what provides integrity,
  and still verifies correctly whether or not the file on disk is
  encrypted.
"""
from __future__ import annotations

import base64
import hashlib
import json
import os
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey,
)
from cryptography.hazmat.primitives import serialization
from cryptography.exceptions import InvalidSignature
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

PBKDF2_ITERATIONS = 600_000


# --------------------------------------------------------------------------
# Key management
# --------------------------------------------------------------------------

def generate_keypair(private_path: Path, public_path: Path) -> None:
    """Generate a new Ed25519 keypair and write it to disk (unencrypted PEM).

    Caller is responsible for keeping private_path safe (e.g. 0600 perms,
    offline machine, backups). This is deliberately not done automatically
    since the right answer depends on the user's environment.
    """
    private_path.parent.mkdir(parents=True, exist_ok=True)
    key = Ed25519PrivateKey.generate()

    priv_bytes = key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    pub_bytes = key.public_key().public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )

    private_path.write_bytes(priv_bytes)
    try:
        os.chmod(private_path, 0o600)
    except OSError:
        pass
    public_path.write_bytes(pub_bytes)


def load_private_key(path: Path) -> Ed25519PrivateKey:
    key = serialization.load_pem_private_key(path.read_bytes(), password=None)
    if not isinstance(key, Ed25519PrivateKey):
        raise ValueError(f"'{path}' is not an Ed25519 private key.")
    return key


def load_public_key(path: Path) -> Ed25519PublicKey:
    key = serialization.load_pem_public_key(path.read_bytes())
    if not isinstance(key, Ed25519PublicKey):
        raise ValueError(f"'{path}' is not an Ed25519 public key.")
    return key


# --------------------------------------------------------------------------
# Plugin hashing + entry signing
# --------------------------------------------------------------------------

@dataclass
class PluginEntry:
    name: str          # plugin folder name
    module: str        # module name, from the manifest
    cls: str           # class name, from the manifest
    sha256: str        # hash of manifest + module file contents
    signed_at: str      # ISO-8601 UTC timestamp
    signature: str = ""  # base64 Ed25519 signature, filled in after signing

    def canonical_bytes(self) -> bytes:
        """Bytes that get signed / verified -- everything except the
        signature itself, in a stable, sorted-key JSON encoding."""
        payload = {
            "name": self.name,
            "module": self.module,
            "cls": self.cls,
            "sha256": self.sha256,
            "signed_at": self.signed_at,
        }
        return json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")


def hash_plugin_files(plugin_dir: Path, name: str, module: str) -> str:
    """Hash the manifest + module file together, in a fixed order, so the
    hash changes if either file is tampered with."""
    manifest_path = plugin_dir / f"{name}.json"
    module_path = plugin_dir / f"{module}.py"
    if not manifest_path.is_file():
        raise FileNotFoundError(f"Manifest not found: {manifest_path}")
    if not module_path.is_file():
        raise FileNotFoundError(f"Module file not found: {module_path}")

    h = hashlib.sha256()
    for p in (manifest_path, module_path):
        h.update(p.name.encode("utf-8"))
        h.update(b"\x00")
        h.update(p.read_bytes())
        h.update(b"\x00")
    return h.hexdigest()


def sign_plugin(private_key: Ed25519PrivateKey, plugin_dir: Path) -> PluginEntry:
    manifest_path_glob = list(plugin_dir.glob("*.json"))
    if len(manifest_path_glob) != 1:
        raise ValueError(
            f"Expected exactly one manifest .json in '{plugin_dir}', found {len(manifest_path_glob)}."
        )
    manifest = json.loads(manifest_path_glob[0].read_text(encoding="utf-8"))
    module = manifest.get("module")
    cls = manifest.get("class")
    if not module or not cls:
        raise ValueError("Manifest is missing 'module' or 'class'.")

    name = plugin_dir.name
    digest = hash_plugin_files(plugin_dir, name, module)
    entry = PluginEntry(
        name=name,
        module=module,
        cls=cls,
        sha256=digest,
        signed_at=datetime.now(timezone.utc).isoformat(timespec="seconds"),
    )
    signature = private_key.sign(entry.canonical_bytes())
    entry.signature = base64.b64encode(signature).decode("ascii")
    return entry


def verify_entry(public_key: Ed25519PublicKey, entry: PluginEntry) -> bool:
    try:
        public_key.verify(base64.b64decode(entry.signature), entry.canonical_bytes())
        return True
    except (InvalidSignature, ValueError):
        return False


def verify_plugin_on_disk(public_key: Ed25519PublicKey, plugin_dir: Path, entry: PluginEntry) -> bool:
    """Full verification: signature is valid AND the files on disk still
    match the hash that was signed."""
    if not verify_entry(public_key, entry):
        return False
    try:
        current_hash = hash_plugin_files(plugin_dir, entry.name, entry.module)
    except FileNotFoundError:
        return False
    return current_hash == entry.sha256


# --------------------------------------------------------------------------
# Whitelist storage (optionally encrypted at rest)
# --------------------------------------------------------------------------

class Whitelist:
    def __init__(self):
        self.entries: dict[str, PluginEntry] = {}

    def add(self, entry: PluginEntry) -> None:
        self.entries[entry.name] = entry

    def remove(self, name: str) -> bool:
        return self.entries.pop(name, None) is not None

    def to_json(self) -> str:
        payload = {name: asdict(e) for name, e in self.entries.items()}
        return json.dumps(payload, indent=2, sort_keys=True)

    @classmethod
    def from_json(cls, data: str) -> "Whitelist":
        wl = cls()
        raw = json.loads(data)
        for name, fields in raw.items():
            wl.entries[name] = PluginEntry(**fields)
        return wl

    # -- disk I/O -----------------------------------------------------

    def save(self, path: Path, passphrase: Optional[str] = None) -> None:
        plaintext = self.to_json().encode("utf-8")
        if passphrase is None:
            path.write_bytes(plaintext)
            return

        salt = os.urandom(16)
        key = _derive_key(passphrase, salt)
        token = Fernet(key).encrypt(plaintext)
        container = {
            "encrypted": True,
            "kdf": "pbkdf2-sha256",
            "iterations": PBKDF2_ITERATIONS,
            "salt": base64.b64encode(salt).decode("ascii"),
            "data": base64.b64encode(token).decode("ascii"),
        }
        path.write_text(json.dumps(container, indent=2), encoding="utf-8")

    @classmethod
    def load(cls, path: Path, passphrase: Optional[str] = None) -> "Whitelist":
        raw = path.read_text(encoding="utf-8")
        try:
            container = json.loads(raw)
        except json.JSONDecodeError:
            container = None

        if isinstance(container, dict) and container.get("encrypted"):
            if passphrase is None:
                raise ValueError("Whitelist is encrypted; a passphrase is required.")
            salt = base64.b64decode(container["salt"])
            key = _derive_key(passphrase, salt, container.get("iterations", PBKDF2_ITERATIONS))
            token = base64.b64decode(container["data"])
            try:
                plaintext = Fernet(key).decrypt(token).decode("utf-8")
            except InvalidToken as e:
                raise ValueError("Wrong passphrase or corrupted whitelist file.") from e
            return cls.from_json(plaintext)

        return cls.from_json(raw)


def _derive_key(passphrase: str, salt: bytes, iterations: int = PBKDF2_ITERATIONS) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=iterations,
    )
    return base64.urlsafe_b64encode(kdf.derive(passphrase.encode("utf-8")))