# This is an asynchronous image crawler written in python
from tornado.httpclient import AsyncHTTPClient
from tornado.queues import Queue
from tornado.locks import Semaphore
from tornado.ioloop import IOLoop
from tornado.web import Application, RequestHandler
from tornado.httpserver import HTTPServer
from tornado import gen

import random
import bs4
import urlparse
import PIL.Image
import ConfigParser
import sys
import time
import re
import os.path

# data structures
links = Queue()
imageurls = Queue()
visited_links = set()
downloaded_images = set()
link_failures = []
download_failures = []
img_counter = 0

class WebHandler(RequestHandler):
    def get(self):
        response_str = "<h1>Visited links: %d<br>"%len(visited_links)
        response_str += "Downloaded images: %d</h1>"%len(downloaded_images)
        response_str += "<hr>"
        response_str += "<b>VISITED:</b><br>"
        for link in visited_links:
            response_str += link+"<br>"
        response_str += "<b>DOWNLOADED:</b><br>"
        for img in downloaded_images:
            response_str += img+"<br>"
        self.write(response_str)
        
class Crawler(object):
    def _init_defaults(self):
        self.start_link = None
        self.link_priority = 2
        self.img_priority = 8
        self.politeness = 2
        self.workers_limit = 10 # allow at most 10 concurrent workers
        self.link_regex = re.compile("^http://.*")
        self.img_regex = re.compile(".*")
        self.fname_digits = 4
        self.min_width = 200
        self.min_height = 200
        self.img_dir = "E:/tmp/"
        self.idle_wait_loops = 100
        self.port = 8888

    def _load_config(self):
        parser = ConfigParser.ConfigParser()
        parser.read("config.ini")

        if parser.has_option("global", "starturl"):
            starturl = parser.get("global", "starturl")
            self.start_link = starturl
            
        if parser.has_option("global", "linkregex"):
            self.link_regex = re.compile(parser.get("global", "linkregex"))
        if parser.has_option("global", "imgregex"):
            self.img_regex = re.compile(parser.get("global", "imgregex"))

        if parser.has_option("global", "politeness"):
            politeness = parser.getint("global", "politeness")
            if politeness <=0:
                print "politeness must be a positive integer"
                raise SystemExit()
            self.politeness = politeness
        if parser.has_option("global", "imgdir"):
            imgdir = parser.get("global", "imgdir")
            if not os.path.exists(imgdir) or not os.path.isdir(imgdir):
                print "invalid imgdir configuration"
                raise SystemExit()
            if not imgdir.endswith("/"):
                imgdir+="/"
            self.img_dir = imgdir

        if parser.has_option("global", "minwidth"):
            width = parser.getint("global", "minwidth")
            self.min_width = width
        if parser.has_option("global", "minheight"):
            height = parser.getint("global", "minheight")
            self.min_height = height
            
    def __init__(self, start_link=None):
        self._init_defaults()
        # Now load the config file to override defaults
        self._load_config()
        
        if start_link:
            self.start_link = start_link
        if not self.start_link:
            raise SystemExit("No start link is provided, exiting now...")
        links.put(self.start_link)
        self.semaphore = Semaphore(self.workers_limit)

    @gen.coroutine
    def run(self):
        # First start an debug server
        app = Application([(r"/", WebHandler)])
        server = HTTPServer(app)
        server.listen(self.port)
        
        idle_loops = 0
        while True:
            if imageurls.qsize()==0 and links.qsize()==0:
                print "Both link and image queues are empty now"
                idle_loops += 1
                if idle_loops == self.idle_wait_loops:
                    break
            else:
                idle_loops = 0 # clear the idle loop counter
                if imageurls.qsize()==0:
                    self.handle_links()
                elif links.qsize()==0:
                    self.handle_imageurls()
                else:
                    choices = [0]*self.link_priority +[1]*self.img_priority
                    choice = random.choice(choices)
                    if choice:
                        self.handle_imageurls()
                    else:
                        self.handle_links()
            yield gen.sleep(0.1 * self.politeness)
        # Wait for all link handlers
        links.join()
        # Handling imageurls if generated by the last few links
        while imageurls.qsize():
            self.handle_imageurls()
        imageurls.join()

    @gen.coroutine
    def handle_links(self):
        yield self.semaphore.acquire()
        newlink = yield links.get()
        
        # Make sure we haven't visited this one
        if newlink in visited_links:
            self.semaphore.release()
            raise gen.Return()
        visited_links.add(newlink)
        
        # use async client to fetch this url
        client = AsyncHTTPClient()
        tries = 3 # Give it 3 chances before putting it in failure
        while tries:
            response = yield client.fetch(newlink)
            if response.code==200:
                break
            tries -= 1
        
        # release the semaphore
        self.semaphore.release()
        if response.code!=200:
            link_failures.append(newlink)
            print "[FAILURE] - %s"%newlink
            raise gen.Return()

        # TODO: replace this with a report api
        print "[VISITED] - %s"%newlink

        # parse url to get the base url
        components = urlparse.urlparse(newlink)
        baseurl = components[0]+"://"+components[1]
        path = components[2]
        
        # parse the html with bs
        soup = bs4.BeautifulSoup(response.body)
        # extract valid links and put into links
        a_tags = soup.find_all("a")
        for tag in a_tags:
            if "href" not in tag.attrs:
                continue
            href = tag['href']
            if href.startswith("#"):
                continue
            if href.startswith("/"): # relative
                href = baseurl+href
            else:
                if not path.endswith("/"):
                    path = path[:path.rfind("/")+1]
                href = baseurl+"/"+path+href
            if not self.link_regex.match(href):
                continue
            if href in visited_links:
                continue
            links.put(href)
            print "NEWLINK:", href
        
        # extract imgs and put into imageurls
        img_tags = soup.find_all("img")
        for tag in img_tags:
            if "src" not in tag.attrs:
                continue
            src = tag['src']
            if src.startswith("/"): # relative
                src = baseurl+src
            if not self.img_regex.match(src):
                continue
            if src in downloaded_images:
                continue
            imageurls.put(src)
            print "NEW IMAGE:", src
                            
        # now the task is done
        links.task_done()

    @gen.coroutine
    def handle_imageurls(self):
        yield self.semaphore.acquire()
        imgurl = yield imageurls.get()

        if imgurl in downloaded_images:
            self.semaphore.release()
            raise gen.Return()
        # mark the image as downloaded
        downloaded_images.add(imgurl)
        
        # use async client to fetch this url
        client = AsyncHTTPClient()
        tries = 3 # Give it 3 chances before putting it in failure
        while tries:
            response = yield client.fetch(imgurl)
            if response.code==200:
                break
            tries -= 1
        # Download is finished, release semaphore
        self.semaphore.release()
        
        if response.code!=200:
            download_failures.append(imgurl)
            print "[FAILURE] - %s"%imgurl
            raise gen.Return()

        # TODO: replace this with a report api
        print "[DOWNLOADED] - %s"%imgurl
        
        # Read the file content
        img = PIL.Image.open(response.buffer)
        w, h = img.size
        if w <self.min_width or h <self.min_height:
            raise gen.Return()
        
        # find out the image extension, default to jpg
        if "." in imgurl:
            ext = imgurl.split(".")[-1].lower()
            if ext not in ["jpg", "png", "gif"]:
                ext = "jpg"
        elif img.format:
            ext = img.format.lower()
        else:
            ext = "jpg"
            
        # increment the counter
        global img_counter
        img_counter += 1
        fname = str(img_counter).zfill(self.fname_digits)+"."+ext
        fpath = self.img_dir + fname
        # save the image file
        f = open(fpath, "wb")
        f.write(response.body)
        
        # now the task is done
        imageurls.task_done()
        
        
def main():
    # allow user to provide start url from commandline
    if len(sys.argv)>1:
        crawler = Crawler(sys.argv[1])
    else:
        crawler = Crawler()
    
    IOLoop.current().run_sync(crawler.run)

    # TODO: replace with reporting api calls
    print "++++++++++++++++++++++++++++"
    print "%d links visited."%len(visited_links)
    print "%d images downloaded"%len(downloaded_images)
    print "Link failures:", link_failures 
    print "Image download failures:", download_failures
    print "++++++++++++++++++++++++++++"
    
if __name__ == "__main__":
    main()

