package cloud.nstart.terravault.ordinis.app.e2e;

import com.nimbusds.jose.JOSEObjectType;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.crypto.RSASSASigner;
import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import com.sun.net.httpserver.HttpServer;

import java.io.IOException;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.time.Duration;
import java.time.Instant;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * Тестовый суррогат Keycloak'а: генерирует RSA keypair, поднимает HTTP-сервер
 * на random port отдающий JWK Set, подписывает JWT с произвольными claim'ами.
 *
 * <p>Используется в e2e auth тестах. Spring Security настраивается на
 * {@code ordinis.auth.jwk-set-uri=<server.url>} и доверяет наши test-JWT.
 */
public final class JwtTestSupport {

  private static final RSAKey RSA_KEY;
  private static final HttpServer JWKS_SERVER;
  private static final String JWKS_URL;

  static {
    try {
      KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
      kpg.initialize(2048);
      KeyPair pair = kpg.generateKeyPair();
      RSA_KEY = new RSAKey.Builder((RSAPublicKey) pair.getPublic())
          .privateKey((RSAPrivateKey) pair.getPrivate())
          .keyID("test-key-1")
          .build();

      String jwksJson = new JWKSet(RSA_KEY.toPublicJWK()).toString();

      JWKS_SERVER = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
      JWKS_SERVER.createContext("/jwks", exchange -> {
        byte[] body = jwksJson.getBytes();
        exchange.getResponseHeaders().add("Content-Type", "application/json");
        exchange.sendResponseHeaders(200, body.length);
        try (OutputStream os = exchange.getResponseBody()) {
          os.write(body);
        }
      });
      JWKS_SERVER.start();

      int port = JWKS_SERVER.getAddress().getPort();
      JWKS_URL = "http://127.0.0.1:" + port + "/jwks";
    } catch (Exception e) {
      throw new ExceptionInInitializerError(e);
    }
  }

  private JwtTestSupport() {}

  public static String jwksUrl() {
    return JWKS_URL;
  }

  /**
   * Подписывает JWT с указанными realm-ролями. Claim {@code realm_access.roles}
   * — Keycloak-стиль.
   */
  public static String signJwt(List<String> realmRoles, String subject) {
    return signJwt(realmRoles, subject, Duration.ofMinutes(30));
  }

  public static String signJwt(List<String> realmRoles, String subject, Duration ttl) {
    try {
      Instant now = Instant.now();
      JWTClaimsSet claims = new JWTClaimsSet.Builder()
          .subject(subject)
          .issuer("https://test-issuer/realms/nstart")
          .issueTime(Date.from(now))
          .expirationTime(Date.from(now.plus(ttl)))
          .jwtID(UUID.randomUUID().toString())
          .claim("preferred_username", subject)
          .claim("realm_access", Map.of("roles", realmRoles))
          .build();

      SignedJWT jwt = new SignedJWT(
          new JWSHeader.Builder(JWSAlgorithm.RS256)
              .keyID(RSA_KEY.getKeyID())
              .type(JOSEObjectType.JWT)
              .build(),
          claims);
      jwt.sign(new RSASSASigner(RSA_KEY));
      return jwt.serialize();
    } catch (Exception e) {
      throw new RuntimeException("Failed to sign test JWT", e);
    }
  }
}
