package org.torproject.onionmasq.circuit;

import androidx.annotation.NonNull;

import org.torproject.onionmasq.events.ClosedConnectionEvent;
import org.torproject.onionmasq.events.FailedConnectionEvent;
import org.torproject.onionmasq.events.NewConnectionEvent;
import org.torproject.onionmasq.events.OnionmasqEvent;
import org.torproject.onionmasq.logging.LogObservable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.stream.Collectors;

public class CircuitStore {
    private final HashMap<ProxyAddressPair, NewConnectionEvent> connectionMap = new HashMap<>();
    private final HashMap<Integer, ArrayList<ProxyAddressPair>> appUIdMap = new HashMap<>();

    public void handleEvent(OnionmasqEvent event) {
        if (event instanceof NewConnectionEvent) {
            handleNewConnectionEvent((NewConnectionEvent) event);
        } else if (event instanceof ClosedConnectionEvent) {
            handleClosedConnectionEvent((ClosedConnectionEvent) event);
        } else if (event instanceof FailedConnectionEvent) {
            handleFailedConnectionEvent((FailedConnectionEvent) event);
        }
        // else ignore
    }

    public void reset() {
        connectionMap.clear();
        appUIdMap.clear();
    }

    public @NonNull ArrayList<Circuit> getCircuitsForAppUid(int appUID) {
        ArrayList<ProxyAddressPair> keys = appUIdMap.get(appUID);
        if (keys == null) {
            return new ArrayList<>();
        }

        HashSet<Circuit> resultList = new HashSet<>();
        ArrayList<ProxyAddressPair> cleanUp = new ArrayList<>();
        for (ProxyAddressPair key : keys) {
            NewConnectionEvent connectionEvent = connectionMap.get(key);
            if (connectionEvent == null) {
                cleanUp.add(key);
                continue;
            }

            Circuit circuit = new Circuit(connectionEvent.circuit, connectionEvent.torDst);
            resultList.add(circuit);
        }


        if (keys.removeAll(cleanUp)) {
            if (keys.size() == 0) {
                appUIdMap.remove(appUID);
            } else {
                appUIdMap.put(appUID, keys);
            }
        }
        return new ArrayList<>(resultList);
    }


    private void handleNewConnectionEvent(NewConnectionEvent event) {
        ProxyAddressPair key = new ProxyAddressPair(event.proxySrc, event.proxyDst);
        connectionMap.put(key, event);
        ArrayList<ProxyAddressPair> proxyAddressPairs = appUIdMap.getOrDefault(event.appId, new ArrayList<>());
        proxyAddressPairs.add(key);
        appUIdMap.put(event.appId, proxyAddressPairs);
    }

    private void handleClosedConnectionEvent(ClosedConnectionEvent event) {
        ProxyAddressPair key = new ProxyAddressPair(event.proxySrc, event.proxyDst);
        NewConnectionEvent connectionEvent = connectionMap.get(key);
        if (connectionEvent == null) {
            // event might have been removed if the connection failed before
            return;
        }

        removeProxyAddressPairForUID(connectionEvent.appId, key);
        connectionMap.remove(key);
    }

    private void handleFailedConnectionEvent(FailedConnectionEvent event) {
        ProxyAddressPair key = new ProxyAddressPair(event.proxySrc, event.proxyDst);
        NewConnectionEvent connectionEvent = connectionMap.get(key);
        if (connectionEvent == null) {
            LogObservable.getInstance().addLog("WARNING: Unknown FailedConnectionEvent with proxySrc:proxyDst tuple " + event.proxySrc + ":" + event.proxyDst);
            return;
        }

        removeProxyAddressPairForUID(event.appId, key);
        connectionMap.remove(key);
    }

    private void removeProxyAddressPairForUID(int uid, ProxyAddressPair addressPair) {
        ArrayList<ProxyAddressPair> proxyAddressPairs = appUIdMap.get(uid);
        if (proxyAddressPairs != null) {
            ArrayList<ProxyAddressPair> updatedList = proxyAddressPairs.
                    stream().
                    filter( proxyAddressPair -> !proxyAddressPair.equals(addressPair)).
                    collect(Collectors.toCollection(ArrayList::new));
            if (updatedList.size() > 0) {
                appUIdMap.put(uid, updatedList);
            } else {
                appUIdMap.remove(uid);
            }
        }
    }
}

