"""
Script to copy and update libtbx_env contents

Usage: libtbx.python update_libtbx_env.py
"""
from __future__ import absolute_import, division, print_function

import argparse
import os
import pickle
import shutil
import sys

from pprint import pprint

from libtbx.path import absolute_path

# =============================================================================
def copy_libtbx_env(dest_dir, default_dir):
  '''
  Function that copies libtbx_env from $LIBTBX_BUILD to $PREFIX
  If $LIBTBX_BUILD is not set, no copy is done. If libtbx_env does not
  exist, an IOError is raised.

  Parameters
  ----------
    default_dir: str
      The directory to copy libtbx_env

  Returns
  -------
    path or None: if the file is copied, the newly created path is returned
  '''
  value = None
  if os.getenv('LIBTBX_BUILD') is not None:
    env_name = 'libtbx_env'
    src_dir = os.getenv('LIBTBX_BUILD')
    dst_dir = os.path.join(dest_dir, os.path.relpath(default_dir, os.sep))
    src = os.path.join(src_dir, env_name)
    dst = os.path.join(dst_dir, env_name)
    print(f"Copying {env_name} from {src_dir} to {dst_dir}")
    if not os.path.isfile(src):
      raise IOError(f'The "{env_name}" file does not exist in {src}.')
    if not os.path.exists(dst_dir):
      # assumes that only the last level needs to be created.
      os.mkdir(dst_dir)
    value = shutil.copy(src, dst)
  return value

def fix_path(env, path):
  return env.as_relocatable_path(path.relocatable.replace("../../../", "lib/python3/dist-packages/"))

# =============================================================================
def update_libtbx_env(dest_dir, default_dir):
  '''
  Function that updates libtbx_env so that modules can be loaded from
  standard locations in $PREFIX

  Parameters
  ----------
    default_dir: str
    The directory to copy libtbx_env

  Returns
  -------
    None
  '''

  dst_dir = os.path.join(dest_dir, os.path.relpath(default_dir, os.sep))
  env_abs_path = os.path.join(dst_dir, 'libtbx_env')

  # set LIBTBX_BUILD and load libtbx_env from ${destdir}/${datadir}
  os.environ['LIBTBX_BUILD'] = dst_dir
  import libtbx.load_env

  sys_prefix = sys.prefix

  # basic path changes
  env = libtbx.env

  pprint(env.__dict__)

  env.build_path = absolute_path(sys.prefix)
  env.set_derived_paths()
  env.exe_path = env.bin_path
  env.pythonpath = list()
  env.python_exe = env.as_relocatable_path(sys.executable)
  env.no_bin_python = True
  site_packages_path = None
  for path in sys.path:
    if path.endswith('site-packages'):
      site_packages_path = env.as_relocatable_path(path)
      break
  relocatable_sys_prefix = env.as_relocatable_path(sys_prefix)
  env.repository_paths = [fix_path(env, p) for p in env.repository_paths]
  env.scons_dist_path = relocatable_sys_prefix

  # update the _dispatcher_registry path maybe temporary
  new = {}
  for k, v in env._dispatcher_registry.items():
    k_new = env.as_relocatable_path(k.relocatable)
    new[k_new] = fix_path(env, v)
  env._dispatcher_registry = new

  # # libtbx.python dispatcher
  # env.write_conda_dispatcher(
  #   source_file=env.python_exe,
  #   target_file='libtbx.python',
  #   source_is_python_exe=True)

  # update module_dict
  for k, v in env.module_dict.items():
    new_dist_paths = [fix_path(env, p) if p is not None else None for p in v.dist_paths]
    v.dist_paths = new_dist_paths

  # update module_dist_paths
  for k, v in env.module_dist_paths.items():
    env.module_dist_paths[k] = fix_path(env, v)

  # update path_utility
  env.path_utility = fix_path(env, env.path_utility)

  # update dispatchers
  #env.reset_dispatcher_bookkeeping()
  #env.write_python_and_show_path_duplicates()
  #for module in env.module_list:
  #  module.process_command_line_directories()

  # save the env
  env.build_path = absolute_path(default_dir)
  env.installed = True
  #env.reset_dispatcher_support()

  with open(env_abs_path, "wb") as f:
     pickle.dump(env, f, 0)
  print(f"Saved {env_abs_path} env")

  pprint(env.__dict__)

  return 0

# =============================================================================
def run():
  parser = argparse.ArgumentParser(description=__doc__,
    formatter_class=argparse.RawDescriptionHelpFormatter)

  # default location is ${PREFIX}/lib/cctbx/pythonX.Y
  default_dir = os.path.join(sys.prefix, "lib", "cctbx", f"python{sys.version_info[0]}.{sys.version_info[1]}")
  dest_dir = os.sep

  parser.add_argument('--default-dir', dest='default_dir', default=default_dir, type=str,
                      help="""The new default for the location of libtbx_env. By default,
                      the new location is ${PREFIX}/lib/cctbx/pythonX.Y. This feature is not
                      fully supported yet.""")

  parser.add_argument('--dest-dir', dest='dest_dir', default=dest_dir, type=str,
                      help="""The destdir for the location of cctbx installation. By default,
                      the new location is ${destdir}/${PREFIX}/lib/cctbx/python/X.Y. This feature is not
                      fully supported yet.""")

  namespace = parser.parse_args()

  copy_libtbx_env(dest_dir=namespace.dest_dir,
                  default_dir=namespace.default_dir)

  update_libtbx_env(dest_dir=namespace.dest_dir,
                    default_dir=namespace.default_dir)

  return 0

# =============================================================================
if __name__ == '__main__':
  sys.exit(run())
