import os
import shutil
import boto3
from s3urls import parse_url

s3 = boto3.resource('s3')

data_dir_root_path = './data'
gold_file_local_path = os.path.join(data_dir_root_path, 'gold.txt')
pred_file_local_path = os.path.join(data_dir_root_path, 'pred.txt')
db_dir_local_root_path = data_dir_root_path
table_file_local_path = os.path.join(data_dir_root_path, 'tables.json')
benchmark_file_local_path = os.path.join(data_dir_root_path, 'benchmark.csv')
model_meta_local_path = os.path.join(data_dir_root_path, 'model.yaml')
dataset_meta_local_path = os.path.join(data_dir_root_path, 'dataset.yaml')
template_local_path =  os.path.join(data_dir_root_path, 'template.txt')

def download_file(s3_url:str, local_root_path:str, local_file_name:str):
    parsed_s3_url = parse_url(s3_url)
    print('download file', parsed_s3_url)

    if local_root_path == '' :
        local_root_path = './'
    
    if local_file_name == '' :
        local_file_name = os.path.basename(parsed_s3_url['key'])

    download_file_path = os.path.join(local_root_path, local_file_name)
    s3.Bucket(parsed_s3_url['bucket']).Object(parsed_s3_url['key']).download_file(download_file_path)

    return download_file_path

def upload_file(local_path:str, s3_url:str):
    parsed_s3_url = parse_url(s3_url)
    print('upload file', parsed_s3_url)

    s3.Bucket(parsed_s3_url['bucket']).upload_file(local_path, parsed_s3_url['key'])

def file_process(path:str, local_path:str):
    if os.path.exists(path) is True:
        return path

    local_root_path = os.path.dirname(local_path)
    if os.path.exists(local_root_path) is False:
        os.makedirs(local_root_path)

    local_file_name = os.path.basename(local_path)
    s3_url = path
    download_file(s3_url, local_root_path, local_file_name)

    return local_path

def db_dir_process(path:str, local_root_path:str):
    if os.path.isdir(path) is True:
        return path

    zip_file_path = ''
    if os.path.exists(path) is True:
        zip_file_path = path
    else:
        s3_url = path
        zip_file_path = download_file(s3_url, local_root_path, '')

    shutil.unpack_archive(zip_file_path, local_root_path)

    return os.path.join(local_root_path, 'database')

def process_file(db:str, table:str, benchmark:str, model_meta:str, dataset_meta:str, template:str, enable_sqlite: bool = False):
    gold_path = gold_file_local_path
    pred_path = pred_file_local_path
    db_dir = ''

    if os.path.exists(data_dir_root_path) is False:
        os.makedirs(data_dir_root_path)

    if enable_sqlite is True:
        db_dir = db_dir_process(db, db_dir_local_root_path)

    table_path = file_process(table, table_file_local_path)
    benchmark_path = file_process(benchmark, benchmark_file_local_path)

    model_meta_path = file_process(model_meta, model_meta_local_path)
    dataset_meta_path = file_process(dataset_meta, dataset_meta_local_path)
    template_path = file_process(template, template_local_path)

    return {
        "db": db_dir,
        "gold": gold_path,
        "pred": pred_path,
        "table": table_path,
        "benchmark": benchmark_path,
        "model_meta": model_meta_path,
        "dataset_meta": dataset_meta_path,
        "template": template_path,
    }

def process_file_for_question_evaluation(model_meta:str, dataset_meta:str, template:str):
    if os.path.exists(data_dir_root_path) is False:
        os.makedirs(data_dir_root_path)

    model_meta_path = file_process(model_meta, model_meta_local_path)
    dataset_meta_path = file_process(dataset_meta, dataset_meta_local_path)
    template_path = file_process(template, template_local_path)

    return {
        "model_meta": model_meta_path,
        "dataset_meta": dataset_meta_path,
        "template": template_path,
    }

if __name__ == "__main__":
    import argparse

    # parse args to var
    parser = argparse.ArgumentParser()
    parser.add_argument('--gold', dest='gold', type=str,
                        help="the path to the gold queries")
    parser.add_argument('--db', dest='db', type=str,
                        help="the directory that contains all the databases and test suites")
    parser.add_argument('--table', dest='table', type=str,
                        help="the tables.json schema file")
    parser.add_argument('--benchmark', dest='benchmark', type=str,
                        help="the benchmark file")
    parser.add_argument("--model_meta",dest='model_meta', type=str, default='undefined',
                        help="meta of model")
    parser.add_argument("--dataset_meta",dest='dataset_meta', type=str, default='undefined',
                        help="meta of dataset")
    parser.add_argument("--template",dest='template', type=str, default='undefined',
                        help="template")


    args = parser.parse_args()
    # s3 path to native path
    processed_path = process_file(args.db, args.table, args.benchmark, args.model_meta, args.dataset_meta, args.template, False)
    print(processed_path)