import psycopg2.extras
import pandas as pd
import numpy as np
import requests
from getpass import getpass
from zeg.__main__ import main as zeg
import sys
import os
import urllib.parse
import logging


#TODO: This whole file is broken.

def upload_imageset(collection_obj, image_path_column, imageset_template, paths, project, token):
    imageset_yaml = imageset_template.format(paths=paths, path_column=image_path_column,
                                             collection_id=collection_obj['id'],
                                             dataset_id=collection_obj['dataset_id'])
    with open("imageset.yaml", "w") as text_file:
        text_file.write(imageset_yaml)
    logging.debug(imageset_yaml)
    with open("imageset-upload.sh", 'r') as imageset_upload_file:
        imageset_upload = imageset_upload_file.read()
    logging.debug(imageset_upload)
    args = imageset_upload.format(imageset_id=collection_obj['imageset_id'], token=token, project=project)
    sys.argv = args.split()
    zeg()


def upload_imageset_from_file(collection_obj, collection_name, token, project):
    input_filename = collection_name[:-4]
    input_file_path = os.path.join("/data_source", collection_name)

    lemnatec_df = pd.read_csv(input_file_path)

    with open("template-imageset.yaml", 'r') as imageset_template_file:
        imageset_template = imageset_template_file.read()

    camera_label = "RGB SV1"
    image_path_column = "{} image path".format(camera_label)
    if image_path_column not in lemnatec_df.columns:
        image_path_column = "RGB_3D_3D_side_far_0"

    for path in lemnatec_df[image_path_column]:
        t_path = os.path.join("/export_images/plantdb/tpa_backup", input_filename, path)
        logging.debug(t_path)
        logging.debug(os.stat(t_path))

    paths = "    - /export_images/plantdb/tpa_backup/" + input_filename + "/" + lemnatec_df[image_path_column].dropna()
    paths = paths.str.cat(sep="\n")
    upload_imageset(collection_obj, image_path_column, imageset_template, paths, project, token)

def upload_dataset_from_file(collection_name, dataset_upload_id, token, project):
    input_file_path = os.path.join("/data_source", collection_name)

    with open("template-dataset-file.yaml", 'r') as dataset_template_file:
        dataset_template = dataset_template_file.read()

    dataset_yaml = dataset_template.format(input_file_path=input_file_path)
    with open("dataset.yaml", "w") as text_file:
        text_file.write(dataset_yaml)

    with open("dataset-upload.sh", 'r') as dataset_upload_file:
        dataset_upload = dataset_upload_file.read()

    args = dataset_upload.format(dataset_upload_id=dataset_upload_id, token=token, project=project)
    sys.argv = args.split()
    zeg()










data_source = int(input("File [1] or Database[2]?"))

if data_source == SRC_FILE:
    files = sorted(os.listdir("/data_source"))
    for i, f in enumerate(files):
        print(i, f)
    file_selection = int(input("Select File:"))

    db_name = ""

    collection_name = files[file_selection]

    collection_obj = find_or_create_collection(token, db_name, collection_name, workspace_id)

    upload_dataset_from_file(collection_name, collection_obj['upload_dataset_id'], token, workspace_id)

    upload_imageset_from_file(collection_obj, collection_name, token, workspace_id)