/*
 * Copyright (C) 2014 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.squareup.okhttp.sample;

import com.squareup.okhttp.Cache;
import com.squareup.okhttp.OkHttpClient;
import com.squareup.okhttp.Request;
import com.squareup.okhttp.Response;
import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;

/**
 * Fetches HTML from a requested URL, follows the links, and repeats.
 */
public final class Crawler {
  private final OkHttpClient client;
  private final Set<URL> fetchedUrls = Collections.synchronizedSet(new LinkedHashSet<URL>());
  private final LinkedBlockingQueue<URL> queue = new LinkedBlockingQueue<>();

  public Crawler(OkHttpClient client) {
    this.client = client;
  }

  private void parallelDrainQueue(int threadCount) {
    ExecutorService executor = Executors.newFixedThreadPool(threadCount);
    for (int i = 0; i < threadCount; i++) {
      executor.execute(new Runnable() {
        @Override public void run() {
          try {
            drainQueue();
          } catch (Exception e) {
            e.printStackTrace();
          }
        }
      });
    }
    executor.shutdown();
  }

  private void drainQueue() throws Exception {
    for (URL url; (url = queue.take()) != null; ) {
      if (!fetchedUrls.add(url)) {
        continue;
      }

      try {
        fetch(url);
      } catch (IOException e) {
        System.out.printf("XXX: %s %s%n", url, e);
      }
    }
  }

  public void fetch(URL url) throws IOException {
    Request request = new Request.Builder()
        .url(url)
        .build();
    Response response = client.newCall(request).execute();
    String responseSource = response.networkResponse() != null
        ? ("(network: " + response.networkResponse().code() + ")")
        : "(cache)";
    int responseCode = response.code();

    System.out.printf("%03d: %s %s%n", responseCode, url, responseSource);

    String contentType = response.header("Content-Type");
    if (responseCode != 200 || contentType == null) {
      response.body().close();
      return;
    }

    Document document = Jsoup.parse(response.body().string(), url.toString());
    for (Element element : document.select("a[href]")) {
      String href = element.attr("href");
      URL link = parseUrl(url, href);
      if (link != null) queue.add(link);
    }
  }

  private URL parseUrl(URL url, String href) {
    try {
      URL result = new URL(url, href);
      return result.getProtocol().equals("http") || result.getProtocol().equals("https")
          ? result
          : null;
    } catch (MalformedURLException e) {
      return null;
    }
  }

  public static void main(String[] args) throws IOException {
    if (args.length != 2) {
      System.out.println("Usage: Crawler <cache dir> <root>");
      return;
    }

    int threadCount = 20;
    long cacheByteCount = 1024L * 1024L * 100L;

    OkHttpClient client = new OkHttpClient();
    Cache cache = new Cache(new File(args[0]), cacheByteCount);
    client.setCache(cache);

    Crawler crawler = new Crawler(client);
    crawler.queue.add(new URL(args[1]));
    crawler.parallelDrainQueue(threadCount);
  }
}

@@ -779,16 +779,67 @@
 sponse;%0A
+import com.squareup.okhttp.internal.NamedRunnable;%0A
 import j
@@ -847,24 +847,24 @@
 va.io.File;%0A
-
 import java.
@@ -1019,24 +1019,71 @@
 a.util.Set;%0A
+import java.util.concurrent.ConcurrentHashMap;%0A
 import java.
@@ -1203,16 +1203,66 @@
 gQueue;%0A
+import java.util.concurrent.atomic.AtomicInteger;%0A
 import o
@@ -1658,16 +1658,112 @@
 eue%3C%3E();
+%0A  private final ConcurrentHashMap%3CString, AtomicInteger%3E hostnames = new ConcurrentHashMap%3C%3E();
 %0A%0A  publ
@@ -2025,16 +2025,21 @@
 ute(new 
+Named
 Runnable
@@ -2039,16 +2039,31 @@
 unnable(
+%22Crawler %25s%22, i
 ) %7B%0A    
@@ -2081,22 +2081,29 @@
 de p
-ublic void run
+rotected void execute
 () %7B
@@ -2621,24 +2621,318 @@
 Exception %7B%0A
+    // Skip hosts that we've visited many times.%0A    AtomicInteger hostnameCount = new AtomicInteger();%0A    AtomicInteger previous = hostnames.putIfAbsent(url.getHost(), hostnameCount);%0A    if (previous != null) hostnameCount = previous;%0A    if (hostnameCount.incrementAndGet() %3E 100) return;%0A%0A
     Request 
@@ -3182,16 +3182,49 @@
 code() +
+ %22 over %22 + response.protocol() +
  %22)%22)%0A  
@@ -3727,19 +3727,40 @@
 arseUrl(
+response.request().
 url
+()
 , href);

Fix some crawler bugs.