<?php

declare(strict_types=1);

namespace Siteworxpro\App\Http\Middleware;

use Carbon\Carbon;
use Carbon\WrapperClock;
use GuzzleHttp\Exception\GuzzleException;
use Lcobucci\JWT\JwtFacade;
use Lcobucci\JWT\Signer\Hmac\Sha256 as Hmac256;
use Lcobucci\JWT\Signer\Key;
use Lcobucci\JWT\Signer\Key\InMemory;
use Lcobucci\JWT\Signer\Rsa\Sha256;
use Lcobucci\JWT\Token\InvalidTokenStructure;
use Lcobucci\JWT\Validation\Constraint\IssuedBy;
use Lcobucci\JWT\Validation\Constraint\LooseValidAt;
use Lcobucci\JWT\Validation\Constraint\PermittedFor;
use Lcobucci\JWT\Validation\Constraint\SignedWith;
use Lcobucci\JWT\Validation\Constraint\StrictValidAt;
use Lcobucci\JWT\Validation\RequiredConstraintsViolated;
use League\Route\Dispatcher;
use Psr\Http\Message\ResponseInterface;
use Psr\Http\Message\ServerRequestInterface;
use Psr\Http\Server\RequestHandlerInterface;
use Siteworxpro\App\Attributes\Guards\Jwt;
use Siteworxpro\App\Controllers\Controller;
use Siteworxpro\App\Http\JsonResponseFactory;
use Siteworxpro\App\Services\Facades\Config;
use Siteworxpro\App\Services\Facades\Guzzle;
use Siteworxpro\App\Services\Facades\Redis;
use Siteworxpro\HttpStatus\CodesEnum;

/**
 * JWT authorization middleware.
 *
 * Applies JWT validation to controller actions annotated with `Jwt` attribute.
 * Flow:
 * - Resolve the targeted controller and method for the current route.
 * - If the method has `Jwt`, read the `Authorization` header and parse the Bearer token.
 * - Validate signature, time constraints, issuer\(\) and audience\(\) based on attribute and config.
 * - On success, attach all token claims to the request as attributes.
 * - On failure, return a 401 JSON response with validation errors.
 *
 * Configuration:
 * - `jwt.signing_key`: key material or `file://` path to key.
 * - `jwt.strict_validation`: bool toggling strict vs loose time validation.
 */
class JwtMiddleware extends Middleware
{
    /**
     * Process the incoming request.
     *
     * If the matched controller method is annotated with `Jwt`, validates the token and
     * augments the request with claims on success. Otherwise, just delegates to the next handler.
     *
     * @param ServerRequestInterface $request PSR-7 request instance.
     * @param RequestHandlerInterface|Dispatcher $handler Next middleware or route dispatcher.
     *
     * @return ResponseInterface Response produced by the next handler or a 401 JSON response.
     *
     * @throws \JsonException On JSON error response encoding issues.
     * @throws \Exception On unexpected reflection or JWT parsing issues.
     */
    public function process(
        ServerRequestInterface $request,
        RequestHandlerInterface|Dispatcher $handler
    ): ResponseInterface {

        // Resolve the callable \[Controller, method] for the current route.
        $callable = $this->extractRouteCallable($handler);
        if ($callable === null) {
            return $handler->handle($request);
        }

        /** @var Controller $class */
        [$class, $method] = $callable;

        if (class_exists($class::class)) {
            $reflectionClass = new \ReflectionClass($class);

            if ($reflectionClass->hasMethod($method)) {
                $reflectionMethod = $reflectionClass->getMethod($method);
                // Read `Jwt` attribute on the controller method.
                $attributes = $reflectionMethod->getAttributes(Jwt::class);

                // If no `Jwt` attribute, do not enforce auth here.
                if (empty($attributes)) {
                    return $handler->handle($request);
                }

                // Extract Bearer token from Authorization header.
                $token = str_replace('Bearer ', '', $request->getHeaderLine('Authorization'));

                if (empty($token)) {
                    return JsonResponseFactory::createJsonResponse([
                        'status_code' => 401,
                        'message' => 'Unauthorized: Missing token',
                    ], CodesEnum::UNAUTHORIZED);
                }

                // Aggregate required issuers and audience from attributes.
                $requiredIssuers = [];
                $requiredAudience = '';

                foreach ($attributes as $attribute) {
                    /** @var Jwt $jwtInstance */
                    $jwtInstance = $attribute->newInstance();

                    if ($jwtInstance->getAudience() !== '') {
                        $requiredAudience = $jwtInstance->getAudience();
                    }

                    $requiredIssuers[] = $jwtInstance->getIssuer();
                }

                try {
                    // Parse and validate the token with signature, time, issuer and audience constraints.
                    $jwt = new JwtFacade()->parse(
                        $token,
                        $this->getSignedWith($token),
                        Config::get('jwt.strict_validation') ?
                            new StrictValidAt(new WrapperClock(Carbon::now())) :
                            new LooseValidAt(new WrapperClock(Carbon::now())),
                        new IssuedBy(...$requiredIssuers),
                        new PermittedFor($requiredAudience)
                    );
                } catch (RequiredConstraintsViolated $exception) {
                    // Collect human-readable violations to return to the client.
                    $violations = [];
                    foreach ($exception->violations() as $violation) {
                        $violations[] = $violation->getMessage();
                    }

                    return JsonResponseFactory::createJsonResponse([
                        'status_code' => CodesEnum::UNAUTHORIZED->value,
                        'message' => 'Unauthorized: Invalid token',
                        'errors' => $violations
                    ], CodesEnum::UNAUTHORIZED);
                } catch (InvalidTokenStructure) {
                    // Token could not be parsed due to malformed structure.
                    return JsonResponseFactory::createJsonResponse([
                        'status_code' => CodesEnum::UNAUTHORIZED->value,
                        'message' => 'Unauthorized: Invalid token',
                    ], CodesEnum::UNAUTHORIZED);
                } catch (GuzzleException | \RuntimeException) {
                    return JsonResponseFactory::createJsonResponse([
                        'status_code' => CodesEnum::INTERNAL_SERVER_ERROR->value,
                        'message' => 'Token validation service unavailable or unknown error',
                    ], CodesEnum::INTERNAL_SERVER_ERROR);
                }

                // Expose all token claims as request attributes for downstream consumers.
                foreach ($jwt->claims()->all() as $item => $value) {
                    $request = $request->withAttribute($item, $value);
                }
            }
        }

        return $handler->handle($request);
    }

    /**
     * Build the signature validation constraint from configured key.
     *
     * - If the configured key content includes the string `PUBLIC KEY`, use RSA SHA-256.
     * - Otherwise assume an HMAC SHA-256 shared secret.
     * - Supports raw key strings or `file://` paths.
     *
     * @return SignedWith Signature constraint used during JWT parsing.
     *
     * @throws \RuntimeException When no signing key is configured.
     * @throws \JsonException
     */
    private function getSignedWith(string $token): SignedWith
    {
        $keyConfig = Config::get('jwt.signing_key');

        if ($keyConfig === null) {
            throw new \RuntimeException('JWT signing key is not configured.');
        }

        // file:// path to key
        if (str_starts_with($keyConfig, 'file://')) {
            $key = InMemory::file(substr($keyConfig, 7));
            // openid jwks url
        } elseif (str_contains($keyConfig, '.well-known/')) {
            $jwt = explode('.', $token);
            if (count($jwt) !== 3) {
                throw new InvalidTokenStructure('Invalid JWT structure for JWKS key retrieval.');
            }
            $header = json_decode(base64_decode($jwt[0]), true, 512, JSON_THROW_ON_ERROR);
            $keyId = $header['kid'] ?? '0'; // Default to '0' if no kid present
            $key = $this->getJwksKey($keyConfig, $keyId);
        } else {
            $key = InMemory::plainText($keyConfig);
        }

        // Heuristic: if PEM public key content is detected, use RSA; otherwise use HMAC.
        if (str_contains($key->contents(), 'PUBLIC KEY')) {
            return new SignedWith(new Sha256(), $key);
        }

        return new SignedWith(new Hmac256(), $key);
    }

    private function getJwksKey(string $url, string $keyId): Key
    {
        $cached = Redis::get('jwks_key_' . $keyId);
        if ($cached !== null) {
            return InMemory::plainText($cached);
        }

        $openIdConfig = Guzzle::get($url);
        $body = json_decode($openIdConfig->getBody()->getContents(), true, JSON_THROW_ON_ERROR);
        $jwksUri = $body['jwks_uri'] ?? '';
        if (empty($jwksUri)) {
            throw new \RuntimeException('JWKS URI not found in OpenID configuration.');
        }

        $jwksResponse = Guzzle::get($jwksUri);
        $jwksBody = json_decode(
            $jwksResponse->getBody()->getContents(),
            true,
            JSON_THROW_ON_ERROR
        );

        // For simplicity, we take the first key in the JWKS.
        $firstKey = array_filter(
            $jwksBody['keys'],
            fn($key) => $key['kid'] === $keyId
        )[0] ?? $jwksBody['keys'][0] ?? null;

        if (empty($firstKey)) {
            throw new \RuntimeException('No matching key found in JWKS for key ID: ' . $keyId);
        }

        $n = $firstKey['n'];
        $e = $firstKey['e'];
        $publicKeyPem = "-----BEGIN PUBLIC KEY-----\n" .
            chunk_split(base64_encode($this->convertJwkToPem($n, $e)), 64) .
            "-----END PUBLIC KEY-----\n";

        Redis::set('jwks_key_' . $keyId, $publicKeyPem, 'EX', 3600);

        return InMemory::plainText($publicKeyPem);
    }

    /**
     * Build a DER-encoded SubjectPublicKeyInfo from JWK 'n' and 'e'.
     * Returns raw DER bytes; caller base64-encodes and wraps with PEM headers.
     */
    private function convertJwkToPem(string $n, string $e): string
    {
        $modulus = $this->base64UrlDecode($n);
        $exponent = $this->base64UrlDecode($e);

        $derN = $this->derEncodeInteger($modulus);
        $derE = $this->derEncodeInteger($exponent);

        // RSAPublicKey (PKCS#1): SEQUENCE { n INTEGER, e INTEGER }
        $rsaPublicKey = $this->derEncodeSequence($derN . $derE);

        // AlgorithmIdentifier for rsaEncryption: 1.2.840.113549.1.1.1 with NULL
        $algId = hex2bin('300d06092a864886f70d0101010500');

        // SubjectPublicKey (SPKI) BIT STRING, 0 unused bits + RSAPublicKey
        $subjectPublicKey = $this->derEncodeBitString($rsaPublicKey);

        // SubjectPublicKeyInfo: SEQUENCE { algorithm AlgorithmIdentifier, subjectPublicKey BIT STRING }
        return $this->derEncodeSequence($algId . $subjectPublicKey);
    }

    private function base64UrlDecode(string $data): string
    {
        $data = strtr($data, '-_', '+/');
        $pad = strlen($data) % 4;
        if ($pad) {
            $data .= str_repeat('=', 4 - $pad);
        }
        return base64_decode($data);
    }

    private function derEncodeLength(int $len): string
    {
        if ($len < 0x80) {
            return chr($len);
        }
        $bytes = '';
        while ($len > 0) {
            $bytes = chr($len & 0xFF) . $bytes;
            $len >>= 8;
        }
        return chr(0x80 | strlen($bytes)) . $bytes;
    }

    private function derEncodeInteger(string $bytes): string
    {
        // Remove leading zeroes
        $bytes = ltrim($bytes, "\x00");
        if ($bytes === '') {
            $bytes = "\x00";
        }
        // Ensure positive INTEGER (prepend 0x00 if MSB set)
        if ((ord($bytes[0]) & 0x80) !== 0) {
            $bytes = "\x00" . $bytes;
        }
        return "\x02" . $this->derEncodeLength(strlen($bytes)) . $bytes;
    }

    private function derEncodeSequence(string $bytes): string
    {
        return "\x30" . $this->derEncodeLength(strlen($bytes)) . $bytes;
    }

    private function derEncodeBitString(string $bytes): string
    {
        // 0 unused bits + data
        $payload = "\x00" . $bytes;
        return "\x03" . $this->derEncodeLength(strlen($payload)) . $payload;
    }
}