package com.cameleer.server.app.security;

import com.cameleer.server.core.security.OidcConfig;
import com.cameleer.server.core.security.OidcConfigRepository;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.JWSVerificationKeySelector;
import com.nimbusds.jose.proc.SecurityContext;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.proc.ConfigurableJWTProcessor;
import com.nimbusds.jwt.proc.DefaultJWTProcessor;
import com.nimbusds.oauth2.sdk.AuthorizationCode;
import com.nimbusds.oauth2.sdk.AuthorizationCodeGrant;
import com.nimbusds.oauth2.sdk.TokenRequest;
import com.nimbusds.oauth2.sdk.TokenResponse;
import com.nimbusds.oauth2.sdk.auth.ClientAuthentication;
import com.nimbusds.oauth2.sdk.auth.ClientSecretBasic;
import com.nimbusds.oauth2.sdk.auth.Secret;
import com.nimbusds.oauth2.sdk.id.ClientID;
import com.nimbusds.openid.connect.sdk.op.OIDCProviderMetadata;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.net.URI;
import java.net.URL;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * Exchanges OIDC authorization codes for validated user information.
 * <p>
 * Reads OIDC configuration from the database via {@link OidcConfigRepository}.
 * Caches provider metadata and JWKS processor per issuer URI; cache is invalidated
 * when {@link #invalidateCache()} is called (e.g. after config update).
 */
@Component
public class OidcTokenExchanger {

    private static final Logger log = LoggerFactory.getLogger(OidcTokenExchanger.class);

    private final OidcConfigRepository configRepository;
    private final SecurityProperties securityProperties;

    private volatile String cachedIssuerUri;
    private volatile OIDCProviderMetadata providerMetadata;
    private volatile ConfigurableJWTProcessor<SecurityContext> jwtProcessor;
    private volatile ConfigurableJWTProcessor<SecurityContext> accessTokenProcessor;

    public OidcTokenExchanger(OidcConfigRepository configRepository,
                               SecurityProperties securityProperties) {
        this.configRepository = configRepository;
        this.securityProperties = securityProperties;
    }

    public record OidcUserInfo(String subject, String email, String name, List<String> roles, String idToken,
                               Map<String, Object> allClaims) {}

    /**
     * Exchanges an authorization code for validated user info.
     */
    public OidcUserInfo exchange(String code, String redirectUri) throws Exception {
        OidcConfig config = getConfig();

        OIDCProviderMetadata metadata = getProviderMetadata(config.issuerUri());

        ClientAuthentication clientAuth = new ClientSecretBasic(
                new ClientID(config.clientId()),
                new Secret(config.clientSecret())
        );

        AuthorizationCodeGrant grant = new AuthorizationCodeGrant(
                new AuthorizationCode(code), new URI(redirectUri));

        TokenRequest tokenRequest = new TokenRequest(
                metadata.getTokenEndpointURI(),
                clientAuth,
                grant
        );

        log.info("OIDC token exchange: tokenEndpoint={}, redirectUri={}", metadata.getTokenEndpointURI(), redirectUri);

        var httpRequest = tokenRequest.toHTTPRequest();
        // RFC 8707: include resource indicator in token exchange to get a JWT access token
        String configAudience = config.audience() != null ? config.audience() : "";
        if (!configAudience.isBlank()) {
            String body = httpRequest.getBody();
            body += "&resource=" + java.net.URLEncoder.encode(configAudience, java.nio.charset.StandardCharsets.UTF_8);
            httpRequest.setBody(body);
        }
        if (securityProperties.getOidc().isTlsSkipVerify()) {
            httpRequest.setSSLSocketFactory(InsecureTlsHelper.socketFactory());
            httpRequest.setHostnameVerifier(InsecureTlsHelper.hostnameVerifier());
        }
        TokenResponse tokenResponse = TokenResponse.parse(httpRequest.send());

        if (!tokenResponse.indicatesSuccess()) {
            var errorObj = tokenResponse.toErrorResponse().getErrorObject();
            log.error("OIDC token exchange failed: code={}, description={}, httpStatus={}",
                    errorObj.getCode(), errorObj.getDescription(), errorObj.getHTTPStatusCode());
            throw new IllegalStateException("OIDC token exchange failed: " + errorObj.getDescription());
        }

        var responseJson = tokenResponse.toSuccessResponse().toJSONObject();

        String idTokenStr = responseJson.getAsString("id_token");
        if (idTokenStr == null) {
            throw new IllegalStateException("OIDC response missing id_token");
        }
        String accessTokenStr = responseJson.getAsString("access_token");

        JWTClaimsSet claims = getJwtProcessor(config.issuerUri()).process(idTokenStr, null);

        String userIdClaim = config.userIdClaim() != null && !config.userIdClaim().isBlank()
                ? config.userIdClaim() : "sub";
        String subject = "sub".equals(userIdClaim)
                ? claims.getSubject()
                : extractStringClaim(claims, userIdClaim);
        if (subject == null || subject.isBlank()) {
            throw new IllegalStateException("OIDC id_token missing user ID claim: " + userIdClaim);
        }
        String email = claims.getStringClaim("email");
        String name = extractStringClaim(claims, config.displayNameClaim());

        // Try roles from access_token first (JWT providers like Logto, Keycloak),
        // then fall back to id_token
        List<String> roles = Collections.emptyList();
        if (accessTokenStr != null && accessTokenStr.contains(".")) {
            try {
                String audience = config.audience() != null ? config.audience() : "";
                JWTClaimsSet atClaims = decodeAccessToken(accessTokenStr, config.issuerUri(), audience);
                if (atClaims != null) {
                    roles = extractRoles(atClaims, config.rolesClaim());
                    if (!roles.isEmpty()) {
                        log.info("OIDC roles from access_token: {}", roles);
                    }
                }
            } catch (Exception e) {
                log.debug("Could not decode access_token as JWT: {}", e.getMessage());
            }
        }
        if (roles.isEmpty()) {
            roles = extractRoles(claims, config.rolesClaim());
        }

        // Merge id_token and access_token claims for claim mapping evaluation
        Map<String, Object> allClaims = new java.util.HashMap<>(claims.getClaims());
        if (accessTokenStr != null && accessTokenStr.contains(".")) {
            try {
                String audience = config.audience() != null ? config.audience() : "";
                JWTClaimsSet atClaims2 = decodeAccessToken(accessTokenStr, config.issuerUri(), audience);
                if (atClaims2 != null) {
                    // Access token claims take precedence (they have scopes, roles)
                    atClaims2.getClaims().forEach(allClaims::putIfAbsent);
                }
            } catch (Exception e) {
                log.debug("Could not merge access_token claims: {}", e.getMessage());
            }
        }

        log.info("OIDC user authenticated: id={}, email={}, rolesClaim='{}', extractedRoles={}, idTokenClaims={}, hasAccessToken={}",
                subject, email, config.rolesClaim(), roles, claims.getClaims().keySet(), accessTokenStr != null);
        return new OidcUserInfo(subject, email != null ? email : "", name != null ? name : "", roles, idTokenStr, allClaims);
    }

    /**
     * Returns the provider's authorization endpoint for the SPA.
     */
    public String getAuthorizationEndpoint() throws Exception {
        OidcConfig config = getConfig();
        return getProviderMetadata(config.issuerUri()).getAuthorizationEndpointURI().toString();
    }

    /**
     * Returns the provider's end-session (logout) endpoint, or {@code null} if not advertised.
     */
    public String getEndSessionEndpoint() throws Exception {
        OidcConfig config = getConfig();
        URI uri = getProviderMetadata(config.issuerUri()).getEndSessionEndpointURI();
        return uri != null ? uri.toString() : null;
    }

    /**
     * Invalidates cached provider metadata and JWKS processor.
     * Call after OIDC configuration is updated in the database.
     */
    public void invalidateCache() {
        synchronized (this) {
            providerMetadata = null;
            jwtProcessor = null;
            accessTokenProcessor = null;
            cachedIssuerUri = null;
            log.info("OIDC provider cache invalidated");
        }
    }

    /**
     * Returns the current OIDC config from the database.
     *
     * @throws IllegalStateException if OIDC is not configured or disabled
     */
    public OidcConfig getConfig() {
        return configRepository.find()
                .filter(OidcConfig::enabled)
                .orElseThrow(() -> new IllegalStateException("OIDC is not configured or disabled"));
    }

    @SuppressWarnings("unchecked")
    private Object extractClaim(JWTClaimsSet claims, String claimPath) {
        if (claimPath == null || claimPath.isBlank()) {
            return null;
        }
        try {
            String[] parts = claimPath.split("\\.");
            Object current = claims.getClaim(parts[0]);
            for (int i = 1; i < parts.length && current instanceof Map; i++) {
                current = ((Map<String, Object>) current).get(parts[i]);
            }
            return current;
        } catch (Exception e) {
            log.debug("Could not extract claim '{}': {}", claimPath, e.getMessage());
            return null;
        }
    }

    private String extractStringClaim(JWTClaimsSet claims, String claimPath) {
        Object value = extractClaim(claims, claimPath);
        return value instanceof String s ? s : null;
    }

    private List<String> extractRoles(JWTClaimsSet claims, String claimPath) {
        Object value = extractClaim(claims, claimPath);
        if (value instanceof List<?> list) {
            return list.stream().map(Object::toString).toList();
        }
        if (value instanceof String s && !s.isBlank()) {
            return List.of(s.split(" "));
        }
        return Collections.emptyList();
    }

    /**
     * Decodes a JWT access token and validates its audience if configured.
     * Returns the claims, or {@code null} if validation fails.
     */
    private JWTClaimsSet decodeAccessToken(String jwt, String issuerUri, String audience) throws Exception {
        ConfigurableJWTProcessor<SecurityContext> processor = getAccessTokenProcessor(issuerUri);
        JWTClaimsSet claims = processor.process(jwt, null);
        if (audience != null && !audience.isBlank()) {
            java.util.List<String> aud = claims.getAudience();
            if (aud == null || !aud.contains(audience)) {
                log.debug("Access token audience {} does not match expected {}", aud, audience);
                return null;
            }
        }
        return claims;
    }

    /**
     * Returns a JWT processor that accepts {@code at+jwt} type headers (RFC 9068),
     * suitable for validating OAuth2 access tokens.
     */
    private ConfigurableJWTProcessor<SecurityContext> getAccessTokenProcessor(String issuerUri) throws Exception {
        if (accessTokenProcessor == null) {
            synchronized (this) {
                if (accessTokenProcessor == null) {
                    OIDCProviderMetadata metadata = getProviderMetadata(issuerUri);
                    URL jwksUrl = metadata.getJWKSetURI().toURL();
                    JWKSource<SecurityContext> jwkSource = OidcProviderHelper.buildJwkSource(
                            jwksUrl, securityProperties.getOidc().isTlsSkipVerify());
                    var keySelector = new JWSVerificationKeySelector<SecurityContext>(
                            OidcProviderHelper.SUPPORTED_ALGORITHMS, jwkSource);
                    ConfigurableJWTProcessor<SecurityContext> processor = new DefaultJWTProcessor<>();
                    processor.setJWSKeySelector(keySelector);
                    // Accept any JWT type header (at+jwt, JWT, etc.)
                    processor.setJWSTypeVerifier((type, ctx) -> { });
                    accessTokenProcessor = processor;
                }
            }
        }
        return accessTokenProcessor;
    }

    private OIDCProviderMetadata getProviderMetadata(String issuerUri) throws Exception {
        if (providerMetadata == null || !issuerUri.equals(cachedIssuerUri)) {
            synchronized (this) {
                if (providerMetadata == null || !issuerUri.equals(cachedIssuerUri)) {
                    providerMetadata = OidcProviderHelper.fetchMetadata(
                            issuerUri, securityProperties.getOidc().isTlsSkipVerify());
                    cachedIssuerUri = issuerUri;
                    jwtProcessor = null; // Reset processor when issuer changes
                    log.info("OIDC provider metadata loaded from {}", issuerUri);
                }
            }
        }
        return providerMetadata;
    }

    private ConfigurableJWTProcessor<SecurityContext> getJwtProcessor(String issuerUri) throws Exception {
        if (jwtProcessor == null) {
            synchronized (this) {
                if (jwtProcessor == null) {
                    OIDCProviderMetadata metadata = getProviderMetadata(issuerUri);
                    URL jwksUrl = metadata.getJWKSetURI().toURL();
                    JWKSource<SecurityContext> jwkSource = OidcProviderHelper.buildJwkSource(
                            jwksUrl, securityProperties.getOidc().isTlsSkipVerify());
                    var keySelector = new JWSVerificationKeySelector<SecurityContext>(
                            OidcProviderHelper.SUPPORTED_ALGORITHMS, jwkSource);

                    ConfigurableJWTProcessor<SecurityContext> processor = new DefaultJWTProcessor<>();
                    processor.setJWSKeySelector(keySelector);
                    jwtProcessor = processor;
                }
            }
        }
        return jwtProcessor;
    }
}
