package com.alttd.altitudeweb.services.limits;

import com.alttd.altitudeweb.controllers.data_from_auth.AuthenticatedUuid;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.time.Duration;
import java.util.Optional;
import java.util.UUID;

@Aspect
@Component
@RequiredArgsConstructor
@Slf4j
public class RateLimitAspect {
    
    private final InMemoryRateLimiterService rateLimiterService;
    private final AuthenticatedUuid authenticatedUuid;
    
    @Around("""
            @annotation(com.alttd.altitudeweb.services.limits.RateLimit)
            || @within(com.alttd.altitudeweb.services.limits.RateLimit)""")
    public Object rateLimit(ProceedingJoinPoint joinPoint) throws Throwable {
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (requestAttributes == null) {
            return joinPoint.proceed();
        }
        
        HttpServletRequest request = requestAttributes.getRequest();
        HttpServletResponse response = requestAttributes.getResponse();

        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit == null) {
            rateLimit = method.getDeclaringClass().getAnnotation(RateLimit.class);
        }

        if (rateLimit == null) {
            return joinPoint.proceed();
        }

        int limit = rateLimit.limit();
        Duration duration = Duration.ofSeconds(rateLimit.timeUnit().toSeconds(rateLimit.timeValue()));
        String customKey = rateLimit.key();

        Optional<UUID> optionalUUID = authenticatedUuid.tryGetAuthenticatedUserUuid();
        if (optionalUUID.isEmpty()) {
            return joinPoint.proceed();
        }
        UUID uuid = optionalUUID.get();
        String key = uuid + "-" + (customKey.isEmpty() ? method.getName() : customKey);

        boolean allowed = rateLimiterService.tryAcquire(key, limit, duration);
        
        if (allowed) {
            if (response != null) {
                response.setHeader("X-Rate-Limit-Limit", String.valueOf(limit));
                response.setHeader("X-Rate-Limit-Remaining", 
                        String.valueOf(rateLimiterService.getRemainingTokens(key, limit, duration)));
            }
            
            return joinPoint.proceed();
        } else {
            log.warn("Rate limit exceeded for uuid: {}, endpoint: {}", uuid, request.getRequestURI());

            Duration nextResetTime = rateLimiterService.getNextResetTime(key, duration);

            return ResponseEntity
                    .status(HttpStatus.TOO_MANY_REQUESTS)
                    .header("X-Rate-Limit-Limit", String.valueOf(limit))
                    .header("X-Rate-Limit-Remaining", "0")
                    .header("Retry-After", String.valueOf(nextResetTime.getSeconds()))
                    .header("Access-Control-Expose-Headers", "Retry-After, X-Rate-Limit-Limit, X-Rate-Limit-Remaining")
                    .body(String.format("Rate limit exceeded. Try again in %d seconds.", nextResetTime.getSeconds()));
        }
    }
}
