package com.cameleer.server.app.security;

import com.cameleer.server.core.security.InvalidTokenException;
import com.cameleer.server.core.security.JwtService;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSHeader;
import com.nimbusds.jose.JWSSigner;
import com.nimbusds.jose.JWSVerifier;
import com.nimbusds.jose.crypto.MACSigner;
import com.nimbusds.jose.crypto.MACVerifier;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.SignedJWT;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.security.SecureRandom;
import java.text.ParseException;
import java.time.Instant;
import java.util.Base64;
import java.util.Date;
import java.util.List;

/**
 * HMAC-SHA256 JWT implementation using Nimbus JOSE+JWT.
 * <p>
 * If {@code security.jwt-secret} is configured, uses that secret (Base64-decoded or raw bytes
 * if not valid Base64). Otherwise generates a random 256-bit secret (ephemeral per server instance).
 * <p>
 * Embeds a {@code roles} array claim in every token for RBAC.
 */
public class JwtServiceImpl implements JwtService {

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

    private final byte[] secret;
    private final JWSSigner signer;
    private final JWSVerifier verifier;
    private final SecurityProperties properties;

    public JwtServiceImpl(SecurityProperties properties) {
        this.properties = properties;

        String configuredSecret = properties.getJwtSecret();
        if (configuredSecret != null && !configuredSecret.isBlank()) {
            this.secret = decodeSecret(configuredSecret);
            log.info("JWT signing secret loaded from configuration");
        } else {
            this.secret = new byte[32]; // 256-bit
            new SecureRandom().nextBytes(secret);
            log.info("JWT signing secret generated randomly (tokens will not survive restarts)");
        }

        try {
            this.signer = new MACSigner(secret);
            this.verifier = new MACVerifier(secret);
        } catch (JOSEException e) {
            throw new IllegalStateException("Failed to initialize JWT signer/verifier", e);
        }
    }

    @Override
    public String createAccessToken(String subject, String application, String environment, List<String> roles) {
        return createToken(subject, application, environment, roles, "access", properties.getAccessTokenExpiryMs());
    }

    @Override
    public String createRefreshToken(String subject, String application, String environment, List<String> roles) {
        return createToken(subject, application, environment, roles, "refresh", properties.getRefreshTokenExpiryMs());
    }

    @Override
    public JwtValidationResult validateAccessToken(String token) {
        return validateToken(token, "access");
    }

    @Override
    public JwtValidationResult validateRefreshToken(String token) {
        return validateToken(token, "refresh");
    }

    @Override
    public String validateAndExtractAgentId(String token) {
        return validateAccessToken(token).subject();
    }

    private String createToken(String subject, String application, String environment,
                               List<String> roles, String type, long expiryMs) {
        Instant now = Instant.now();
        JWTClaimsSet claims = new JWTClaimsSet.Builder()
                .subject(subject)
                .claim("group", application)
                .claim("env", environment)
                .claim("type", type)
                .claim("roles", roles)
                .issueTime(Date.from(now))
                .expirationTime(Date.from(now.plusMillis(expiryMs)))
                .build();

        SignedJWT signedJWT = new SignedJWT(new JWSHeader(JWSAlgorithm.HS256), claims);
        try {
            signedJWT.sign(signer);
        } catch (JOSEException e) {
            throw new IllegalStateException("Failed to sign JWT", e);
        }
        return signedJWT.serialize();
    }

    @SuppressWarnings("unchecked")
    private JwtValidationResult validateToken(String token, String expectedType) {
        try {
            SignedJWT signedJWT = SignedJWT.parse(token);

            if (!signedJWT.verify(verifier)) {
                throw new InvalidTokenException("Invalid JWT signature");
            }

            JWTClaimsSet claims = signedJWT.getJWTClaimsSet();

            Date expiration = claims.getExpirationTime();
            if (expiration == null || expiration.before(new Date())) {
                throw new InvalidTokenException("Token has expired");
            }

            String type = claims.getStringClaim("type");
            if (!expectedType.equals(type)) {
                throw new InvalidTokenException(
                        "Expected token type '" + expectedType + "' but got '" + type + "'");
            }

            String subject = claims.getSubject();
            if (subject == null || subject.isBlank()) {
                throw new InvalidTokenException("Token has no subject");
            }

            String application = claims.getStringClaim("group");

            // Extract roles — may be absent in legacy tokens
            List<String> roles;
            Object rolesClaim = claims.getClaim("roles");
            if (rolesClaim instanceof List<?>) {
                roles = ((List<?>) rolesClaim).stream()
                        .map(Object::toString)
                        .toList();
            } else {
                roles = List.of();
            }

            String environment = claims.getStringClaim("env");

            Date issueTime = claims.getIssueTime();
            Instant issuedAt = issueTime != null ? issueTime.toInstant() : null;

            return new JwtValidationResult(subject, application, environment, roles, issuedAt);
        } catch (ParseException e) {
            throw new InvalidTokenException("Failed to parse JWT", e);
        } catch (JOSEException e) {
            throw new InvalidTokenException("Failed to verify JWT signature", e);
        }
    }

    private static byte[] decodeSecret(String secret) {
        try {
            byte[] decoded = Base64.getDecoder().decode(secret);
            if (decoded.length >= 32) {
                return decoded;
            }
        } catch (IllegalArgumentException ignored) {
            // Not valid Base64 — use raw bytes
        }
        byte[] raw = secret.getBytes(java.nio.charset.StandardCharsets.UTF_8);
        if (raw.length < 32) {
            // Pad to 32 bytes with SHA-256 hash
            try {
                return java.security.MessageDigest.getInstance("SHA-256").digest(raw);
            } catch (java.security.NoSuchAlgorithmException e) {
                throw new IllegalStateException("SHA-256 not available", e);
            }
        }
        return raw;
    }
}
