#!/usr/bin/env python3
# Converts an IMAS database (all its data-entries) to the specified folder and backend

import argparse
import logging
import os
from datetime import datetime
from glob import glob

try:
    import imaspy as imas
except ImportError:
    import imas
import pandas as pd
from rich.console import Console
from rich_argparse import RichHelpFormatter

from idstools.database import DBMaster
from idstools.utils.clihelper import get_backend_id, imas_parser
from idstools.utils.idshelper import (
    get_available_ids_and_occurrences,
    idsdiff,
)
from idstools.utils.idslogger import setup_logger

logger = setup_logger("module", stdout_level=logging.INFO)
console = Console()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Copy all data-entries from one database into another one",
        parents=[imas_parser],
        formatter_class=RichHelpFormatter,
    )
    parser.add_argument(
        "-do",
        "--database_out",
        type=str,
        required=True,
        help="Name of destination database",
    )
    parser.add_argument(
        "-bo",
        "--backend_output",
        type=str,
        required=True,
        help="Desired backend for destination data-entry",
    )
    parser.add_argument(
        "--skip-obsolete",
        action="store_true",
        help="Do not copy data that have been marked obsolete (ITER scenarios only)",
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        help="Performs diff and validation of the migrated data",
    )
    args = parser.parse_args()
    locpath = DBMaster.get_db_Path(args.user, args.database, args.version)

    backend = get_backend_id(args.backend)

    if backend == imas.ids_defs.MDSPLUS_BACKEND:
        files = DBMaster.mds_list_pulse_run(locpath, with_status="active" if args.skip_obsolete else None)
    elif backend == imas.ids_defs.HDF5_BACKEND:
        files = DBMaster.hdf5_list_pulse_run(locpath)

    schema = []
    current_fpath = os.path.dirname(os.path.realpath(__file__))
    schema_dir = os.path.join(current_fpath, "../../../../bin/validation_schemas")
    if not os.path.isdir(schema_dir):
        current_fpath = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
        schema_dir = os.path.join(current_fpath, "idstools/validation_schemas")

    if os.path.isdir(schema_dir):
        schema += sorted(glob(schema_dir + "/generic/*.y*ml", recursive=True))

    schema_dict = {}

    for file_path in schema:
        filename = os.path.splitext(os.path.basename(file_path))[0]
        schema_dict[filename] = file_path
    log_messages = []

    for pulse in files:
        error = False
        console.print(f"\nProcessing {pulse}")

        idsinf = []
        run = pulse[1]
        try:
            src = imas.DBEntry(imas.ids_defs.MDSPLUS_BACKEND, args.database, pulse[0], run, args.user)
            src.open()
        except Exception as e:
            console.print(f"Error opening pulse!error={e} Please check existence.")
            log_messages.append("Error opening pulse")
            continue

        dest = imas.DBEntry(get_backend_id(args.backend_output), args.database_out, pulse[0], run)
        try:
            dest.create()
        except Exception as e:
            console.print(
                f"Error creating pulse! error={e} backend={args.backend_output} "
                "database={args.database_out} shot={pulse[0]} run={run}"
            )
            error = True
            log_messages.append("Error creating pulse")
            continue
        avids = get_available_ids_and_occurrences(src)
        complete_description = ""
        identical = True
        for ids in avids:
            idsname = ids[0]
            inocc = ids[1]
            idsobj = src.get(idsname, occurrence=inocc)
            try:
                dest.put(idsobj, occurrence=inocc)
            except Exception as e:
                console.print(f"Error while storing ids !error={e}")
                error = True
                continue
            if args.validate:
                idsobj2 = dest.get(idsname, occurrence=inocc)
                compare_output, _, compare_output_description = idsdiff(idsobj, idsobj2, "", "", print_result=False)
                if not compare_output:
                    identical = False
                    print(compare_output_description)
        if args.validate:
            if not identical:
                log_messages.append(complete_description)
            else:
                log_messages.append("identical")

        # #TODO IDS validator is not yet updated to imas-python
        # if idsname in schema_dict.keys():
        #     idsinf.append(
        #         (
        #             idsname,
        #             same,
        #             ids_validator(idsobj, schema_dict[idsname])[0],
        #         )
        #     )
        # else:
        #     idsinf.append(
        #         (
        #             idsname,
        #             same,
        #             "No Schema",
        #         )
        #     )
        #     console.print(f"could not find schema for {idsname}")

        src.close()
        dest.close()
        if error:
            console.print(
                f"Could not convert : backend={args.backend_output} database={args.database_out} "
                f"pulse={pulse[0]} run={run}"
            )
        else:
            console.print(
                f"successfully converted, backend={args.backend_output} database={args.database_out} "
                f"pulse={pulse[0]} run={run}"
            )

    if args.validate:
        df = pd.DataFrame(files, columns=["PULSE", "RUN"])
        df["Validation"] = log_messages
        date = datetime.now().strftime("%Y_%m_%d-%I:%M:%S_%p")
        df.to_csv(
            args.database_out + "_migration_log--" + date + ".csv",
            na_rep="None",
            index=False,
            header=True,
        )
        print(args.database_out + "_migration_log--" + date + ".csv")
        print(df)
