import os
import urllib.request
import urllib.parse
import re
import time
import random
import datetime
import json
from threading import Thread
import queue

import tqdm

directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/"
stored_images_dir = r"..\Data"
# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/suvi-l2-ci094/"
# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16\l2\suvi-l2-ci094"
ignore_folder_names = ["l1b", "goes17", "2017", "2018", "2019", "2020"]
file_database_path = r"..\file_database.json"
fetch_interval = 0 # 60*60
nworkers = 8
randomize_order = False

def recursive_find_links(url, attempt_count = 3, randomize_order = True):
    links_regex_pattern = r'(?<=<a href=")([^ ?:]*)(?=">.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace
    times_regex_pattern = r'(?<=<td align="right">)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm
    attempts = 0
    html_content = None
    while attempts < attempt_count:
        try:
            with urllib.request.urlopen(url) as response:
                html_content = response.read().decode('utf-8')
                break
        except Exception as e:
            print(f"Exception while fetching links: {e}")
            time.sleep(1 + random.random())
            attempts += 1
        if (attempt_count > 1) and (attempts == attempt_count):
            print(f'\nAfter {attempt_count} retries, could not fetch: {url}')
            return

    links = re.findall(links_regex_pattern, html_content)
    times = re.findall(times_regex_pattern, html_content)
    for i in range(len(times)):
        dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M")
        dt.replace(tzinfo=datetime.timezone.utc)
        times[i] = time.mktime(dt.timetuple())
    
    if len(links) - 1 == len(times):
        links = links[1:]
    elif len(links) == len(times):
        pass
    else:
        raise(ValueError)

    results = list(zip(links, times))
    if randomize_order:
        random.shuffle(results)
    for _link, _time in results:
        if _link.endswith("/"):
            if _link.split(r"/")[-2] in ignore_folder_names:
                continue
            else:
                yield from recursive_find_links(urllib.parse.urljoin(url,_link), attempt_count, randomize_order)
        else:
            yield urllib.parse.urljoin(url,_link), _time
    

# Fetch image from url and store it to path, retrying on failure
def image_fetch_worker(work_queue, result_queue, attempt_count = 1):
    while True:
        job = work_queue.get()
        if job == None:
            return
        url, path, t = job

        attempts = 0
        while attempts < attempt_count:
            try:
                req = urllib.request.Request(url, data=None)
                image_data = urllib.request.urlopen(req).read()
                os.makedirs(os.path.split(path)[0], exist_ok=True)
                open(path, 'wb').write(image_data)
                result_queue.put((True, url, t))
                break
            except Exception as e:
                if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov)
                    attempts += attempt_count
                elif attempts == 0:
                    print(f"\nA problem occurred on image: {url} | {e}")
                    time.sleep(1 + random.random())
                    attempts += 1
            if (attempt_count > 1) and (attempts == attempt_count):
                print(f'\nAfter {attempt_count} retries, could not fetch: {url}')
                result_queue.put((False, url, t))
                break


if __name__ == "__main__":
    file_info_cache = {}
    try:
        print(f"Attempting to load file records from cache: {file_database_path}")
        with open(file_database_path, 'r') as f:
            file_info_cache = json.loads(f.read())
        print(f"File records loaded from cache: {len(file_info_cache)} records found.")
    except Exception as e:
        print(f"Load failed, starting with empty cache")
        file_info_cache = {}

    work_queue = queue.Queue(maxsize=nworkers)
    result_queue = queue.Queue()
    workers = []
    for _ in range(nworkers):
        t = Thread(target=image_fetch_worker, args=(work_queue, result_queue), daemon=True)
        t.start()
        workers.append(t)
    try:
        while True:
            fetched_image_count = 0
            already_had_image_count = 0
            failed_image_count = 0
            urllen = len(directory_url)
            for l, t in tqdm.tqdm(recursive_find_links(directory_url, randomize_order = randomize_order), desc="Downloading files"):
                # Collect completed jobs and record completion status
                while True:
                    try:
                        r_success, r_url, r_t = result_queue.get_nowait()
                        if r_success: 
                            fetched_image_count += 1
                            file_info_cache[r_url] = r_t
                        else:
                            failed_image_count += 1
                    except queue.Empty:
                        break
                # If we dont have the file or the file at the link is newer than the one we previously fetched
                if (not (l in file_info_cache)) or t > file_info_cache[l]: 
                    file_portion_of_link = l[urllen:]
                    filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep))
                    if os.path.exists(filepath):
                        if l in file_info_cache: # If we have record of this file, it must be out of date, rename it and download the new version.
                            new_name = filepath + f"_{int(file_info_cache[l])}"
                            if os.path.exists(new_name):
                                os.remove(new_name)
                            os.rename(filepath, new_name)
                        else: # If we have no record of this file, update the file info cache and don't redownload
                            file_info_cache[l] = t
                            already_had_image_count += 1
                            continue
                    work_queue.put((l, filepath, t))
                else:
                    already_had_image_count += 1
            with open(file_database_path, 'w') as f:
                f.write(json.dumps(file_info_cache))
            print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")
            if fetch_interval > 0:
                print(f"Run complete!, sleeping for {fetch_interval} seconds.")
                with open(file_database_path, 'w') as f:
                    f.write(json.dumps(file_info_cache))
                time.sleep(fetch_interval)
            else:
                print("Run complete!, exiting...")
                break
    except KeyboardInterrupt:
        print("Saving file database and shutting down.")

    for _ in range(nworkers):
        try:
            work_queue.put(None, timeout=5.0)
        except:
            break

    for w in workers:
        w.join(5.0)

    with open(file_database_path, 'w') as f:
        f.write(json.dumps(file_info_cache))

    print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")