from multiprocessing import Process
from multiprocessing import Queue
import multiprocessing
import os
import re
import sys
import time
import urllib2

# --- DESCRIPTION ---
# A script for circumventing the image "protection" from http://celebforum.to/
# specifically for imagevenue.com links (will also support imagebam in the future)
# 
# USAGE:
# python grabber.py <html source of gallery>


# http://img184.imagevenue.com/loc663/th_74708_0005_blood_orange_dress05_123_663lo.jpg
# http://img184.imagevenue.com/img.php?loc=loc663&image=74708_0005_blood_orange_dress05_123_663lo.jpg
IMAGEVENUE_PAGE_RE = r'"http://(img\d{3})\.imagevenue\.com/(loc\d{3})/(.+?\.jpg)"'
IMAGEVENUE_SRC_RE = r'<img id="thepic" .+? SRC="(.+?)"'

IMAGEBAM_PAGE_RE = r'"http://thumbnails\d+\.imagebam\.com/\d+/(.+?\.jpg)"'
IMAGEBAM_SRC_RE = r'<img .+? onclick="scale\(this\);" src="(.+?\.jpg)"'

# Location where the images should be saved
IMAGE_SAVE_PATH = 'images'


class ImageVenueProcessor(Process):

	def __init__(self, input_file_name, image_src_queue):
		Process.__init__(self)
		self.input_file_name = input_file_name
		self.image_src_queue = image_src_queue
		self.running = False

	def run(self):
		self.running = True

		# Iterate over the input source file and look for all thumbnail image matches
		matches = []
		with open(self.input_file_name, 'r') as input_file:
			for line in input_file:
				line_matches = re.findall(IMAGEVENUE_PAGE_RE, line)
				for line_match in line_matches:
					matches.append(line_match)

		# Iterate over each match and build the correct image page URL
		# Major thanks to http://vintage-erotica-forum.com/showthread.php?p=694632
		# for explaining the secret sauce!
		total = len(matches)
		for idx, match in enumerate(matches):
			server = 'http://{:s}.imagevenue.com/'.format(match[0])
			image_name = match[2][3:]
			image_loc = 'img.php?loc={:s}&image={:s}'.format(match[1], image_name)
			link = '{:s}{:s}'.format(server, image_loc)

			# Open the image page and download the source HTML
			response = urllib2.urlopen(link)
			html = response.read()

			# Iterate over the source HTML looking for the image source link
			for line in html.split('\n'):
				img = re.search(IMAGEVENUE_SRC_RE, line)
				if img:
					location = '{:s}{:s}'.format(server, img.group(1))
					print '{:3d} of {:3d} - {:s}'.format(idx + 1, total, location)
					self.image_src_queue.put(location)

		self.running = False

class ImageBamProcessor(Process):

	def __init__(self, input_file_name, image_src_queue):
		Process.__init__(self)
		self.input_file_name = input_file_name
		self.image_src_queue = image_src_queue
		self.running = False

	def run(self):
		self.running = True

		# Iterate over the input source file and look for all thumbnail image matches
		matches = []
		with open(self.input_file_name, 'r') as input_file:
			for line in input_file:
				line_matches = re.findall(IMAGEBAM_PAGE_RE, line)
				for line_match in line_matches:
					matches.append(line_match)

		# Iterate over each match and build the correct image page URL
		# Thanks to http://board.jdownloader.org/showthread.php?t=26255
		total = len(matches)
		for idx, match in enumerate(matches):
			image_name = match
			server = 'http://www.imagebam.com/image/'
			link = '{:s}{:s}'.format(server, image_name)

			# Open image page and download the source HTML
			response = urllib2.urlopen(link)
			html = response.read()

			# Iterate over the source HTML looking for the image source link
			for line in html.split('\n'):
				img_match = re.search(IMAGEBAM_SRC_RE, line)
				if img_match:
					location = img_match.group(1)
					print '{:3d} of {:3d} - {:s}'.format(idx + 1, total, location)
					self.image_src_queue.put(location)

		self.running = False

class Downloader(Process):

	def __init__(self, image_src_queue):
		Process.__init__(self)
		self.image_src_queue = image_src_queue

	def run(self):

		while True:
			image_path = self.image_src_queue.get()
			image_name = image_path.rsplit('/', 1)[1]

			print 'Downloading "{:s}"...'.format(image_path),

			# Check if the image has already been downloaded before 
			# attempting to download it
			contents = os.listdir(IMAGE_SAVE_PATH)
			if image_name not in os.listdir(IMAGE_SAVE_PATH):

				# Download the freakin image!
				try:
					response = urllib2.urlopen(image_path, timeout=4)
					data = response.read()
				except:
					# print 'Unable to download: {:s}'.format(location)
					print 'ERROR'
					continue

				# Save the image to the desired location
				with open(os.path.join(IMAGE_SAVE_PATH, image_name), 'wb') as new_image:
					new_image.write(data)

				print 'DONE'

			else:
				print 'SKIPPED'

def main():

	if len(sys.argv) < 2:
		print 'Not enough arguments specified!'
		return -1

	# Create the image source queue that the "Processing" objects will produce
	# to and the "Downloader" objects will consume from
	image_src_queue = Queue(2048)

	# Based on the number of cores in the system determine how many 
	# Downloader worker objects should be created
	num_workers = max(1, multiprocessing.cpu_count() - 2)

	# Create the Downloader worker objects
	dl_workers = []
	for ctr in xrange(num_workers):
		dl_worker = Downloader(image_src_queue)
		dl_worker.start()
		dl_workers.append(dl_worker)

	# Create the ImageBam Processsor
	ib_proc = ImageBamProcessor(sys.argv[1], image_src_queue)
	ib_proc.start()

	# Create the ImageVenue Processor
	iv_proc = ImageVenueProcessor(sys.argv[1], image_src_queue)
	iv_proc.start()

	# Wait for the Processor objects to complete
	ib_proc.join()
	iv_proc.join()

	# Wait until the image source queue is empty, then kill all the
	# Downloader workers
	while True:
		if image_src_queue.empty():
			# Wait 3 seconds to allow any downloads to finish
			time.sleep(3)
			for dl_worker in dl_workers:
				dl_worker.terminate()
			break

	print '\n\n'


if __name__ == '__main__':
	sys.exit(main())