package com.r35157.cryptowallet.solana.impl.ref;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.r35157.cryptowallet.solana.SolanaWallet;
import com.r35157.cryptowallet.solana.MissingRecipientTokenAccountPolicy;
import com.r35157.libs.solana.SolanaAccountInfo;
import com.r35157.libs.solana.SPLTokenHolding;
import com.r35157.libs.solana.SolanaBlockChain;
import com.r35157.libs.solana.SolanaSignedTransaction;
import com.r35157.libs.solana.SolanaUnsignedTransaction;
import com.r35157.libs.solana.SPLTokenSupply;
import com.r35157.libs.solana.valuetypes.economic.SolanaSPLTokenProgram;
import com.r35157.libs.valuetypes.basic.MoneyAmount;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.IOException;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Objects;

public class SolanaWalletImpl implements SolanaWallet {
    public SolanaWalletImpl(
            ΩSolanaWalletIdΩ address,
            String signerKeyName,
            SolanaBlockChain solanaBlockChain
    ) {
        this.address = address;
        this.signerKeyName = signerKeyName;
        this.solanaBlockChain = solanaBlockChain;
    }

    @Override
    public @NotNull ΩSolanaWalletIdΩ getAddress() {
        return address;
    }

    @Override
    public @NotNull ΩSolanaAmountΩ getSolanaBalance()
            throws IOException, InterruptedException {
        return solanaBlockChain.getBalanceInSolana(address);
    }

    @Override
    public @Nullable ΩAmountΩ getSPLTokenBalance(
            @NotNull ΩSPLMintAddressΩ mintAddress,
            @NotNull SolanaSPLTokenProgram splProgram
    ) throws IOException, InterruptedException {
        SPLTokenHolding tokenHolding = solanaBlockChain
                .getSPLTokenHoldings(address, splProgram)
                .get(mintAddress);

        return tokenHolding == null
                ? null
                : tokenHolding.uiAmount();
    }

    @Override
    public @NotNull ΩSolanaTransactionSignatureΩ sendSolana(
            @NotNull ΩSolanaWalletIdΩ recipient,
            @NotNull ΩSolanaAmountΩ amount
    ) throws IOException, InterruptedException {
        SolanaUnsignedTransaction unsignedTransaction =
                solanaBlockChain.buildSolanaTransferTransaction(
                        address,
                        recipient,
                        amount
                );

        return signAndSendTransaction(unsignedTransaction);
    }

    @Override
    public @NotNull ΩSolanaTransactionSignatureΩ sendAllSolana(
            @NotNull ΩSolanaWalletIdΩ recipient
    ) throws IOException, InterruptedException {
        SolanaUnsignedTransaction unsignedTransaction =
                solanaBlockChain.buildSolanaTransferAllTransaction(
                        address,
                        recipient
                );

        return signAndSendTransaction(unsignedTransaction);
    }

    @Override
    public @NotNull ΩSolanaTransactionSignatureΩ sendSPLToken(
            @NotNull ΩSPLMintAddressΩ mintAddress,
            @NotNull ΩSolanaWalletIdΩ recipient,
            @NotNull ΩAmountΩ amount,
            @NotNull MissingRecipientTokenAccountPolicy policy
    ) throws IOException, InterruptedException {
        if (mintAddress == null || mintAddress.isBlank()) {
            throw new IllegalArgumentException(
                    "SPL token mint address must not be blank"
            );
        }
        if (recipient == null
                || recipient.isBlank()
                || recipient.equals(address)) {
            throw new IllegalArgumentException(
                    "SPL token recipient must be present and different "
                            + "from the sender"
            );
        }
        if (amount == null || amount.signum() <= 0) {
            throw new IllegalArgumentException(
                    "SPL token transfer amount must be greater than zero"
            );
        }
        Objects.requireNonNull(
                policy,
                "Missing recipient token account policy must not be null"
        );

        SolanaAccountInfo mintAccount = solanaBlockChain.getAccountInfo(
                mintAddress
        );
        if (mintAccount == null) {
            throw new IllegalStateException(
                    "SPL token mint account does not exist: " + mintAddress
            );
        }

        SolanaSPLTokenProgram splProgram = detectTokenProgram(mintAccount);
        SPLTokenSupply tokenSupply = solanaBlockChain.getSPLTokenSupply(
                mintAddress,
                splProgram
        );
        SPLTokenHolding senderHolding = solanaBlockChain
                .getSPLTokenHoldings(address, splProgram)
                .get(mintAddress);

        if (senderHolding == null) {
            throw new IllegalStateException(
                    "Sender does not have a token account for mint "
                            + mintAddress
            );
        }
        if (senderHolding.decimals() != tokenSupply.decimals()) {
            throw new IllegalStateException(
                    "Sender token account decimals do not match the mint"
            );
        }

        ΩRawAmountΩ rawAmount;
        try {
            rawAmount = amount
                    .movePointRight(tokenSupply.decimals())
                    .toBigIntegerExact()
                    .toString();
        } catch (ArithmeticException e) {
            throw new IllegalArgumentException(
                    "SPL token transfer amount has more than "
                            + tokenSupply.decimals()
                            + " decimal places",
                    e
            );
        }

        if (new java.math.BigInteger(rawAmount).compareTo(
                new java.math.BigInteger(senderHolding.rawAmount())
        ) > 0) {
            throw new IllegalStateException(
                    "Sender has insufficient SPL token balance"
            );
        }

        ΩSPLTokenAccountΩ recipientTokenAccount =
                solanaBlockChain.findAssociatedTokenAccount(
                        recipient,
                        mintAddress,
                        splProgram
                );
        SolanaAccountInfo recipientAccount = solanaBlockChain.getAccountInfo(
                recipientTokenAccount
        );
        boolean createRecipientAccount = recipientAccount == null;

        if (createRecipientAccount
                && policy == MissingRecipientTokenAccountPolicy.FAIL) {
            throw new IllegalStateException(
                    "Recipient Associated Token Account does not exist: "
                            + recipientTokenAccount
            );
        }
        if (recipientAccount != null) {
            validateRecipientTokenAccount(
                    recipientAccount,
                    mintAddress,
                    recipient,
                    splProgram
            );
        }

        SolanaUnsignedTransaction transaction =
                solanaBlockChain.buildSPLTokenTransferTransaction(
                        address,
                        senderHolding.tokenAccount(),
                        mintAddress,
                        recipient,
                        recipientTokenAccount,
                        rawAmount,
                        tokenSupply.decimals(),
                        splProgram,
                        createRecipientAccount
                );

        return signAndSendTransaction(transaction);
    }

    @Override
    public @NotNull ΩSolanaTransactionSignatureΩ burnSPLToken(
            @NotNull ΩSPLMintAddressΩ mintAddress,
            @NotNull ΩAmountΩ amount
    ) throws IOException, InterruptedException {
        if (mintAddress == null || mintAddress.isBlank()) {
            throw new IllegalArgumentException(
                    "SPL token mint address must not be blank"
            );
        }
        if (amount == null || amount.signum() <= 0) {
            throw new IllegalArgumentException(
                    "SPL token burn amount must be greater than zero"
            );
        }

        SolanaAccountInfo mintAccount = solanaBlockChain.getAccountInfo(
                mintAddress
        );
        if (mintAccount == null) {
            throw new IllegalStateException(
                    "SPL token mint account does not exist: " + mintAddress
            );
        }

        SolanaSPLTokenProgram splProgram = detectTokenProgram(mintAccount);
        SPLTokenSupply tokenSupply = solanaBlockChain.getSPLTokenSupply(
                mintAddress,
                splProgram
        );
        if (tokenSupply == null) {
            throw new IllegalStateException(
                    "SPL token mint supply does not exist: " + mintAddress
            );
        }
        if (tokenSupply.decimals() < 0 || tokenSupply.decimals() > 255) {
            throw new IllegalStateException(
                    "SPL token mint decimals do not fit in an unsigned byte"
            );
        }

        SPLTokenHolding holding = solanaBlockChain
                .getSPLTokenHoldings(address, splProgram)
                .get(mintAddress);
        if (holding == null) {
            throw new IllegalStateException(
                    "Wallet does not have a token account for mint "
                            + mintAddress
            );
        }
        if (holding.decimals() != tokenSupply.decimals()) {
            throw new IllegalStateException(
                    "Wallet token account decimals do not match the mint"
            );
        }

        BigInteger burnAmount;
        try {
            burnAmount = amount
                    .movePointRight(tokenSupply.decimals())
                    .toBigIntegerExact();
        } catch (ArithmeticException e) {
            throw new IllegalArgumentException(
                    "SPL token burn amount has more than "
                            + tokenSupply.decimals()
                            + " decimal places",
                    e
            );
        }

        BigInteger holdingAmount;
        try {
            holdingAmount = new BigInteger(holding.rawAmount());
        } catch (RuntimeException e) {
            throw new IllegalStateException(
                    "Wallet token account contains an invalid raw balance",
                    e
            );
        }
        if (holdingAmount.signum() < 0) {
            throw new IllegalStateException(
                    "Wallet token account contains a negative raw balance"
            );
        }
        if (burnAmount.compareTo(holdingAmount) > 0) {
            throw new IllegalStateException(
                    "Wallet has insufficient SPL token balance"
            );
        }

        SolanaUnsignedTransaction transaction =
                solanaBlockChain.buildSPLTokenBurnTransaction(
                        address,
                        holding.tokenAccount(),
                        mintAddress,
                        burnAmount.toString(),
                        tokenSupply.decimals(),
                        splProgram
                );

        return signAndSendTransaction(transaction);
    }

    private void validateRecipientTokenAccount(
            SolanaAccountInfo tokenAccount,
            ΩSPLMintAddressΩ expectedMint,
            ΩSolanaWalletIdΩ expectedOwner,
            SolanaSPLTokenProgram expectedProgram
    ) {
        if (!expectedProgram.getAddress().equals(tokenAccount.owner())) {
            throw new IllegalStateException(
                    "Recipient token account is owned by an unexpected program"
            );
        }

        byte[] accountData;
        try {
            accountData = Base64.getDecoder().decode(tokenAccount.data());
        } catch (IllegalArgumentException e) {
            throw new IllegalStateException(
                    "Recipient token account contains invalid Base64 data",
                    e
            );
        }
        if (accountData.length < 64) {
            throw new IllegalStateException(
                    "Recipient token account data is too short"
            );
        }

        ΩSolanaAddressΩ actualMint = solanaBlockChain.encodeSolanaAddress(
                java.util.Arrays.copyOfRange(accountData, 0, 32)
        );
        ΩSolanaAddressΩ actualOwner = solanaBlockChain.encodeSolanaAddress(
                java.util.Arrays.copyOfRange(accountData, 32, 64)
        );

        if (!expectedMint.equals(actualMint)) {
            throw new IllegalStateException(
                    "Recipient token account belongs to an unexpected mint"
            );
        }
        if (!expectedOwner.equals(actualOwner)) {
            throw new IllegalStateException(
                    "Recipient token account belongs to an unexpected owner"
            );
        }
    }

    private SolanaSPLTokenProgram detectTokenProgram(
            SolanaAccountInfo mintAccount
    ) {
        for (SolanaSPLTokenProgram program
                : SolanaSPLTokenProgram.values()) {
            if (program.getAddress().equals(mintAccount.owner())) {
                return program;
            }
        }

        throw new IllegalStateException(
                "Mint account is not owned by a supported SPL token program: "
                        + mintAccount.owner()
        );
    }

    @Override
    public @NotNull SolanaSignedTransaction signTransaction(
            @NotNull SolanaUnsignedTransaction transaction
    ) throws IOException, InterruptedException {
        if (transaction.serializedTransaction() == null
                || transaction.serializedTransaction().isBlank()) {
            throw new IllegalArgumentException(
                    "Unsigned serialized transaction must not be blank"
            );
        }

        Process process = new ProcessBuilder(
                "jup",
                "sign",
                "-f",
                "json",
                "--key",
                signerKeyName,
                "--tx",
                transaction.serializedTransaction()
        ).start();

        String standardOutput = new String(
                process.getInputStream().readAllBytes(),
                StandardCharsets.UTF_8
        );

        String standardError = new String(
                process.getErrorStream().readAllBytes(),
                StandardCharsets.UTF_8
        );

        int exitCode = process.waitFor();

        if (exitCode != 0) {
            throw new IOException(
                    "Solana transaction signing failed with exit code "
                            + exitCode
                            + ": "
                            + standardError
            );
        }

        SignerResponse response;

        try {
            response = OBJECT_MAPPER.readValue(
                    standardOutput,
                    SignerResponse.class
            );
        } catch (JsonProcessingException e) {
            throw new IOException(
                    "Unable to decode signer response: " + standardOutput,
                    e
            );
        }

        if (response.signer() == null || response.signer().isBlank()) {
            throw new IOException(
                    "Signer response did not contain a signer address"
            );
        }

        if (!address.equals(response.signer())) {
            throw new IOException(
                    "Unexpected signer address. Expected "
                            + address
                            + " but received "
                            + response.signer()
            );
        }

        if (response.signedTransaction() == null
                || response.signedTransaction().isBlank()) {
            throw new IOException(
                    "Signer response did not contain a signed transaction"
            );
        }

        ΩBase64StringΩ signedTransaction = response.signedTransaction();
        SolanaSignedTransaction sst = new SolanaSignedTransaction(signedTransaction);

        return sst;
    }

    @Override
    public @NotNull ΩSolanaTransactionSignatureΩ sendTransaction(
            @NotNull SolanaSignedTransaction transaction
    ) throws IOException, InterruptedException {
        return solanaBlockChain.sendTransaction(transaction);
    }

    private record SignerResponse(
            String signer,
            ΩBase64StringΩ signedTransaction
    ) {
    }

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private final ΩSolanaWalletIdΩ address;
    private final String signerKeyName;
    private final SolanaBlockChain solanaBlockChain;
}
