#!python

from PW_explorer.export import PWEExport
from PW_explorer.helper import (
    pw_slicer,
    rel_slicer,
)
from PWE_CLI.helper import (
    load_from_temp_pickle,
    set_current_project_name,
    get_current_project_name,
    get_save_folder,
    save_to_txt_file,
)

import argparse
import os


def __main__():

    ASP_TRIPLES_DB_OPTIONS = ['csv', 'sql', 'pkl']
    ASP_TRIPLES_OPTIONS = ['str'] + ASP_TRIPLES_DB_OPTIONS

    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--project_name", type=str, help="provide session/project name used while parsing")
    parser.add_argument("-s", "--schema", action="store_true", help="generate sql schemas", default=False)
    parser.add_argument("-sql", action="store_true", help="include if you want to export a sql db", default=False)
    parser.add_argument("-csv", action="store_true", help="include if you want to export in csv", default=False)
    parser.add_argument("-pkl", action="store_true", help="include if you want to export in pickle format",
                        default=False)
    parser.add_argument("-asp_simple", action='store_true', help="Output the facts as asp simple tuples in a .lp4 file",
                        default=False)
    parser.add_argument("-asp_triples", "--asp_triples", type=str, choices=ASP_TRIPLES_OPTIONS,
                        help="Output the facts as asp facts in form of triples. Choose a valid format. 'str' "
                             "outputs to a asp_triples.lp4 file", default=None)
    parser.add_argument("-rels_to_use", type=str, nargs='*', default=[],
                        help="Enlist the relation names to output. All are used by default.")
    parser.add_argument("-pws_to_use", type=int, nargs='*', default=[],
                        help="Enlist the pw_ids to output. All are used by default.")
    parser.add_argument("--rel_facts_with_arity", dest='rel_facts_with_arity', action='store_true',
                        help="Use with -asp_simple or -asp_triples to include arity information in exported asp "
                             "relation names.")
    parser.add_argument("--no-rel_facts_with_arity", dest='rel_facts_with_arity', action='store_false',
                        help="Use with -asp_simple or -asp_triples to include arity information in exported asp "
                             "relation names.")
    parser.set_defaults(rel_facts_with_arity=False)
    parser.add_argument("--include_pw_ids", dest='include_pw_ids', action='store_true',
                        help="Use with -asp_simple or -asp_triples to include pw_ids in exported asp "
                             "relation facts.")
    parser.add_argument("--no-include_pw_ids", dest='include_pw_ids', action='store_false',
                        help="Use with -asp_simple or -asp_triples to include pw_ids in exported asp "
                             "relation facts.")
    parser.set_defaults(include_pw_ids=True)
    args = parser.parse_args()

    project_name = ''
    if args.project_name is None:
        project_name = get_current_project_name()
        if project_name is None:
            print("Couldn't find current project. Please provide a project name.")
            exit(1)
    else:
        project_name = args.project_name

    export_to_sql = args.sql
    export_to_csv = args.csv
    export_to_pkl = args.pkl

    dfs = load_from_temp_pickle(project_name, 'dfs')
    pws = load_from_temp_pickle(project_name, 'pws')

    if args.pws_to_use:
        dfs, pws = pw_slicer(dfs, pws, args.pws_to_use)
    if args.rels_to_use:
        dfs, _, pws = rel_slicer(dfs, None, pws, args.rels_to_use)

    if export_to_sql:
        if PWEExport.export_to_sqlite_db(get_save_folder(project_name, 'sql_export'), dfs, project_name):
            print("Successfully exported to sql")
        else:
            print("SQL Export Failed")
    if export_to_csv:
        if PWEExport.export('csv', get_save_folder(project_name, 'csv_export'), dfs):
            print("Successfully exported to csv")
        else:
            print("CSV Export Failed")
    if export_to_pkl:
        if PWEExport.export('pkl', get_save_folder(project_name, 'pkl_export'), dfs):
            print("Successfully exported to pkl")
        else:
            print("PICKLE Export Failed")

    if args.schema:
        schemas = PWEExport.get_sqlite_schema(dfs)
        print('\n'.join(schemas))

    if args.asp_simple:
        meta_data = load_from_temp_pickle(project_name, 'meta_data')
        attr_defs = meta_data['attr_defs'] if 'attr_defs' in meta_data else {}
        asp_facts = PWEExport.export_as_asp_facts(pws=pws, rel_facts_with_arity=args.rel_facts_with_arity,
                                  include_pw_ids=args.include_pw_ids, attr_defs=attr_defs)
        save_to_txt_file(asp_facts, os.path.join(get_save_folder(project_name, 'asp_simple'), 'asp_simple.lp4'))
        print("Succesfully exported as asp simple facts.")

    if args.asp_triples:
        out_type = args.asp_triples
        if out_type in ASP_TRIPLES_DB_OPTIONS:
            out_type = 'db'
        asp_facts = PWEExport.export_as_asp_triples(pws=pws, rel_facts_with_arity=args.rel_facts_with_arity,
                                                    include_pw_ids=args.include_pw_ids, output_type=out_type)

        if args.asp_triples == 'str':
            save_to_txt_file(asp_facts, os.path.join(get_save_folder(project_name, 'asp_triples'), 'asp_triples.lp4'))
            print("Successfully exported in asp triples format.")
        elif args.asp_triples in ['csv', 'pkl']:
            PWEExport.export(args.asp_triples, get_save_folder(project_name, 'asp_triples'), {'asp_triples': asp_facts})
            print("Successfully exported in asp triples format.")
        elif args.asp_triples == 'sql':
            PWEExport.export_to_sqlite_db(get_save_folder(project_name, 'asp_triples'), {'asp_triples': asp_facts},
                                          project_name)
            print("Successfully exported in asp triples format.")
        else:
            print("Error! Unrecognized asp_triples argument.")

    set_current_project_name(project_name)


if __name__ == '__main__':
    __main__()
