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 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:
        return
    tbl = profile.to_apache_tbl()
    write_apache_arrow(profile, tbl)
    os.remove(filename)
    return 

def main():
    bufr_files = glob.glob("/ldm/data/upperair/hires_sonde/*/*/*.bufr")
    print(len(bufr_files))
    #for bfile in bufr_files:
    #    process_file(bfile)
    with Pool(processes=12) as pool:
        pool.map(process_file, bufr_files)

if __name__ == "__main__":
    main()