package com.iota.iri.service.ledger.impl;

import com.iota.iri.BundleValidator;
import com.iota.iri.controllers.MilestoneViewModel;
import com.iota.iri.controllers.StateDiffViewModel;
import com.iota.iri.controllers.TransactionViewModel;
import com.iota.iri.model.Hash;
import com.iota.iri.service.ledger.LedgerException;
import com.iota.iri.service.ledger.LedgerService;
import com.iota.iri.service.milestone.MilestoneService;
import com.iota.iri.service.snapshot.Snapshot;
import com.iota.iri.service.snapshot.SnapshotException;
import com.iota.iri.service.snapshot.SnapshotProvider;
import com.iota.iri.service.snapshot.SnapshotService;
import com.iota.iri.service.snapshot.impl.SnapshotStateDiffImpl;
import com.iota.iri.service.spentaddresses.SpentAddressesService;
import com.iota.iri.storage.Tangle;
import com.iota.iri.utils.dag.DAGHelper;

import java.util.*;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * <p>
 * Creates a service instance that allows us to perform ledger state specific operations.
 * </p>
 * <p>
 * This class is stateless and does not hold any domain specific models.
 * </p>
 */
public class LedgerServiceImpl implements LedgerService {

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

    /**
     * Holds the tangle object which acts as a database interface.
     */
    private final Tangle tangle;

    /**
     * Holds the snapshot provider which gives us access to the relevant snapshots.
     */
    private final SnapshotProvider snapshotProvider;

    /**
     * Holds a reference to the service instance containing the business logic of the snapshot package.
     */
    private final SnapshotService snapshotService;

    /**
     * Holds a reference to the service instance containing the business logic of the milestone package.
     */
    private final MilestoneService milestoneService;

    private final SpentAddressesService spentAddressesService;

    private final BundleValidator bundleValidator;

    /**
     * @param tangle Tangle object which acts as a database interface
     * @param snapshotProvider snapshot provider which gives us access to the relevant snapshots
     * @param snapshotService service instance of the snapshot package that gives us access to packages' business logic
     * @param milestoneService contains the important business logic when dealing with milestones
     */
    public LedgerServiceImpl(Tangle tangle, SnapshotProvider snapshotProvider, SnapshotService snapshotService,
            MilestoneService milestoneService, SpentAddressesService spentAddressesService,
                                  BundleValidator bundleValidator) {
        this.tangle = tangle;
        this.snapshotProvider = snapshotProvider;
        this.snapshotService = snapshotService;
        this.milestoneService = milestoneService;
        this.spentAddressesService = spentAddressesService;
        this.bundleValidator = bundleValidator;
    }

    @Override
    public void restoreLedgerState() throws LedgerException {
        try {
            Optional<MilestoneViewModel> milestone = milestoneService.findLatestProcessedSolidMilestoneInDatabase();
            if (milestone.isPresent()) {
                snapshotService.replayMilestones(snapshotProvider.getLatestSnapshot(), milestone.get().index());
            }
        } catch (Exception e) {
            throw new LedgerException("unexpected error while restoring the ledger state", e);
        }
    }

    @Override
    public boolean applyMilestoneToLedger(MilestoneViewModel milestone) throws LedgerException {
        log.debug("applying milestone {}", milestone.index());
        if(generateStateDiff(milestone)) {
            try {
                snapshotService.replayMilestones(snapshotProvider.getLatestSnapshot(), milestone.index());
            } catch (SnapshotException e) {
                throw new LedgerException("failed to apply the balance changes to the ledger state", e);
            }

            return true;
        }

        return false;
    }

    @Override
    public boolean tipsConsistent(List<Hash> tips) throws LedgerException {
        Set<Hash> visitedHashes = new HashSet<>();
        Map<Hash, Long> diff = new HashMap<>();
        for (Hash tip : tips) {
            if (!isBalanceDiffConsistent(visitedHashes, diff, tip)) {
                return false;
            }
        }

        return true;
    }

    @Override
    public boolean isBalanceDiffConsistent(Set<Hash> approvedHashes, Map<Hash, Long> diff, Hash tip) throws
            LedgerException {

        try {
            if (!TransactionViewModel.fromHash(tangle, tip).isSolid()) {
                return false;
            }
        } catch (Exception e) {
            throw new LedgerException("failed to check the consistency of the balance changes", e);
        }

        if (approvedHashes.contains(tip)) {
            return true;
        }
        Set<Hash> visitedHashes = new HashSet<>(approvedHashes);
        Map<Hash, Long> currentState = generateBalanceDiff(visitedHashes, tip,
                snapshotProvider.getLatestSnapshot().getIndex(), true);
        if (currentState == null) {
            return false;
        }
        diff.forEach((key, value) -> {
            if (currentState.computeIfPresent(key, ((hash, aLong) -> value + aLong)) == null) {
                currentState.putIfAbsent(key, value);
            }
        });
        boolean isConsistent = snapshotProvider.getLatestSnapshot().patchedState(new SnapshotStateDiffImpl(
                currentState)).isConsistent();
        if (isConsistent) {
            diff.putAll(currentState);
            approvedHashes.addAll(visitedHashes);
        }
        return isConsistent;
    }

    @Override
    public Map<Hash, Long> generateBalanceDiff(Set<Hash> visitedTransactions, Hash startTransaction, int milestoneIndex,
            boolean enforceExtraRules)
            throws LedgerException {

        Map<Hash, Long> state = new HashMap<>();

        Snapshot initialSnapshot = snapshotProvider.getInitialSnapshot();
        Map<Hash, Integer> solidEntryPoints = initialSnapshot.getSolidEntryPoints();
        solidEntryPoints.keySet().forEach(solidEntryPointHash -> {
            visitedTransactions.add(solidEntryPointHash);
        });

        final Queue<Hash> nonAnalyzedTransactions = new LinkedList<>(Collections.singleton(startTransaction));
        Hash transactionPointer;
        while ((transactionPointer = nonAnalyzedTransactions.poll()) != null) {
            try {
                final TransactionViewModel transactionViewModel = TransactionViewModel.fromHash(tangle,
                        transactionPointer);
                if (transactionViewModel.getCurrentIndex() == 0 && visitedTransactions.add(transactionPointer)) {
                    if (transactionViewModel.getType() == TransactionViewModel.PREFILLED_SLOT) {
                        return null;
                    }
                    if (!milestoneService.isTransactionConfirmed(transactionViewModel, milestoneIndex)) {

                        final List<TransactionViewModel> bundleTransactions = bundleValidator.validate(tangle,
                                enforceExtraRules, snapshotProvider.getInitialSnapshot(),
                                transactionViewModel.getHash());

                        if (bundleTransactions.isEmpty()) {
                            return null;
                        }

                        // ISSUE 1008: generateBalanceDiff should be refactored so we don't have those hidden
                        // concerns
                        spentAddressesService.persistValidatedSpentAddressesAsync(bundleTransactions);

                        if (BundleValidator.isInconsistent(bundleTransactions)) {
                            log.error("Encountered an inconsistent bundle with tail {} and bundle hash {}",
                                    bundleTransactions.get(0).getHash(), bundleTransactions.get(0).getBundleHash());
                            return null;
                        }

                        for (final TransactionViewModel bundleTransactionViewModel : bundleTransactions) {
                            if (bundleTransactionViewModel.value() != 0) {

                                final Hash address = bundleTransactionViewModel.getAddressHash();
                                final Long value = state.get(address);
                                state.put(address, value == null ? bundleTransactionViewModel.value()
                                        : Math.addExact(value, bundleTransactionViewModel.value()));
                            }
                        }
                        nonAnalyzedTransactions.addAll(DAGHelper.get(tangle).findTails(transactionViewModel));
                    }

                }

            } catch (Exception e) {
                throw new LedgerException("unexpected error while generating the balance diff", e);
            }
        }

        return state;
    }


    /**
     * <p>
     * Generates the {@link com.iota.iri.model.StateDiff} that belongs to the given milestone in the database and marks
     * all transactions that have been approved by the milestone accordingly by setting their {@code snapshotIndex}
     * value.
     * </p>
     * <p>
     * It first checks if the {@code snapshotIndex} of the transaction belonging to the milestone was correctly set
     * already (to determine if this milestone was processed already) and proceeds to generate the {@link
     * com.iota.iri.model.StateDiff} if that is not the case. To do so, it calculates the balance changes, checks if
     * they are consistent and only then writes them to the database.
     * </p>
     * <p>
     * If inconsistencies in the {@code snapshotIndex} are found it issues a reset of the corresponding milestone to
     * recover from this problem.
     * </p>
     *
     * @param milestone the milestone that shall have its {@link com.iota.iri.model.StateDiff} generated
     * @return {@code true} if the {@link com.iota.iri.model.StateDiff} could be generated and {@code false} otherwise
     * @throws LedgerException if anything unexpected happens while generating the {@link com.iota.iri.model.StateDiff}
     */
    private boolean generateStateDiff(MilestoneViewModel milestone) throws LedgerException {
        try {
            TransactionViewModel transactionViewModel = TransactionViewModel.fromHash(tangle, milestone.getHash());

            if (!transactionViewModel.isSolid()) {
                return false;
            }

            final int transactionSnapshotIndex = transactionViewModel.snapshotIndex();
            boolean successfullyProcessed = transactionSnapshotIndex == milestone.index();
            if (!successfullyProcessed) {
                // if the snapshotIndex of our transaction was set already, we have processed our milestones in
                // the wrong order (i.e. while rescanning the db)
                if (transactionSnapshotIndex != 0) {
                    milestoneService.resetCorruptedMilestone(milestone.index());
                }

                snapshotProvider.getLatestSnapshot().lockRead();
                try {
                    Hash tail = transactionViewModel.getHash();
                    Map<Hash, Long> balanceChanges = generateBalanceDiff(new HashSet<>(), tail,
                            snapshotProvider.getLatestSnapshot().getIndex(), false);
                    successfullyProcessed = balanceChanges != null;
                    if (successfullyProcessed) {
                        successfullyProcessed = snapshotProvider.getLatestSnapshot().patchedState(
                                new SnapshotStateDiffImpl(balanceChanges)).isConsistent();
                        if (successfullyProcessed) {
                            milestoneService.updateMilestoneIndexOfMilestoneTransactions(milestone.getHash(),
                                    milestone.index());

                            if (!balanceChanges.isEmpty()) {
                                new StateDiffViewModel(balanceChanges, milestone.getHash()).store(tangle);
                            }
                        }
                    }
                } finally {
                    snapshotProvider.getLatestSnapshot().unlockRead();
                }
            }

            return successfullyProcessed;
        } catch (Exception e) {
            throw new LedgerException("unexpected error while generating the StateDiff for " + milestone, e);
        }
    }
}
