#!/usr/bin/python3

import os
import pickle

from libtbx.env_config import unpickle
from libtbx.path import relocatable_path, absolute_path

PREFIX = "/usr"
DIST_PACKAGE = "lib/python3/dist-packages"

def to_rel(s):
    if s.startswith('/'):
        s = s[1:]
    return s

def fix_path(p: relocatable_path,
             prefix: str,
             dist_package: str) -> relocatable_path:
    if p is None:
        return None

    ps = p.relocatable.split('modules')
    if len(ps) == 2:
        anchor = absolute_path(prefix)
        relocatable = os.path.join(dist_package,
                                   to_rel(ps[1].split('cctbx_project')[-1]))
        return relocatable_path(anchor, relocatable)

def dispatcher_registry(reg, prefix, dist_package):
    res = {}
    for k, v in reg.items():
        path = v
        vs = v.relocatable.split('modules/')
        if len(vs) == 2:
            anchor = absolute_path(prefix)
            relocatable = os.path.join(dist_package,
                                       vs[1].split('cctbx_project/')[-1])
            path = relocatable_path(anchor, relocatable)
        res[k] = path

    return res


def module_dict(d, prefix, dist_package):
    res = d.copy()
    for k, v in res.items():
        v.dist_paths = [fix_path(p, prefix, dist_package) for p in v.dist_paths]
    return res


def module_dist_paths(d, prefix, dist_package):
    res = d.copy()
    for k, v in res.items():
        res[k] = fix_path(v, prefix, dist_package)
    return res


def main():
    env = unpickle()
    env.build_path.reset(PREFIX)

    env._dispatcher_registry = dispatcher_registry(env._dispatcher_registry,
                                                   PREFIX,
                                                   DIST_PACKAGE)
    env.module_dict = module_dict(env.module_dict,
                                  PREFIX,
                                  DIST_PACKAGE)

    env.module_dist_paths = module_dist_paths(env.module_dist_paths,
                                              PREFIX,
                                              DIST_PACKAGE)

    env.path_utility = fix_path(env.path_utility, PREFIX, DIST_PACKAGE)

    env.pythonpath = []

    env.repository_paths = [fix_path(p, PREFIX, DIST_PACKAGE)
                            for p in env.repository_paths]

    with open("libtbx_env", "wb") as f:
        pickle.dump(env, f, 0)

    print("done!")

main()
