<?php

namespace io\billdesk\client\hmacsha256;

use Jose\Component\Signature\JWSVerifier;
use Jose\Component\Core\JWK;
use Jose\Component\Core\AlgorithmManager;
use Jose\Component\Encryption\JWEBuilder;
use Jose\Component\Encryption\Serializer\CompactSerializer as JWECompactSerializer;
use Jose\Component\Signature\Serializer\CompactSerializer as JWSCompactSerializer;
use Jose\Component\Encryption\Algorithm\ContentEncryption\A256GCM;
use Jose\Component\Encryption\Algorithm\KeyEncryption\Dir;
use Jose\Component\Encryption\JWEDecrypter;
use Jose\Component\Signature\JWSBuilder;
use Jose\Component\Signature\Algorithm\HS256;
use Psr\Log\LoggerInterface;

class JWESymmetricHelper
{
    private string $encryptionKey;
    private string $clientId;
    private string $encryptionKeyId;
    private LoggerInterface $logger;
    private string $signingKey;
    private string $signingKeyId;

    public function __construct(
        string $encryptionKey,
        string $clientId,
        string $encryptionKeyId,
        LoggerInterface $logger,
        string $signingKey = '',
        string $signingKeyId = ''
    ) {
        $this->encryptionKey = $encryptionKey;
        $this->clientId = $clientId;
        $this->encryptionKeyId = $encryptionKeyId;
        $this->logger = $logger;
        $this->signingKey = $signingKey;
        $this->signingKeyId = $signingKeyId;
    }

    public function encryptPayload(array $payload): string
    {
        $jsonPayload = json_encode($payload, JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE);
        $this->logger->info("Encrypting payload", ['payload' => $jsonPayload]);
        return $this->encrypt($jsonPayload);
    }

    private function encrypt(string $payload): string
    {
        $decodedKey = base64_decode($this->encryptionKey, true);
        if ($decodedKey === false || strlen($decodedKey) !== 32) {
            throw new \RuntimeException("Encryption key must be a valid 32-byte AES key");
        }

        $jwk = new JWK([
            'kty' => 'oct',
            'k' => $this->base64UrlEncode($decodedKey),
        ]);

        $headers = [
            'alg' => 'dir',
            'enc' => 'A256GCM',
            'kid' => $this->encryptionKeyId,
            'clientid' => $this->clientId,
        ];

        try {
            $jweBuilder = new JWEBuilder(
                new AlgorithmManager([new Dir()]),
                new AlgorithmManager([new A256GCM()])
            );

            $jwe = $jweBuilder
                ->create()
                ->withPayload($payload)
                ->withSharedProtectedHeader($headers)
                ->addRecipient($jwk)
                ->build();

            return (new JWECompactSerializer())->serialize($jwe, 0);
        } catch (\Exception $e) {
            $this->logger->error("Encryption failed", ['error' => $e->getMessage()]);
            throw new \RuntimeException("Encryption failed: " . $e->getMessage());
        }
    }

    public function sign(string $encryptedPayload): string
    {
        $decodedKey = base64_decode($this->signingKey, true);
        if ($decodedKey === false || strlen($decodedKey) < 32) {
            throw new \RuntimeException("Invalid signing key");
        }

        $jwk = new JWK([
            'kty' => 'oct',
            'k' => $this->base64UrlEncode($decodedKey),
        ]);

        $header = [
            'alg' => 'HS256',
            'kid' => $this->signingKeyId,
            'clientid' => $this->clientId,
        ];

        try {
            $jwsBuilder = new JWSBuilder(new AlgorithmManager([new HS256()]));

            $jws = $jwsBuilder
                ->create()
                ->withPayload($encryptedPayload)
                ->addSignature($jwk, $header)
                ->build();

            return (new JWSCompactSerializer())->serialize($jws, 0);
        } catch (\Exception $e) {
            $this->logger->error("Signing failed", ['error' => $e->getMessage()]);
            throw new \RuntimeException("Signing failed: " . $e->getMessage());
        }
    }

    public function verifyAndDecrypt(string $jwsString): ?string
    {
        try {
            $jweString = $this->verifyJWS($jwsString); // this is a compact JWE
            if (!$jweString) {
                return null;
            }

            return $this->decryptPayload($jweString); // decrypt normally
        } catch (\Exception $e) {
            $this->logger->error("Verify and decrypt failed", ['error' => $e->getMessage()]);
            return null;
        }
    }

    public function verifyJWS(string $jwsString): ?string
    {
        $keyMaterial = base64_decode($this->signingKey, true);
        if ($keyMaterial === false) {
            $this->logger->error("Invalid signing key encoding");
            return null;
        }

        $jwk = new JWK([
            'kty' => 'oct',
            'k' => $this->base64UrlEncode($keyMaterial),
            'kid' => $this->signingKeyId,
        ]);

        $jwsVerifier = new JWSVerifier(new AlgorithmManager([new HS256()]));

        try {
            $serializer = new JWSCompactSerializer();
            $jws = $serializer->unserialize($jwsString);

            if ($jwsVerifier->verifyWithKey($jws, $jwk, 0)) {
                // ✅ This is a compact JWE string, not JSON!
                return $jws->getPayload();
            }

            $this->logger->error("JWS verification failed");
        } catch (\Exception $e) {
            $this->logger->error("JWS verification error", ['error' => $e->getMessage()]);
        }

        return null;
    }


    private function decryptPayload(string $jweString): ?string
    {
        try {
            $decodedKey = base64_decode($this->encryptionKey);
            if ($decodedKey === false || strlen($decodedKey) !== 32) {
                $this->logger->error("Invalid encryption key");
                return null;
            }

            $parts = explode('.', $jweString);
            if (count($parts) !== 5) {
                $this->logger->error("Invalid JWE format");
                return null;
            }

            $jwk = new JWK([
                'kty' => 'oct',
                'k' => $this->base64UrlEncode($decodedKey),
            ]);

            $jweDecrypter = new JWEDecrypter(
                new AlgorithmManager([new Dir()]),
                new AlgorithmManager([new A256GCM()])
            );

            $serializer = new JWECompactSerializer();
            $jwe = $serializer->unserialize($jweString);

            if (!$jweDecrypter->decryptUsingKey($jwe, $jwk, 0)) {
                $this->logger->error("JWE decryption failed");
                return null;
            }

            $payload = $jwe->getPayload();

            // Optional: validate it's JSON
            json_decode($payload, true);
            if (json_last_error() !== JSON_ERROR_NONE) {
                $this->logger->error("Invalid JSON payload");
                return null;
            }

            return $payload; // ✅ No re-encoding
        } catch (\Exception $e) {
            $this->logger->error("Decryption failed", ['error' => $e->getMessage()]);
            return null;
        }
    }


    private function safeBase64UrlDecode($data)
    {
        if ($data === '') {
            return '';
        }

        $decoded = base64_decode(strtr($data, '-_', '+/'), true);
        return $decoded === false ? false : $decoded;
    }

    private function base64UrlEncode(string $data): string
    {
        return rtrim(strtr(base64_encode($data), '+/', '-_'), '=');
    }
}
