package com.jscape.inet.ssh.util;

import java.io.Serializable;
import java.net.InetAddress;
import java.security.PublicKey;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import static com.jscape.util.Parameters.assertNotNull;

/**
 * Associative host keys store. Contains host addresses with
 * associated host keys as one-to-one map.
 * <p/>
 * Supports serialization for simplifying physical save/restore process.
 * <p/>
 * Methods of this class are not thread-safe.
 * <p/>
 * All methods throw <code>NullPointerException</code> if some object
 * parameter is <code>null</code> and <code>IllegalArgumentException</code>
 * if some parameter doesn't correspond to the method contract.
 *
 * @author Alex
 */
public class SshHostKeys
        implements Serializable {

    private Map<InetAddress, List<Entry>> entries;

    /**
     * Creates a new store instance.
     *
     * @param entries
     */
    public SshHostKeys(Entry... entries) {
        this.entries = new ConcurrentHashMap<>();
        set(entries);
    }

    /**
     * Sets the store entries.
     *
     * @param entries entries
     *
     * @return current store instance
     */
    public SshHostKeys set(Entry... entries) {
        for (Entry entry : entries) {
            List<Entry> hostEntries = this.entries.computeIfAbsent(entry.host, k -> new LinkedList<>());
            hostEntries.add(entry);
        }

        return this;
    }

    /**
     * Puts a key to the store.
     *
     * @param host a server host
     * @param key a server key
     *
     * @return current store instance
     */
    public SshHostKeys put(InetAddress host, PublicKey key) {
        set(Entry.entryFor(host, key), Entry.entrySha256For(host, key));

        return this;
    }

    /**
     * Puts a key to the store.
     *
     * @param host a server host
     * @param keyFingerprint a key fingerprint
     *
     * @return current store instance
     */
    public SshHostKeys put(InetAddress host, String keyFingerprint) {
        Entry entry = new Entry(host, null, keyFingerprint);
        set(entry);

        return this;
    }

    /**
     * Adds a new key entry to the store.
     *
     * @param host           the host address
     * @param keyFingerprint host key fingerprint string
     *
     * @deprecated use put()
     */
    public SshHostKeys addKey(InetAddress host, String keyFingerprint) {
        return put(host, keyFingerprint);
    }

    /**
     * Removes a key from the store.
     *
     * @param host a server host
     * @param key a server key
     *
     * @return current store instance
     */
    public SshHostKeys remove(InetAddress host, PublicKey key) {
        List<Entry> hostEntries = this.entries.get(host);

        if (hostEntries == null) {
            return this;
        }

        for (Entry entry : hostEntries) {
            if (entry.key != null && entry.key.equals(key)) {
                hostEntries.remove(entry);
                break;
            }
        }

        if (hostEntries.isEmpty()) {
            this.entries.remove(host);
        }

        return this;
    }

    /**
     * Removes a key from the store.
     *
     * @param host a server host
     * @param keyFingerprint a key fingerprint
     *
     * @return current store instance
     */
    public SshHostKeys remove(InetAddress host, String keyFingerprint) {
        List<Entry> hostEntries = this.entries.get(host);

        if (hostEntries == null) {
            return this;
        }

        for (Entry entry : hostEntries) {
            if (entry.fingerprint.equals(keyFingerprint)) {
                hostEntries.remove(entry);
                break;
            }
        }

        if (hostEntries.isEmpty()) {
            this.entries.remove(host);
        }

        return this;
    }

    /**
     * Removes the key entry for the specified host.
     *
     * @param host           target host
     * @param keyFingerprint key fingerprint to remove
     *
     * @deprecated use remove()
     */
    public void removeKey(InetAddress host, String keyFingerprint) {
        remove(host, keyFingerprint);
    }

    /**
     * Clears the store contents.
     */
    public void clear() {
        this.entries.clear();
    }

    /**
     * Returns known server hosts.
     *
     * @return known server hosts
     */
    public Collection<InetAddress> hosts() {
        return this.entries.keySet();
    }

    /**
     * Returns current hosts collection contained in the store.
     *
     * @return current hosts collection
     *
     * @deprecated use hosts()
     */
    public Iterator getHosts() {
        return hosts().iterator();
    }

    /**
     * Returns a collection of known key entries for specified host.
     *
     * @param host a server host
     *
     * @return a collection of known key entries
     */
    public Collection<Entry> entriesOf(InetAddress host) {
        List<Entry> hostEntries = this.entries.get(host);
        return hostEntries != null ? hostEntries : Collections.<Entry>emptyList();
    }

    /**
     * Returns keys iterator for the specified host.
     *
     * @param host target host
     *
     * @return host key iterator
     *
     * @deprecated use entriesOf()
     */
    public Iterator<String> getKeys(InetAddress host) {
        return fingerprintsOf(host).iterator();
    }

    /**
     * Returns a collection of known keys for specified host.
     *
     * @param host a server host
     *
     * @return a collection of known keys
     */
    public List<PublicKey> keysOf(InetAddress host) {
        List<PublicKey> keys = new ArrayList<>();
        for (Entry entry : entriesOf(host)) {
            if (entry.key != null) {
                keys.add(entry.key);
            }
        }

        return keys;
    }

    /**
     * Returns a collection of known key fingerprints for specified host.
     *
     * @param host a server host
     *
     * @return a collection of known key fingerprints
     */
    public List<String> fingerprintsOf(InetAddress host) {
        List<String> fingerprints = new ArrayList<>();
        for (Entry entry : entriesOf(host)) {
            fingerprints.add(entry.fingerprint);
        }

        return fingerprints;
    }

    /**
     * Checks if specified key is already known.
     *
     * @param host a server host
     * @param key a server key
     *
     * @return <code>true</code> if key is known; <code>false</code> otherwise
     */
    public boolean contains(InetAddress host, PublicKey key) {
        Fingerprint fingerprint = new Fingerprint(key);
        for (Entry entry : entriesOf(host)) {
            if (fingerprint.sameAsIn(entry)) {
                return true;
            }
        }

        return false;
    }

    @Override
    public String toString() {
        return "SshHostKeys {"
               + "entries=" + this.entries
               + '}';
    }

    /**
     * A store entry.
     */
    public static class Entry {

        public final InetAddress host;
        public final PublicKey key;
        public final String fingerprint;

        public static Entry entryFor(InetAddress host, PublicKey key) {
            return new Entry(host, key, KeyFingerprintFormat.format(key));
        }

        public static Entry entrySha256For(InetAddress host, PublicKey key) {
            return new Entry(host, key, KeyFingerprintFormat.formatSha256(key));
        }

        public Entry(InetAddress host, PublicKey key, String fingerprint) {
            assertNotNull(host);
            this.host = host;

            this.key = key;

            assertNotNull(fingerprint);
            this.fingerprint = fingerprint;
        }

        @Override
        public String toString() {
            return "Entry {"
                   + "key=" + this.key
                   + ", fingerprint='" + this.fingerprint + '\''
                   + '}';
        }
    }

    public static class Fingerprint {

        private final PublicKey key;

        public Fingerprint(PublicKey key) {
            this.key = key;
        }

        public boolean sameAsIn(Entry entry) {
            return this.key.equals(entry.key) || sameMd5(entry.fingerprint) || sameSha256(entry.fingerprint);
        }

        private boolean sameMd5(String fingerprint) {
            if (fingerprint == null) {
                return false;
            }

            try {
                return KeyFingerprintFormat.format(this.key).equals(fingerprint);
            } catch (Exception e) {
                return false;
            }
        }

        private boolean sameSha256(String fingerprint) {
            if (fingerprint == null) {
                return false;
            }

            try {
                byte[] hash = KeyFingerprintFormat.sha256HashOf(this.key);
                byte[] fingerprintHash = KeyFingerprintFormat.parseSha256(fingerprint);

                return Arrays.equals(hash, fingerprintHash);
            } catch (Exception e) {
                return false;
            }
        }
    }
}
