import numpy as np
from BUFRObserved import BUFRObservedProfile
import gcsfs
import pyarrow as pa
from pyarrow import fs
import pyarrow.dataset as ds
import pyarrow.parquet as pq
from multiprocessing import Pool

import shutil
import argparse
import requests
import glob
import bs4
import os

from fsspec.asyn import sync_wrapper
class GCSFileSystem(gcsfs.GCSFileSystem):
    # https://github.com/fsspec/gcsfs/issues/404
    def _mkdir(self, bucket, acl="projectPrivate", default_acl="bucketOwnerFullControl", **kwargs):
        return super()._mkdir(bucket, acl, default_acl)
    mkdir = sync_wrapper(_mkdir)

PATH = "/mrms-tiledata/BUFR-archive"

def read_processed_files():
    """
    Reads a text file that tells us what we've processed so far
    in the tar archive. 
    """
    files = []
    with open("./processed_files.txt", "r") as pfiles:
        files = pfiles.readlines()
    files = [f.strip("\n") for f in files]
    return files

def get_available_global_bufr_tar():
    """
    Get the list of available B_BUFR tar files from NCEI
    and return a list of their URLs.
    """
    processed_files = read_processed_files()
    url = "https://www.ncei.noaa.gov/data/nws-global-upper-air-bufr/archive/"
    r = requests.get(url)
    data = bs4.BeautifulSoup(r.text, "html.parser")
    tarfiles = []
    for l in data.find_all("a"):
        data_url = url + l["href"]
        if l["href"] not in processed_files and l["href"].endswith("gz"):
            tarfiles.append(url+l["href"])
    return tarfiles

def download_tar(url):
    """
    Download a tar file from the given URL to the 
    unprocessed tar directory.
    """
    print("Downloading: {url}".format(url=url))
    os.system("cd {path}; wget {url}".format(path=PATH+"/tar/unprocessed", url=url))
    return os.path.join(PATH, "tar", "unprocessed", url.split("/")[-1])

def untar(tar_path, working_dir):
    """
    Untar a B_BUFR archive into the given working dir
    """
    os.system("cd {working_dir}; tar -xzvf {tar}".format(tar=tar_path, working_dir=working_dir))

def split_bufr(filename):
    with open(filename, "rb") as bufrfile:
        data = bufrfile.read()
        delim = b"7777"
        split_bufrs = [b+delim for b in data.split(delim) if b]
        count = 0
        for bufr in split_bufrs:
            with open("{filename}.{count}".format(filename=filename, count=count), "wb") as out:
                out.write(bufr)
            count += 1
    os.remove(filename)

def write_apache_arrow(profile, table):

    fs = GCSFileSystem(project="sharp-328923", token="/home/ldm/.config/gcloud/legacy_credentials/wxbyte@gmail.com/adc.json")
    table_names = ["Site ID", "WMO ID", "Site Name", "State", "Country", "Latitude", "Longitude", "Elevation", "Flag"]

    partition = ds.partitioning(pa.schema([
        ("year", pa.int32()),
        ("month", pa.int32()),
        ("day", pa.int32()),
        ("hour", pa.int32()),
        ("WMO", pa.string())
    ]))

    hires_path = "global-radiosondes/hires-sonde/"
    lores_path = "global-radiosondes/lores-sonde/"
    bucket_path = None

    ## Some of the BUFR profiles are not high resolution, 1-second profiles.
    ## This is a kind of naive check, as theres no lon/lat data for such
    ## profiles
    if len(profile["lon"]) < 200 and np.isnan(profile["lon"]).all():
        bucket_path = lores_path
    else:
        bucket_path = hires_path

    pq_format_opts = ds.ParquetFileFormat().make_write_options(compression="SNAPPY") 
    ds.write_dataset(table, bucket_path, filesystem=fs, format="parquet",\
                        partitioning=partition, existing_data_behavior='overwrite_or_ignore',
                        file_options=pq_format_opts)

def process_file(filename):
    print("Working on: ", filename)
    profile = BUFRObservedProfile(filename)
    if profile.profile is None:
        os.remove(filename)
        return
    tbl = profile.to_apache_tbl()
    write_apache_arrow(profile, tbl)
    os.remove(filename)
    return 

def main():
    tarfiles = sorted(get_available_global_bufr_tar())
    for tarfile in tarfiles:
        tarfile = download_tar(tarfile)
        untar(tarfile, os.path.join(PATH, "working_dir"))

        bufr_files = glob.glob(os.path.join(PATH, "working_dir", "*.bin"))
        for bfile in bufr_files:
            split_bufr(bfile)
            split_files = glob.glob("{bfile}.*".format(bfile=bfile))
            #for filename in split_files:
            #    process_file(bfile)
            with Pool(processes=12) as pool:
                pool.map(process_file, split_files)

        with open("processed_files.txt", "a+") as textfile:
            textfile.write(os.path.split(tarfile)[-1] + "\n")
        os.remove(tarfile)

def test():
    filename = "working_dir/2017060700_GTS_sn.0011.bin"
    split_bufr(filename)


if __name__ == "__main__":
    main()
    #test()