import os
import shutil
import glob
import json
import csv
import yaml
from git import Repo 
import gzip


fieldnames = ['original','compressed']

def to_csv_record(writer, buffer):
  record = json.loads(buffer)
  writer.writerow(dict(
    original=record['graph']['sentence'],
    compressed=record['compression']['text']))

def build_dataset(rawdata_dir, preprocessed_data_dir):
    print("Data Preparation...")
    os.makedirs(preprocessed_data_dir, exist_ok=True)
    with open(os.path.join(preprocessed_data_dir, 'training_data.csv'),'w') as csvfile:
      writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
      writer.writeheader()
      for rawdata_files in glob.glob(f'{rawdata_dir}/data/**train**.json'):
        with open(rawdata_files) as raw_contents:
          buffer = ''
          for line in raw_contents:
            if line.strip()=='':
                to_csv_record(writer, buffer)
                buffer = ''
            else:
              buffer += line
          if len(buffer)>0: 
            to_csv_record(writer, buffer)

    with open(os.path.join(preprocessed_data_dir, 'eval_data.csv'),'w') as csvfile:
      writer = csv.DictWriter(csvfile, fieldnames=['original','compressed'])
      writer.writeheader()
      with open(f'{rawdata_dir}/data/comp-data.eval.json') as raw_contents:
        buffer = ''
        for line in raw_contents:
          if line.strip()=='':
              to_csv_record(writer, buffer)
              buffer = ''
          else: buffer += line
        if len(buffer)>0: to_csv_record(writer, buffer)
    
def decompressing_rawdata(rawdata_dir):
    print("Decompression...")
    compressed_files = glob.glob(rawdata_dir + "/data/*.json.gz")
    for compressed_file_path in compressed_files:
        output_file_path = os.path.splitext(compressed_file_path)[0] 
        with gzip.open(compressed_file_path, 'rb') as comp_file:
            compressed_content = comp_file.read()
        with open(output_file_path, 'wb') as output_file:
            output_file.write(compressed_content)
        os.remove(compressed_file_path)

def download_rawdata(git_url, rawdata_dir):
    os.makedirs(rawdata_dir, exist_ok=True)
    print("Data Cloning...")
    current_dir = os.getcwd()
    try:
        os.chdir(rawdata_dir)
        Repo.clone_from(git_url, '.')
    except Exception as e:
        print("Error:", e)
    finally:
        os.chdir(current_dir)
    
if __name__ == "__main__":
    config = yaml.safe_load(open("config.yaml", "r"))
    PROJECT_DIR = eval(config["SENTENCE_COMPRESSION"]["PROJECT_DIR"])
    rawdata_git = config["SENTENCE_COMPRESSION"]["DATA"]["RAW_DATA"]
    preprocessed_data_dir = os.path.join(PROJECT_DIR, config["SENTENCE_COMPRESSION"]["DATA"]["CLEAN_DATA"])
    rawdata_dir = os.path.join(PROJECT_DIR, config["SENTENCE_COMPRESSION"]["DATA"]["RAW_DIR"])
    download_rawdata(rawdata_git, rawdata_dir)
    decompressing_rawdata(rawdata_dir)
    build_dataset(rawdata_dir, preprocessed_data_dir)
    shutil.rmtree(rawdata_dir)