#!python
import os
import re
import sys
import shutil
import subprocess
import hifiyaml as hy

# Use local package when running from the repo
here = os.path.dirname(os.path.abspath(__file__))
yamltools_path = os.path.join(here, "..")
if os.path.exists(os.path.join(yamltools_path, "yamltools4jedi", "__init__.py")):
    if yamltools_path not in sys.path:
        sys.path.insert(0, yamltools_path)

import yamltools4jedi as yj  # noqa: E402

args = sys.argv
nargs = len(args) - 1
if nargs < 2:
    print(f"Usage: {os.path.basename(args[0])} <operator> <file|dir> [dir|file|obs_str|querystr] [file|newblock(string_or_file)]")
    print("where 'operator' can be: split1, split2, pack, compare1, compare2, listobs, removeobs, keepobs, dump, drop, modify")
    print("check https://github.com/hifiyaml/yamltools4jedi/wiki/yj for more details")
    exit()

dedent = os.getenv("YJ_DEDENT", "False").upper() == "TRUE"
split_level = int(os.getenv("YJ_SPLIT_LEVEL", "1"))
listIndent = os.getenv("YJ_LIST_INDENT", "False").upper() == "TRUE"
plain_pack = os.getenv("YJ_PLAIN_PACK", "True").upper() == "TRUE"

newblock_file, querystr, fpack, yobsout = "", "", "", ""

# ------------------------------------------------------------------------------------------------------------------------------------
#  process the command line options
# ------------------------------------------------------------------------------------------------------------------------------------
operator = args[1]
yfile = args[2]
basename = os.path.basename(yfile)
if operator == "split1":
    dirname = f'split1.{basename}'   # default value for the split operators
elif operator == "split2":
    dirname = f'split2.{basename}'   # default value for the split operators
elif operator == "split":
    dirname = f'split{split_level}.{basename}'   # default value for the split operators

if nargs > 2:
    querystr = args[3]
    dirname = args[3] # split operators
    fpack = args[3]  # pack operator
    yfile2 = args[3] # compare operators
    obs_str = args[3] # listobs, removeobs, keepobs
if nargs > 3:
    newblock_file = args[4]  # modify operator
    yobsout = args[4]  # listobs, removeobs, keepobs, optional

if operator == "modify":
    if not newblock_file:
        print("newblock_file cannot be empty for the modify operator")
    if not querystr:
        print("querystr cannot be empty for the modify operator")
elif operator == "pack":  # pack operator
    dirname = yfile

# ------------------------------------------------------------------------------------------------------------------------------------
#  load the yaml file for all non-pack operations; pack will read a directory separately
# ------------------------------------------------------------------------------------------------------------------------------------
if operator != "pack":
    data = hy.load(yfile)

# ------------------------------------------------------------------------------------------------------------------------------------
#  if specified, convert JEDI specific querystr (#obs:filter#/path/to/a/block) to be generic
# ------------------------------------------------------------------------------------------------------------------------------------
if "#" in querystr:
    querystr = querystr.strip("#").strip()
    parts = querystr.split("#")
    part0 = parts[0]
    part1 = ""
    if len(parts) > 1:
        part1 = parts[1]
    myfilter = ""
    if ":" in part0:
        myobs, myfilter = part0.strip().split(":")
    else:
        myobs = part0.strip()
    if myfilter:  # need to read all obs/filters
        dcObs = yj.get_all_obs(data, shallow=False)
    else:
        dcObs = yj.get_all_obs(data, shallow=True)
    # ~~~~
    obsIndex = -1
    filterstr = ""
    for knt, (name, observer) in enumerate(dcObs.items()):
        sname = observer["sname"]
        if sname in myobs:
            obsIndex = knt
            if myfilter:  # find the specified filter id
                def get_filterstr(key):
                    if key in observer:
                        for knt, dcFilter in enumerate(observer[key]):
                            if dcFilter["id"] == myfilter:
                                filterstr = f"/obs {key}/{knt}"
                                return filterstr
                if not filterstr:
                    filterstr = get_filterstr("filters")
                if not filterstr:
                    filterstr = get_filterstr("prior filters")
                if not filterstr:
                    filterstr = get_filterstr("pre filters")
                if not filterstr:
                    filterstr = get_filterstr("post filters")
                if not filterstr:
                    sys.stderr.write(f'filterID "{myfilter}" not found!\n')
                    sys.exit(1)
            # ~~~~~~~~~~
            break
    # ~~~~~~~~
    if obsIndex >= 0:
        querystr = f"observations/observers/{obsIndex}{filterstr}/" +  part1.strip().strip("/")
    else:
        sys.stderr.write(f'observer "{myobs}" not found!\n')
        sys.exit(1)

# ------------------------------------------------------------------------------------------------------------------------------------
#  dump, drop, modify, split, split1, split2, pack
# ------------------------------------------------------------------------------------------------------------------------------------
if operator == "dump":
    hy.dump(data, querystr, do_dedent=dedent)

elif operator == "drop":
    hy.drop(data, querystr)
    hy.dump(data)

elif operator == "modify":
    if os.path.exists(newblock_file):
        newblock = hy.load(newblock_file)
    else:
        newblock = hy.text_to_yblock(newblock_file)
    hy.modify(data, querystr, newblock)
    hy.dump(data)

elif operator == "split":
    yj.split(yfile, level=split_level, dirname=dirname, do_dedent=dedent)
    print(f"Done, split files are available under {dirname}/")

elif operator == "split1":
    yj.split(yfile, level=1, dirname=dirname,  do_dedent=dedent)
    print(f"Done, split1 files are available under {dirname}/")

elif operator == "split2":
    yj.split(yfile, level=2, dirname=dirname, do_dedent=dedent)
    print(f"Done, split2 files are available under {dirname}/")

elif operator == "pack":
    if not fpack:
        if "split1." in dirname:
            fpack = "pack." + dirname.rstrip('/').replace("split1.", "")
        else:
            fpack = "pack." + dirname.rstrip('/').replace("split2.", "")
    yj.pack(dirname, fpack, listIndent=listIndent, plain_pack=plain_pack)
    print(f"Done, the final packed file: {fpack}")

# ------------------------------------------------------------------------------------------------------------------------------------
# compare1 and compare2
# ------------------------------------------------------------------------------------------------------------------------------------
elif operator == "compare1":  # compare by observers
    tmpdir = "tmp_yamltoos4jedi_compare"
    shutil.rmtree(tmpdir, ignore_errors=True)
    os.makedirs(tmpdir, exist_ok=True)
    yj.split(yfile, level=1, dirname=f"{tmpdir}/{yfile}", do_dedent=False)
    yj.split(yfile2, level=1, dirname=f"{tmpdir}/{yfile2}", do_dedent=False)
    with open(f'{tmpdir}/compare.sh', 'w') as outfile:
        outfile.write(f'''
cd {tmpdir}
cat <<EOF
A series of vimdiff commands will run to show the differences between common observers
and the observers existing in one file but not in another

In the vimdiff windows,
Use ":qa" to quit the current vimdiff and move to the next
Use ":cq" to quit the current vimdiff and end this compare1 session
press Enter to continue ...
EOF
read whatever
echo "searching for the next batch of differences ..."
for myfile in {yfile}/*; do
  filename=$(basename ${{myfile}})
  if [[ -s {yfile2}/${{filename}} ]] && ! diff ${{myfile}} {yfile2}/${{filename}} 2>1 1>/dev/null; then
    vimdiff ${{myfile}} {yfile2}/${{filename}} || exit $?
    echo "searching for the next batch of differences ..."
  fi
done
echo "Done!"
        ''')
    subprocess.run(["bash", f"{tmpdir}/compare.sh"])
    shutil.rmtree(tmpdir, ignore_errors=True)

elif operator == "compare2":  # compare by observers/filters
    tmpdir = "tmp_yamltoos4jedi_compare"
    shutil.rmtree(tmpdir, ignore_errors=True)
    os.makedirs(tmpdir, exist_ok=True)
    yj.split(yfile, level=2, dirname=f"{tmpdir}/{yfile}", do_dedent=False)
    yj.split(yfile2, level=2, dirname=f"{tmpdir}/{yfile2}", do_dedent=False)
    with open(f'{tmpdir}/compare.sh', 'w') as outfile:
        outfile.write(f'''
cd {tmpdir}
cat <<EOF
A series of vimdiff commands will run to show the differences between common observers/filters
and the observers/filters existing in one file but not in another

In the vimdiff windows,
Use ":qa" to quit the current vimdiff and move to the next
Use ":cq" to quit the current vimdiff and end this compare2 session
press Enter to continue ...
EOF
read whatever
echo "searching for the next batch of differences ..."
for mydir in {yfile}/*; do
  filename=$(basename ${{mydir}})
  if [[ ! -d ${{mydir}} ]]; then
    if [[ -s {yfile2}/${{filename}} ]] && ! diff ${{mydir}} {yfile2}/${{filename}} 2>1 1>/dev/null; then
      vimdiff ${{mydir}} {yfile2}/${{filename}} || exit $?
      echo "searching for the next batch of differences ..."
    fi
  else
    for myfile in ${{mydir}}/*; do
      filename2=$(basename ${{myfile}})
      if [[ -s {yfile2}/${{filename}}/${{filename2}} ]] && ! diff ${{myfile}} {yfile2}/${{filename}}/${{filename2}} 2>1 1>/dev/null; then
        vimdiff ${{myfile}} {yfile2}/${{filename}}/${{filename2}} || exit $?
        echo "searching for the next batch of differences ..."
      fi
    done
  fi
done
echo "Done!"
        ''')
    subprocess.run(["bash", f"{tmpdir}/compare.sh"])
    shutil.rmtree(tmpdir, ignore_errors=True)

# ------------------------------------------------------------------------------------------------------------------------------------
# listobs, removeobs, keepobs
# ------------------------------------------------------------------------------------------------------------------------------------
elif operator in ("listobs", "removeobs", "keepobs"):  # listobs, removeobs, keepobs
    if operator == "listobs":
        shallow = False
    else:
        shallow = True
    dcObs = yj.get_all_obs(data, shallow=shallow)
    if operator == "listobs":
        obs_text = ''
        filter_text = ''
        for name, observer in dcObs.items():
            obs_text = obs_text + observer['sname'] + ', '
            filter_knt = len(observer['pre filters']) + len(observer['filters']) + len(observer['prior filters']) + len(observer['post filters'])
            filter_text += f'{name.ljust(35)}: {filter_knt} filters\n'
        nobs = len(dcObs)
        obs_text = f'{nobs} observers:\n{obs_text.rstrip(", ")}'
        print(f'{filter_text}\n\n{obs_text}\n')
    else:
        pos1, _ = hy.get_start_pos(data, "observations/observers")  # get the YAML head section
        pos2 = hy.next_pos(data, pos1)
        output = data[0:pos1 + 1]  # header part, including the line "observers:"
        # assemble observers
        obs_list = [myobs.strip() for myobs in re.split(r'[,\s]', obs_str)]
        for _, observer in dcObs.items():
            sname = observer["sname"]
            #  determine whether to keep an observer based on removeobs or keepobs operators and the obs_list
            if operator == "removeobs":
                keep = True
                for myobs in obs_list:
                    if myobs == sname:
                        keep = False
                        break
            else:  # keepobs
                keep = False
                for myobs in obs_list:
                    if myobs == sname:
                        keep = True
                        break
            # ~~~~~~~~~~~~
            if keep:
                output.extend(data[observer["pos1"]:observer["pos2"]])
        # ~~~~~~~~
        # concatenate the trailing part if any
        output.extend(data[pos2:])
        if not yobsout:  # output file not specified
            hy.dump(output)
        else:
            hy.dump(output, fpath=yobsout)

else:
    hy.printd(f"unknown operator: {operator}")
