package cn.octopusyan.dmt.common.manager.http;

import java.net.*;
import java.util.*;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Cookie 管理
 *
 * @author octopus_yan
 */
public class CookieManager {

    private static final InMemoryCookieStore inMemoryCookieStore = new InMemoryCookieStore();
    private static final java.net.CookieManager cookieManager =
            new java.net.CookieManager(inMemoryCookieStore, CookiePolicy.ACCEPT_ALL);

    public static java.net.CookieManager get() {
        return cookieManager;
    }

    public static InMemoryCookieStore getStore() {
        return inMemoryCookieStore;
    }

    public static class InMemoryCookieStore implements CookieStore {
        // the in-memory representation of cookies
        private final List<HttpCookie> cookieJar;

        // the cookies are indexed by its domain and associated uri (if present)
        // CAUTION: when a cookie removed from main data structure (i.e. cookieJar),
        //          it won't be cleared in domainIndex & uriIndex. Double-check the
        //          presence of cookie when retrieve one form index store.
        private final Map<String, List<HttpCookie>> domainIndex;
        private final Map<URI, List<HttpCookie>> uriIndex;

        // use ReentrantLock instead of synchronized for scalability
        private final ReentrantLock lock;


        /**
         * The default ctor
         */
        public InMemoryCookieStore() {
            cookieJar = new ArrayList<>();
            domainIndex = new HashMap<>();
            uriIndex = new HashMap<>();

            lock = new ReentrantLock(false);
        }

        /**
         * Add one cookie into cookie store.
         */
        public void add(URI uri, HttpCookie cookie) {
            // pre-condition : argument can't be null
            if (cookie == null) {
                throw new NullPointerException("cookie is null");
            }


            lock.lock();
            try {
                // remove the ole cookie if there has had one
                cookieJar.remove(cookie);

                // add new cookie if it has a non-zero max-age
                if (cookie.getMaxAge() != 0) {
                    cookieJar.add(cookie);
                    // and add it to domain index
                    if (cookie.getDomain() != null) {
                        addIndex(domainIndex, cookie.getDomain(), cookie);
                    }
                    if (uri != null) {
                        // add it to uri index, too
                        addIndex(uriIndex, getEffectiveURI(uri), cookie);
                    }
                }
            } finally {
                lock.unlock();
            }
        }


        /**
         * Get all cookies, which:
         * 1) given uri domain-matches with, or, associated with
         * given uri when added to the cookie store.
         * 3) not expired.
         * See RFC 2965 sec. 3.3.4 for more detail.
         */
        public List<HttpCookie> get(URI uri) {
            // argument can't be null
            if (uri == null) {
                throw new NullPointerException("uri is null");
            }

            List<HttpCookie> cookies = new ArrayList<>();
            boolean secureLink = "https".equalsIgnoreCase(uri.getScheme());
            lock.lock();
            try {
                // check domainIndex first
                getInternal1(cookies, domainIndex, uri.getHost(), secureLink);
                // check uriIndex then
                getInternal2(cookies, uriIndex, getEffectiveURI(uri), secureLink);
            } finally {
                lock.unlock();
            }

            return cookies;
        }

        /**
         * Get all cookies in cookie store, except those have expired
         */
        public List<HttpCookie> getCookies() {
            List<HttpCookie> rt;

            lock.lock();
            try {
                cookieJar.removeIf(HttpCookie::hasExpired);
            } finally {
                rt = Collections.unmodifiableList(cookieJar);
                lock.unlock();
            }

            return rt;
        }

        /**
         * Get all URIs, which are associated with at least one cookie
         * of this cookie store.
         */
        public List<URI> getURIs() {
            List<URI> uris;

            lock.lock();
            try {
                Iterator<URI> it = uriIndex.keySet().iterator();
                while (it.hasNext()) {
                    URI uri = it.next();
                    List<HttpCookie> cookies = uriIndex.get(uri);
                    if (cookies == null || cookies.isEmpty()) {
                        // no cookies list or an empty list associated with
                        // this uri entry, delete it
                        it.remove();
                    }
                }
            } finally {
                uris = new ArrayList<>(uriIndex.keySet());
                lock.unlock();
            }

            return uris;
        }


        /**
         * Remove a cookie from store
         */
        public boolean remove(URI uri, HttpCookie ck) {
            // argument can't be null
            if (ck == null) {
                throw new NullPointerException("cookie is null");
            }

            boolean modified;
            lock.lock();
            try {
                modified = cookieJar.remove(ck);
            } finally {
                lock.unlock();
            }

            return modified;
        }


        /**
         * Remove all cookies in this cookie store.
         */
        public boolean removeAll() {
            lock.lock();
            try {
                if (cookieJar.isEmpty()) {
                    return false;
                }
                cookieJar.clear();
                domainIndex.clear();
                uriIndex.clear();
            } finally {
                lock.unlock();
            }

            return true;
        }


        /* ---------------- Private operations -------------- */


        /*
         * This is almost the same as HttpCookie.domainMatches except for
         * one difference: It won't reject cookies when the 'H' part of the
         * domain contains a dot ('.').
         * I.E.: RFC 2965 section 3.3.2 says that if host is x.y.domain.com
         * and the cookie domain is .domain.com, then it should be rejected.
         * However that's not how the real world works. Browsers don't reject and
         * some sites, like yahoo.com do actually expect these cookies to be
         * passed along.
         * And should be used for 'old' style cookies (aka Netscape type of cookies)
         */
        private boolean netscapeDomainMatches(String domain, String host) {
            if (domain == null || host == null) {
                return false;
            }

            // if there's no embedded dot in domain and domain is not .local
            boolean isLocalDomain = ".local".equalsIgnoreCase(domain);
            int embeddedDotInDomain = domain.indexOf('.');
            if (embeddedDotInDomain == 0) {
                embeddedDotInDomain = domain.indexOf('.', 1);
            }
            if (!isLocalDomain && (embeddedDotInDomain == -1 || embeddedDotInDomain == domain.length() - 1)) {
                return false;
            }

            // if the host name contains no dot and the domain name is .local
            int firstDotInHost = host.indexOf('.');
            if (firstDotInHost == -1 && isLocalDomain) {
                return true;
            }

            int domainLength = domain.length();
            int lengthDiff = host.length() - domainLength;
            if (lengthDiff == 0) {
                // if the host name and the domain name are just string-compare equal
                return host.equalsIgnoreCase(domain);
            } else if (lengthDiff > 0) {
                // need to check H & D component
                String H = host.substring(0, lengthDiff);
                String D = host.substring(lengthDiff);

                return (D.equalsIgnoreCase(domain));
            } else if (lengthDiff == -1) {
                // if domain is actually .host
                return (domain.charAt(0) == '.' &&
                        host.equalsIgnoreCase(domain.substring(1)));
            }

            return false;
        }

        private void getInternal1(List<HttpCookie> cookies, Map<String, List<HttpCookie>> cookieIndex,
                                  String host, boolean secureLink) {
            // Use a separate list to handle cookies that need to be removed so
            // that there is no conflict with iterators.
            ArrayList<HttpCookie> toRemove = new ArrayList<>();
            for (Map.Entry<String, List<HttpCookie>> entry : cookieIndex.entrySet()) {
                String domain = entry.getKey();
                List<HttpCookie> lst = entry.getValue();
                for (HttpCookie c : lst) {
                    if ((c.getVersion() == 0 && netscapeDomainMatches(domain, host)) ||
                            (c.getVersion() == 1 && HttpCookie.domainMatches(domain, host))) {
                        if ((cookieJar.contains(c))) {
                            // the cookie still in main cookie store
                            if (!c.hasExpired()) {
                                // don't add twice and make sure it's the proper
                                // security level
                                if ((secureLink || !c.getSecure()) &&
                                        !cookies.contains(c)) {
                                    cookies.add(c);
                                }
                            } else {
                                toRemove.add(c);
                            }
                        } else {
                            // the cookie has been removed from main store,
                            // so also remove it from domain indexed store
                            toRemove.add(c);
                        }
                    }
                }
                // Clear up the cookies that need to be removed
                for (HttpCookie c : toRemove) {
                    lst.remove(c);
                    cookieJar.remove(c);

                }
                toRemove.clear();
            }
        }

        // @param cookies           [OUT] contains the found cookies
        // @param cookieIndex       the index
        // @param comparator        the prediction to decide whether or not
        //                          a cookie in index should be returned
        private <T> void getInternal2(List<HttpCookie> cookies,
                                      Map<T, List<HttpCookie>> cookieIndex,
                                      Comparable<T> comparator, boolean secureLink) {
            for (T index : cookieIndex.keySet()) {
                if (comparator.compareTo(index) == 0) {
                    List<HttpCookie> indexedCookies = cookieIndex.get(index);
                    // check the list of cookies associated with this domain
                    if (indexedCookies != null) {
                        Iterator<HttpCookie> it = indexedCookies.iterator();
                        while (it.hasNext()) {
                            HttpCookie ck = it.next();
                            if (cookieJar.contains(ck)) {
                                // the cookie still in main cookie store
                                if (!ck.hasExpired()) {
                                    // don't add twice
                                    if ((secureLink || !ck.getSecure()) &&
                                            !cookies.contains(ck))
                                        cookies.add(ck);
                                } else {
                                    it.remove();
                                    cookieJar.remove(ck);
                                }
                            } else {
                                // the cookie has been removed from main store,
                                // so also remove it from domain indexed store
                                it.remove();
                            }
                        }
                    } // end of indexedCookies != null
                } // end of comparator.compareTo(index) == 0
            } // end of cookieIndex iteration
        }

        // add 'cookie' indexed by 'index' into 'indexStore'
        private <T> void addIndex(Map<T, List<HttpCookie>> indexStore,
                                  T index,
                                  HttpCookie cookie) {
            if (index != null) {
                List<HttpCookie> cookies = indexStore.get(index);
                if (cookies != null) {
                    // there may already have the same cookie, so remove it first
                    cookies.remove(cookie);

                    cookies.add(cookie);
                } else {
                    cookies = new ArrayList<>();
                    cookies.add(cookie);
                    indexStore.put(index, cookies);
                }
            }
        }


        //
        // for cookie purpose, the effective uri should only be http://host
        // the path will be taken into account when path-match algorithm applied
        //
        private URI getEffectiveURI(URI uri) {
            URI effectiveURI;
            try {
                effectiveURI = new URI("http",
                        uri.getHost(),
                        null,  // path component
                        null,  // query component
                        null   // fragment component
                );
            } catch (URISyntaxException ignored) {
                effectiveURI = uri;
            }

            return effectiveURI;
        }
    }
}
