#!/usr/bin/env python3
import argparse
import logging
import os
import sys

try:
    import imaspy as imas
except ImportError:
    import imas
from rich_argparse import RichHelpFormatter

from idstools.compute.equilibrium import EquilibriumCompute
from idstools.database import DBMaster
from idstools.utils.idshelper import get_ids_types
from idstools.utils.idslogger import setup_logger

logger = setup_logger("module", stdout_level=logging.INFO)
if __name__ == "__main__":
    # Management of input arguments
    parser = argparse.ArgumentParser(
        description="Rescaling an equilibrium magnetic field, storing the output into another entry of the same "
        "DB. replaced by ids_rescale_eq",
        formatter_class=RichHelpFormatter,
    )
    parser.add_argument(
        "-s",
        "--src",
        type=str,
        required=True,
        help="source uri (e.g. imas:hdf5?path=./testdb1 or ./testpulse1.nc)",
    )
    parser.add_argument(
        "-d",
        "--dest",
        type=str,
        required=True,
        help="destination uri (e.g. imas:hdf5?path=./testdb2 or ./testpulse2.nc)",
    )

    parser.add_argument(
        "--rescale",
        help="Rescaling factor of the equilibrium magnetic field",
        type=float,
        required=True,
    )
    parser.add_argument(
        "--dd-update",
        action="store_true",
        help=(
            "Convert IDS to the default version of the data dictionary if enabled"
            "otherwise, use the original IDS stored on disk."
        ),
    )
    args = parser.parse_args()

    rescale = args.rescale

    if rescale == 0:
        logger.critical("Rescale factor cannot be zero!")
        sys.exit(1)

    # OPEN INPUT
    inputargs = argparse.Namespace()
    inputargs.uri = args.src
    connection = DBMaster.get_connection(inputargs)
    if connection is None:
        print("Error opening source pulse! Please check existence.")
        sys.exit()

    # hack to get underlying data dictionary version used by idses
    ids = get_ids_types()
    src_dd_version = connection.dd_version
    if not args.dd_update:
        for idsname in ids:
            try:
                _dummy = connection.get(idsname, lazy=True, autoconvert=False)
                src_dd_version = _dummy.ids_properties.version_put.data_dictionary.value
                break
            except Exception as e:
                pass

    ids_equilibrium = connection.get("equilibrium", autoconvert=False)

    outputargs = argparse.Namespace()
    outputargs.mode = "w"
    outputargs.uri = args.dest

    if "mode" not in outputargs.__dict__:
        outputargs.mode = "a"
    output_connection = None
    if outputargs.uri != "" and outputargs.uri is not None:
        if args.dd_update:
            output_connection = imas.DBEntry(outputargs.uri, outputargs.mode)
        else:
            output_connection = imas.DBEntry(outputargs.uri, outputargs.mode, dd_version=src_dd_version)

    if output_connection is None:
        logger.error("Error opening destination pulse! Please check parameters and permissions.")
        exit(1)

    logger.info(f"Rescaling equilibrium magnetic field by {rescale}")
    equilibrium_compute = EquilibriumCompute(ids_equilibrium)
    equout = equilibrium_compute.rescale(rescale, dd_update=args.dd_update)
    if args.dd_update:
        dest_ids_object = imas.convert_ids(equout, output_connection.factory.version)
    # PUT IDS INTO OUTPUT
    output_connection.put(equout)

    # CLOSE FILES
    connection.close()
    output_connection.close()

    logger.info("Equilibrium IDS is rescaled successfully.")
    if args.dest is not None:
        logger.info(f"Output database details {args.dest}")
    else:
        logger.info(
            f"Output database details database={args.dest_database}, pulse={args.dest_pulse}, "
            f"run={args.dest_run}, user={os.environ['USER']}"
        )
