"""
signed_tool — dekorator @signed dla serwerów MCP.

Użycie:
    from mcp_helpers import SignedTool, signed

    # przy starcie serwera:
    SignedTool.load_keys(config["authorized_keys"])
    # {"alice": "base64_klucz_publiczny_ed25519"}

    # dekorowanie narzędzia:
    @signed("alice")
    async def update_code(path: str, content: str) -> str:
        ...

Dekorator automatycznie:
- dodaje pola `timestamp: int` i `signature: str` do sygnatury funkcji
- dodaje do docstringa instrukcję dla modelu
- weryfikuje okno czasowe (domyślnie ±5 minut)
- weryfikuje podpis Ed25519
- wywołuje oryginalną funkcję bez timestamp/signature

Wiadomość do podpisania:
    "{tool_name}:{timestamp}:{sha256(json(args))}"
    gdzie args to argumenty bez pól timestamp i signature,
    JSON deterministyczny (sort_keys=True, bez spacji).

Wymagania: cryptography
"""
from __future__ import annotations

import base64
import functools
import hashlib
import inspect
import json
import time
from typing import Any, Callable

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey


class SignedTool:
    _keys: dict[str, Ed25519PublicKey] = {}

    @classmethod
    def load_keys(cls, authorized_keys: dict[str, str]) -> None:
        """Ładuje klucze publiczne Ed25519.

        Args:
            authorized_keys: słownik {nazwa: base64(raw 32-bajtowy klucz publiczny Ed25519)}
        """
        cls._keys = {}
        for name, b64_key in authorized_keys.items():
            raw = base64.b64decode(b64_key)
            cls._keys[name] = Ed25519PublicKey.from_public_bytes(raw)

    @classmethod
    def get_key(cls, name: str) -> Ed25519PublicKey:
        if name not in cls._keys:
            raise ValueError(f"Nieznany klucz: {name!r}. Załadowane klucze: {list(cls._keys)}")
        return cls._keys[name]


def signed(key_name: str, window_seconds: int = 300) -> Callable:
    """Dekorator wymagający podpisu Ed25519 przy wywołaniu narzędzia MCP.

    Args:
        key_name:       nazwa klucza z authorized_keys (np. "alice")
        window_seconds: dopuszczalne odchylenie timestampu w sekundach (domyślnie 300 = ±5 min)
    """
    def decorator(fn: Callable) -> Callable:
        tool_name = fn.__name__

        @functools.wraps(fn)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            try:
                timestamp: int = kwargs.pop("timestamp")
                signature_b64: str = kwargs.pop("signature")
            except KeyError as e:
                raise ValueError(f"Brakujące pole wymagane przez @signed: {e}") from e

            # weryfikacja okna czasowego
            now = int(time.time())
            if abs(now - timestamp) > window_seconds:
                raise ValueError(
                    f"Timestamp poza oknem czasowym (±{window_seconds}s). "
                    f"Teraz: {now}, timestamp: {timestamp}, różnica: {abs(now - timestamp)}s"
                )

            # budowanie wiadomości
            args_hash = _hash_args(kwargs)
            message = f"{tool_name}:{timestamp}:{args_hash}".encode()

            # weryfikacja podpisu
            public_key = SignedTool.get_key(key_name)
            try:
                sig_bytes = base64.b64decode(signature_b64)
            except Exception as e:
                raise ValueError(f"Nieprawidłowy format base64 podpisu: {e}") from e

            try:
                public_key.verify(sig_bytes, message)
            except InvalidSignature:
                raise ValueError(
                    f"Nieprawidłowy podpis Ed25519 dla klucza {key_name!r}. "
                    f"Podpisywana wiadomość: {message.decode()!r}"
                )

            return await fn(*args, **kwargs)

        # rozszerzenie sygnatury o timestamp i signature
        orig_sig = inspect.signature(fn)
        extra = [
            inspect.Parameter("timestamp", inspect.Parameter.KEYWORD_ONLY, annotation=int),
            inspect.Parameter("signature", inspect.Parameter.KEYWORD_ONLY, annotation=str),
        ]
        wrapper.__signature__ = orig_sig.replace(parameters=list(orig_sig.parameters.values()) + extra)
        wrapper.__annotations__ = {**fn.__annotations__, "timestamp": int, "signature": str}

        # rozszerzenie docstringa
        sign_note = (
            f"\n\n[WYMAGA PODPISU od: {key_name!r}]\n"
            f"Przed wywołaniem poproś użytkownika o podpis Ed25519:\n"
            f"  1. Zbierz argumenty wywołania (bez timestamp i signature)\n"
            f"  2. Oblicz: sha256 = SHA-256(JSON(args, sort_keys=True, bez spacji))\n"
            f"  3. Utwórz wiadomość: \"{tool_name}:{{timestamp}}:{{sha256}}\"\n"
            f"     gdzie timestamp to bieżący czas Unix (int)\n"
            f"  4. Poproś użytkownika o podpisanie wiadomości kluczem prywatnym {key_name!r}\n"
            f"  5. Przekaż: timestamp (int) i signature (base64 podpisu)"
        )
        wrapper.__doc__ = (fn.__doc__ or "") + sign_note

        return wrapper
    return decorator


def _hash_args(args: dict) -> str:
    """SHA-256 z deterministycznego JSON argumentów."""
    serialized = json.dumps(args, sort_keys=True, ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha256(serialized.encode()).hexdigest()
