source: backend/web3-Ednevnik-Backend/src/main/java/com/nikola/web3ednevnikbackend/controller/MetamuskController.java

package com.nikola.web3ednevnikbackend.controller;


import com.nikola.web3ednevnikbackend.config.security.CustomUserDetails;
import com.nikola.web3ednevnikbackend.dto.users.BlagajnikResponseDTO;
import com.nikola.web3ednevnikbackend.dto.users.RoditelResponseDTO;
import com.nikola.web3ednevnikbackend.service.users.BlagajnikService;
import com.nikola.web3ednevnikbackend.service.users.RoditelService;

import org.springframework.http.ResponseEntity;
import org.springframework.security.core.Authentication;
import org.springframework.web.bind.annotation.*;
import org.web3j.crypto.*;
import org.web3j.utils.Numeric;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/metamuskAuth")
public class MetamuskController {
    private final Map<String, String> nonces = new ConcurrentHashMap<>();


    private final RoditelService roditelService;
    private final BlagajnikService blagajnikService;

    public MetamuskController(RoditelService roditelService, BlagajnikService blagajnikService) {
        this.roditelService = roditelService;
        this.blagajnikService = blagajnikService;
    }


    @GetMapping("/nonce")
    public Map<String, String> getNonce(@RequestParam String address) {
        String normalized = address.toLowerCase();

        String nonce = UUID.randomUUID().toString();
        nonces.putIfAbsent(normalized, nonce);

        return Map.of("nonce", nonce);
    }

    //TODO: @PostMapping("/unlink-wallet") //logika samo za roditel (roditel moze da pravi unlink)
    @PostMapping("/unlink-wallet")
    public ResponseEntity<?> unlinkWallet(Authentication authentication) {

        if (authentication == null || !authentication.isAuthenticated()) {
            return ResponseEntity.status(401)
                    .body(Map.of("message", "Unauthorized"));
        }

        CustomUserDetails userDetails =
                (CustomUserDetails) authentication.getPrincipal();

        UUID roditelId = userDetails.getId();

        RoditelResponseDTO roditel = roditelService.getRoditelById(roditelId);

        if (roditel.getMetamusk_adresa() == null) {
            return ResponseEntity.badRequest()
                    .body(Map.of("message", "No wallet linked"));
        }

        roditelService.unlinkMetamuskAddress(roditelId);

        return ResponseEntity.ok(Map.of("message", "Wallet unlinked successfully"));
    }



    @PostMapping("/link-wallet-roditel") //logika samo za roditel (roditel moze da go menuva wallet)
    public ResponseEntity<?> linkWalletParent(
            @RequestBody Map<String, String> payload,
            Authentication authentication
    ) {

        if (authentication == null || !authentication.isAuthenticated()) {
            return ResponseEntity.status(401)
                    .body(Map.of("message", "Unauthorized"));
        }

        String address = payload.get("address");
        String signature = payload.get("signature");

        if (address == null || signature == null) {
            return ResponseEntity.badRequest()
                    .body(Map.of("message", "Missing data"));
        }

        address = address.trim();
        if (!WalletUtils.isValidAddress(address)) {
            return ResponseEntity.badRequest()
                    .body(Map.of("message", "Invalid Ethereum address"));
        }
        String normalized = address.toLowerCase();

        String nonce = nonces.get(normalized);

        if (nonce == null) {
            return ResponseEntity.badRequest()
                    .body(Map.of("message", "Nonce not found"));
        }

        boolean valid = verifySignature(address, signature, nonce);

        if (!valid) {
            return ResponseEntity.status(401)
                    .body(Map.of("message", "Invalid signature"));
        }

        nonces.remove(normalized);

        CustomUserDetails userDetails =
                (CustomUserDetails) authentication.getPrincipal();

        UUID roditelId = userDetails.getId();

        RoditelResponseDTO roditel = roditelService.getRoditelById(roditelId);


        roditelService.updateMetamuskAddress(roditelId, normalized);
        System.out.println("Linked wallet attempt for roditel id " + roditelId + " with address" + normalized);


        return ResponseEntity.ok(Map.of("message", "Wallet linked successfully"));
    }


    public boolean verifySignature(String expectedAddress, String signatureHex, String message) {
        try {
            byte[] messageBytes = message.getBytes(StandardCharsets.UTF_8);
            byte[] sigBytes = Numeric.hexStringToByteArray(signatureHex);

            if (sigBytes.length != 65) {
                return false;
            }

            byte v = sigBytes[64];
            if (v < 27) {
                v += 27;
            }

            Sign.SignatureData sigData = new Sign.SignatureData(
                    v,
                    Arrays.copyOfRange(sigBytes, 0, 32),
                    Arrays.copyOfRange(sigBytes, 32, 64)
            );


            BigInteger publicKey = Sign.signedPrefixedMessageToKey(messageBytes, sigData);

            String recoveredAddress = "0x" + Keys.getAddress(publicKey);

            System.out.println("Recovered: " + recoveredAddress);

            return recoveredAddress.equalsIgnoreCase(expectedAddress);

        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }


}