""" Download Pushshift Data """

import os

from urllib import request as req

import re

import pycurl

import hashlib

# Define values

URLS=["https://files.pushshift.io/reddit/comments/","https://files.pushshift.io/reddit/submissions"]

BZ2_LINK_RE_PATTERN = r"<a\s.*href=[\"'](\S+)[\"'][^>]*>\S*.bz2<\/a>"

SHA256_LINK_RE_PATTERN = r"<a\s.*href=[\"'](\S+)[\"'][^>]*>sha256\S*<\/a>"

OUTPUT_DIR = "reddit_submissions"

# Define functions

def main():

    """The main entrypoint."""

    for BASE_URL in URLS:

        submissions_page=req.urlopen(BASE_URL).read().decode("utf-8")

        # Get BZ2 Links

        raw_links = re.findall(BZ2_LINK_RE_PATTERN,submissions_page)

        filtered_links = [link[2:] for link in raw_links if link.startswith("./")]

        individual_links = list(set(filtered_links))

        # Download files

        if not os.path.exists(OUTPUT_DIR):

            os.makedirs(OUTPUT_DIR)

        else:

            # get first match and remove the ./ from the start of the link

            sha256_link = re.findall(SHA256_LINK_RE_PATTERN,submissions_page)[0][2:]

            hash_file=(req.urlopen("%s/%s"%(BASE_URL,sha256_link)).read().decode("utf-8"))

            hash_file_pairs=({entry.split("  ")[1]:entry.split("  ")[0] for entry in hash_file.split("\n") if len(entry.split("  "))>1})

            for file in hash_file_pairs.keys():

                file_path=os.path.join(OUTPUT_DIR,file)

                if(os.path.exists(file_path) and hashlib.sha256(open(file_path,'rb').read()).hexdigest()!=hash_file_pairs[file.split("/")[-1]]):

                    print("File is corrput, deleting %s"%file_path)

                    os.remove(file_path)

        curl = pycurl.Curl()

        for link in sorted(individual_links):

            url = BASE_URL + "/" + link

            file_path=os.path.join(OUTPUT_DIR, link)

            if not os.path.exists(file_path):

                with open(file_path, "wb") as file:

                    curl.setopt(curl.URL, url)

                    curl.setopt(curl.WRITEDATA, file)

                    curl.perform()

                print("Downloaded", link)

        curl.close()

# Execute main function

if __name__ == "__main__":

    main()