import errno
import os
import shutil
from pathlib import Path

from absl import app
from absl import flags

FLAGS = flags.FLAGS

flags.DEFINE_string("tff", default=None, help="Relative Path to tff repo folder.")
flags.mark_flag_as_required("tff")


def copydir(src, dest):
    try:
        shutil.copytree(src, dest)
    except OSError as e:
        # If the error was caused because the source wasn't a directory
        if e.errno == errno.ENOTDIR:
            shutil.copy(src, dest)
        else:
            print('Directory not copied. Error: %s' % e)


def main(argv):
    del argv  # Unused.
    dirname = os.path.dirname(__file__)

    tff_python_dir = Path(FLAGS.tff).resolve() / 'tensorflow_federated' / 'python'

    def copy_tff_subdir(subdir):
        srcdir = str(tff_python_dir / subdir)
        dstdir = os.path.join(dirname, f'tensorflow_federated_{subdir}')

        if os.path.exists(dstdir):
            shutil.rmtree(dstdir)

        copydir(srcdir, dstdir)

        for dname, dirs, files in os.walk(dstdir):
            for fname in files:
                ext = fname.split(".")[-1]
                if ext in ['BUILD', 'py', 'ipynb', 'md', 'txt', 'sh']:
                    fpath = os.path.join(dname, fname)
                    with open(fpath, encoding="utf8") as f:
                        s = f.read()
                    s = s.replace(f'tensorflow_federated.python.{subdir}', f'tensorflow_federated_{subdir}').replace(
                        f'tensorflow_federated/python/{subdir}', f'tensorflow_federated_{subdir}')
                    with open(fpath, "w", encoding='utf8') as f:
                        f.write(s)

    copy_tff_subdir('examples')
    copy_tff_subdir('research')


if __name__ == '__main__':
    app.run(main)