package com.gnx.telemedicine.service;

import com.gnx.telemedicine.model.TwoFactorAuth;
import com.gnx.telemedicine.model.UserModel;
import com.gnx.telemedicine.repository.TwoFactorAuthRepository;
import com.gnx.telemedicine.repository.UserRepository;
import dev.samstevens.totp.code.*;
import dev.samstevens.totp.exceptions.QrGenerationException;
import dev.samstevens.totp.qr.QrData;
import dev.samstevens.totp.qr.QrGenerator;
import dev.samstevens.totp.qr.ZxingPngQrGenerator;
import dev.samstevens.totp.secret.DefaultSecretGenerator;
import dev.samstevens.totp.secret.SecretGenerator;
import dev.samstevens.totp.time.SystemTimeProvider;
import dev.samstevens.totp.time.TimeProvider;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.*;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Slf4j
public class TwoFactorAuthService {
    
    private final TwoFactorAuthRepository twoFactorAuthRepository;
    private final UserRepository userRepository;
    private final SecretGenerator secretGenerator = new DefaultSecretGenerator();
    private final QrGenerator qrGenerator = new ZxingPngQrGenerator();
    private final TimeProvider timeProvider = new SystemTimeProvider();
    private final CodeGenerator codeGenerator = new DefaultCodeGenerator();
    private final CodeVerifier codeVerifier = new DefaultCodeVerifier(codeGenerator, timeProvider);

    @Transactional
    public TwoFactorAuth setup2FA(String userEmail) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElse(new TwoFactorAuth());

        if (twoFactorAuth.getSecretKey() == null) {
            // Generate new secret key
            String secret = secretGenerator.generate();
            twoFactorAuth.setSecretKey(secret);
            
            // Generate backup codes
            List<String> backupCodes = generateBackupCodes();
            twoFactorAuth.setBackupCodes(backupCodes);
            
            twoFactorAuth.setUser(user);
            twoFactorAuth.setEnabled(false);
        }

        twoFactorAuth = twoFactorAuthRepository.save(twoFactorAuth);
        return twoFactorAuth;
    }

    @Transactional
    public void enable2FA(String userEmail, String code) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElseThrow(() -> new IllegalArgumentException("2FA not set up. Please set up first."));

        if (!verifyCode(twoFactorAuth.getSecretKey(), code)) {
            throw new IllegalArgumentException("Invalid verification code");
        }

        twoFactorAuth.setEnabled(true);
        twoFactorAuthRepository.save(twoFactorAuth);
        log.info("2FA enabled for user: {}", userEmail);
    }

    @Transactional
    public void disable2FA(String userEmail, String code) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElseThrow(() -> new IllegalArgumentException("2FA not enabled"));

        // Try TOTP code first, then backup code
        boolean isValidCode = verifyCode(twoFactorAuth.getSecretKey(), code);
        if (!isValidCode) {
            // If TOTP code fails, try backup code
            isValidCode = verifyBackupCode(userEmail, code);
        }

        if (!isValidCode) {
            throw new IllegalArgumentException("Invalid verification code");
        }

        twoFactorAuth.setEnabled(false);
        twoFactorAuth.setBackupCodes(null);
        twoFactorAuthRepository.save(twoFactorAuth);
        log.info("2FA disabled for user: {}", userEmail);
    }

    public boolean verifyCode(String secretKey, String code) {
        if (code == null || code.isEmpty()) {
            return false;
        }

        // Try TOTP code
        try {
            if (codeVerifier.isValidCode(secretKey, code)) {
                return true;
            }
        } catch (Exception e) {
            log.debug("Error verifying TOTP code: {}", e.getMessage());
        }

        return false;
    }

    public boolean verifyBackupCode(String userEmail, String code) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElse(null);

        if (twoFactorAuth == null || twoFactorAuth.getBackupCodes() == null) {
            return false;
        }

        List<String> backupCodes = twoFactorAuth.getBackupCodes();
        if (backupCodes.contains(code)) {
            // Remove used backup code
            backupCodes.remove(code);
            twoFactorAuth.setBackupCodes(backupCodes);
            twoFactorAuthRepository.save(twoFactorAuth);
            return true;
        }

        return false;
    }

    @Transactional
    public List<String> regenerateBackupCodes(String userEmail, String code) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElseThrow(() -> new IllegalArgumentException("2FA not set up"));

        if (!verifyCode(twoFactorAuth.getSecretKey(), code)) {
            throw new IllegalArgumentException("Invalid verification code");
        }

        List<String> backupCodes = generateBackupCodes();
        twoFactorAuth.setBackupCodes(backupCodes);
        twoFactorAuthRepository.save(twoFactorAuth);
        return backupCodes;
    }

    public boolean is2FAEnabled(String userEmail) {
        return userRepository.findByEmail(userEmail)
                .map(user -> twoFactorAuthRepository.findByUserId(user.getId())
                        .map(TwoFactorAuth::getEnabled)
                        .orElse(false))
                .orElse(false);
    }

    public String generateQrCodeUrl(String userEmail) {
        UserModel user = userRepository.findByEmail(userEmail)
                .orElseThrow(() -> new RuntimeException("User not found"));

        TwoFactorAuth twoFactorAuth = twoFactorAuthRepository.findByUserId(user.getId())
                .orElseThrow(() -> new IllegalArgumentException("2FA not set up"));

        String issuer = "Telemedicine API";
        String accountName = user.getEmail();

        QrData data = new QrData.Builder()
                .label(accountName)
                .secret(twoFactorAuth.getSecretKey())
                .issuer(issuer)
                .algorithm(HashingAlgorithm.SHA256)
                .digits(6)
                .period(30)
                .build();

        try {
            byte[] qrCodeBytes = qrGenerator.generate(data);
            String base64 = java.util.Base64.getEncoder().encodeToString(qrCodeBytes);
            return "data:image/png;base64," + base64;
        } catch (QrGenerationException e) {
            throw new RuntimeException("Failed to generate QR code", e);
        }
    }

    private List<String> generateBackupCodes() {
        List<String> codes = new ArrayList<>();
        Random random = new Random();
        
        for (int i = 0; i < 10; i++) {
            // Generate 8-digit backup code
            int code = 10000000 + random.nextInt(90000000);
            codes.add(String.valueOf(code));
        }
        
        return codes;
    }

    public Optional<TwoFactorAuth> get2FASetup(String userEmail) {
        return userRepository.findByEmail(userEmail)
                .flatMap(user -> twoFactorAuthRepository.findByUserId(user.getId()));
    }
}

