# -*- coding: utf-8 -*-

import os
import requests
import wget
import numpy as np
import pandas as pd
import json
import time
from threading import Thread
from queue import Queue
from bs4 import BeautifulSoup
from tqdm import tqdm


# CLASS Downloader

class Downloader(object):
	"""A downloader using multi-threading tech

	----------------------------------------
	source: A list of zip of names and urls, like:
		[(name, url), (name, url), ... (name, url)]
		A source could be generated by calling "list(zip(name_list, url_list))",
		where "name_list" is a list of names and "url_list" is a list of urls.

	file_extension: 'normal' by default or other file extensions like
		'jpg', 'png', 'mp4', etc.
		'normal': Automatically recognize the file extension.
	"""
	def __init__(self, source, file_extension='normal'):
		self.source = source
		self.file_extension = file_extension
		self.crawler = Crawler(self.download_single_file, self.source)
		self.check('initialization')

	def check(self, item):
		# initial check
		if item == "initialization":
			self.num_files = len(self.source)
		# check if the output folder already exists.
		elif item == "out_folder":
			if os.path.isdir(self.out_folder):
				print("[WARNING]: Folder '%s' already exists. Automatically download into this folder." % self.out_folder)
				return True
			else:
				os.mkdir(self.out_folder)

	def download_single_file(self, index, source_item):
		name, url = source_item
		if self.file_extension == 'normal':
			filex = url.split('.')[-1]
		else:
			filex = self.file_extension
		out_path = "./%s/%s.%s" %(self.out_folder, name, filex)

		if os.path.isfile(out_path):
			print("\r[INFO]: Arealy has this file, skipping...".ljust(73), end="")
		else:
			if self.engine == "wget":
				wget.download(url, out_path)
			elif self.engine == "wget-cli":
				os.system('wget %s -O %s' %(url, out_path))
			elif self.engine == "you-get":
				os.system('you-get %s -O %s' %(url, out_path))

	def download(self, out_folder="./", engine="wget"):
		self.out_folder = out_folder
		self.check('out_folder')
		self.engine = engine
		self.crawl()

	def crawl(self):
		self.crawler.crawl()
		print("[INFO]: %d urls failed." %(len(self.crawler.failed_urls), ))
		if len(self.crawler.failed_urls) > 0:
			print("[FAILED]: \n", self.crawler.failed_urls)
			go_on = input("[INPUT]: Recrawl the faield urls, (y/n): ")
			if go_on == 'y':
				self.crawler = Crawler(self.download_single_file, self.crawler.failed_urls)
				self.crawl()
			else:
				print('[INFO]: Task done.')
				return


# A TEST FOR CLASS Downloader

def test_downloader():
	names = [str(i) for i in range(100)]
	urls = ['https://www.baidu.com/img/baidu_resultlogo@2.png']*100
	source = list(zip(names, urls))

	downloader = Downloader(source)
	downloader.download('test', engine="wget")




# CLASS Crawler

class Crawler(object):
	"""A wapper of multi-threading crawler for human

	----------------------------------------
	basic_func: Any function that you need it to be multi-threaded.
		Note that the parameters of this function are fixed, which
		means you should define this function like:
			def basic_func_name(index, src_item):
				pass
		where "index" represents the index of the "src_item", and
		"src_item" represents the information to be handled in this
		function.

	urls: A list of sources that you'd like to be handled multi-threadingly
		Given that crawler always request urls, so typically the source
		list is often a list of urls.
		Note that this parameter could be anything iterable, as long
		as you'd like to deal with the item in your basic_func.

	has_result: (Boolean) "False" by default.
		Set this to be "True" when you'd like to export data after
		all of crawling task done. And remember to return the data you
		wanted in your basic_func.
		However, this way of exporting data is not recommended. The better
		way is to export your data in your basic_func, so that every time
		the basic_func finished in multi-threading tasks, the associated
		data will be exported to your own dataset.
	"""

	def __init__(self, basic_func, urls, has_result=False, batch_size=None):
		super(Crawler, self).__init__()
		self.basic_func = basic_func
		self.has_result = has_result
		self.batch_size = batch_size
		self.urls = urls
		self.num_urls = len(self.urls)
		self.failed_urls = []

	def crawl(self):
		# initialization
		print("[INFO]: %d urls in total." %(self.num_urls))
		if self.batch_size is None:
			batch_size = int(input("[INPUT]: BATCH SIZE: "))
		else:
			batch_size = self.batch_size
			print("[INFO]: Using BATCH SIZE: %d" %(batch_size))
		self.num_thrd = int(np.ceil(self.num_urls / batch_size))
		if self.has_result == True:
			self.queue = Queue()

		thds = []
		t1 = time.time()
		# for i in range(self.num_thrd):
		for i in tqdm(range(self.num_thrd), desc="[INFO]: Open threads"):
			frm = i * batch_size
			# batch_size of the last thread
			if i == self.num_thrd - 1 and self.num_urls % batch_size != 0:
				batch_size = self.num_urls % batch_size
			thd = Thread(target = self.crawl_batch,
						 args=(frm, batch_size))
			thd.start()
			thds.append(thd)

		for thd in thds:
			thd.join()
		print("\r[INFO]: Task done.".ljust(74))
		print("[INFO]: The task costs %.4f sec." %(time.time()-t1))

		if self.has_result == True:
			# Load data
			result = []
			for i in tqdm(range(len(thds)), desc="[INFO]: Load data"):
				result += self.queue.get()
			print("[INFO]: All data are loaded.")
			return result

	def crawl_batch(self, frm, batch_size):
		thd_num = frm // batch_size
		urls = self.urls[frm: frm + batch_size]
		batch_result = []
		for i, url in enumerate(urls):
			try:
				res = self.basic_func(frm + i, url)
				print("\r[INFO]: Thread %d, url %d is done.".ljust(70)
					%(thd_num + 1, i+1), end='')
			except:
				res = None
				self.failed_urls.append(url)
				print('\r[ERROR]: Thread %d, url %d is failed, which is stored in failed_urls'.rjust(70)
				 	%(thd_num + 1, i+1), end='')
			batch_result.append(res)

		if self.has_result == True:
			self.queue.put(batch_result)


if __name__=="__main__":
	test_downloader()
