#!/usr/bin/env python3
"""Turn on a chosen set of tracks in UCSC trackDb.txt files.

For each requested track the visibility is set, and every ancestor is made
visible too: a superTrack becomes 'superTrack on show', a composite (or view)
is raised to at least the requested visibility, and a 'parent X off' is
changed to 'parent X'.

Stanzas declaring both 'superTrack on' and 'compositeTrack on' are fixed by
removing the compositeTrack line, keeping the superTrack.

Example:
  trackDbDefaultOn.py -t censatAnn,ont5mC_CpG_G638,feature,wfmashCHM13,issues \\
      -i */contrib/T2TPrimates/T2TPrimates.trackDb.txt

A track given as name=vis gets that visibility; otherwise bigWig tracks get
'full' and all others 'dense'.  With -i each file is rewritten in place and
the original is saved once as file.bak; without -i a single file (or - for
stdin) is written to stdout.  Changes and warnings are reported on stderr.
A requested track missing from a file is a warning, or with --strict an
error: that file is not written and the exit status is 1.
"""

import argparse
import os
import shutil
import sys

validVis = ("hide", "dense", "squish", "pack", "full")
visRank = {vis: rank for rank, vis in enumerate(validVis)}

# settings that only make sense on a real composite; warn if seen when the
# compositeTrack line is removed
compositeOnly = ("subGroup", "dimensions", "sortOrder", "allButtonPair",
                 "dragAndDrop", "view", "filterComposite")


def readTrackDb(path):
    """Return a list of items: a stanza is a list of lines, a blank
    separator line is a str."""
    if path == "-":
        return readLines(sys.stdin)
    with open(path) as fh:
        return readLines(fh)


def readLines(fh):
    items = []
    stanza = None
    for line in fh:
        line = line.rstrip("\n")
        if line.strip() == "":
            if stanza is not None:
                items.append(stanza)
                stanza = None
            items.append(line)
        else:
            if stanza is None:
                stanza = []
            stanza.append(line)
    if stanza is not None:
        items.append(stanza)
    return items


def writeTrackDb(items, fh):
    for item in items:
        if isinstance(item, list):
            for line in item:
                fh.write(line + "\n")
        else:
            fh.write(item + "\n")


def splitSetting(line):
    stripped = line.lstrip()
    if stripped.startswith("#"):
        return None, None
    parts = stripped.split(None, 1)
    return parts[0], (parts[1].strip() if len(parts) > 1 else "")


def getSetting(stanza, key):
    for line in stanza:
        name, value = splitSetting(line)
        if name == key:
            return value
    return None


def setSetting(stanza, key, value):
    """Replace the key's line keeping its indent, or append a new line."""
    for i, line in enumerate(stanza):
        name, _ = splitSetting(line)
        if name == key:
            indent = line[:len(line) - len(line.lstrip())]
            stanza[i] = f"{indent}{key} {value}"
            return
    trackLine = next((l for l in stanza if splitSetting(l)[0] == "track"),
                     stanza[0])
    indent = trackLine[:len(trackLine) - len(trackLine.lstrip())]
    stanza.append(f"{indent}{key} {value}")


def removeSetting(stanza, key):
    stanza[:] = [l for l in stanza if splitSetting(l)[0] != key]


def isSuperTrack(stanza):
    value = getSetting(stanza, "superTrack")
    return value is not None and value.split()[:1] == ["on"]


def isComposite(stanza):
    value = getSetting(stanza, "compositeTrack")
    return value is not None and value.split()[:1] == ["on"]


def getParentName(stanza):
    """Parent from 'parent X [on|off]', or old style 'superTrack X [vis]'."""
    value = getSetting(stanza, "parent")
    if value:
        return value.split()[0]
    value = getSetting(stanza, "superTrack")
    if value and value.split()[0] != "on":
        return value.split()[0]
    return None


def defaultVis(stanza):
    trackType = getSetting(stanza, "type") or ""
    return "full" if trackType.startswith("bigWig") else "dense"


def fixConflicts(stanzaIndex, log):
    for name, stanza in stanzaIndex.items():
        if isSuperTrack(stanza) and isComposite(stanza):
            removeSetting(stanza, "compositeTrack")
            log(f"{name}: removed 'compositeTrack on', kept superTrack")
            leftover = [k for k in compositeOnly
                        if getSetting(stanza, k) is not None]
            if leftover:
                log(f"WARNING {name}: has composite settings "
                    f"{', '.join(leftover)}; check by hand")


def turnOn(stanzaIndex, name, vis, log):
    """Set the track visible and make every ancestor visible.  Returns the
    set of ancestor names."""
    stanza = stanzaIndex[name]
    ancestors = set()
    if isSuperTrack(stanza):
        setSetting(stanza, "superTrack", "on show")
        log(f"{name}: superTrack on show")
    else:
        vis = vis or defaultVis(stanza)
        setSetting(stanza, "visibility", vis)
        log(f"{name}: visibility {vis}")
    while True:
        parentName = getParentName(stanza)
        if parentName is None:
            break
        parentValue = getSetting(stanza, "parent")
        if parentValue and "off" in parentValue.split()[1:]:
            setSetting(stanza, "parent", parentName)
            log(f"{getSetting(stanza, 'track')}: parent {parentName} (was off)")
        oldStyle = getSetting(stanza, "superTrack")
        if oldStyle and oldStyle.split()[0] == parentName and vis:
            setSetting(stanza, "superTrack", f"{parentName} {vis}")
        parent = stanzaIndex.get(parentName)
        if parent is None:
            log(f"WARNING {name}: parent {parentName} not found")
            break
        ancestors.add(parentName)
        if isSuperTrack(parent):
            if getSetting(parent, "superTrack").split() != ["on", "show"]:
                setSetting(parent, "superTrack", "on show")
                log(f"{parentName}: superTrack on show")
        elif vis:
            current = getSetting(parent, "visibility")
            if current not in visRank or visRank[current] < visRank[vis]:
                setSetting(parent, "visibility", vis)
                log(f"{parentName}: visibility {vis}")
        stanza = parent
    return ancestors


def hideOthers(stanzaIndex, keep, log):
    for name, stanza in stanzaIndex.items():
        if name in keep:
            continue
        parentName = getParentName(stanza)
        parent = stanzaIndex.get(parentName) if parentName else None
        if isSuperTrack(stanza):
            if getSetting(stanza, "superTrack").split() != ["on"]:
                setSetting(stanza, "superTrack", "on")
                log(f"{name}: superTrack on (hidden)")
        elif parent is not None and isComposite(parent):
            # only uncheck subtracks of a composite being turned on, so a
            # hidden composite still shows its subtracks when enabled later
            if parentName not in keep:
                continue
            if getSetting(stanza, "parent") not in (None, f"{parentName} off"):
                setSetting(stanza, "parent", f"{parentName} off")
                log(f"{name}: parent {parentName} off")
        elif getSetting(stanza, "visibility") != "hide":
            setSetting(stanza, "visibility", "hide")
            log(f"{name}: visibility hide")


def parseTrackList(trackArg):
    tracks = []
    for item in trackArg.split(","):
        item = item.strip()
        if not item:
            continue
        name, _, vis = item.partition("=")
        if vis and vis not in validVis:
            sys.exit(f"invalid visibility '{vis}' for {name}, "
                     f"use one of: {' '.join(validVis)}")
        tracks.append((name, vis or None))
    return tracks


def processFile(path, tracks, args, out):
    label = "stdin" if path == "-" else path

    def log(msg):
        print(f"{label}: {msg}", file=sys.stderr)

    items = readTrackDb(path)
    stanzaIndex = {}
    for item in items:
        if isinstance(item, list):
            name = getSetting(item, "track")
            if name:
                stanzaIndex[name.split()[0]] = item

    missing = [name for name, _ in tracks if name not in stanzaIndex]
    for name in missing:
        log(f"{'ERROR' if args.strict else 'WARNING'} track {name} not found")
    if missing and args.strict:
        return False

    fixConflicts(stanzaIndex, log)
    keep = set()
    for name, vis in tracks:
        if name in missing:
            continue
        keep.add(name)
        keep |= turnOn(stanzaIndex, name, vis, log)
    if args.hideOthers:
        hideOthers(stanzaIndex, keep, log)

    if args.inPlace:
        backup = path + ".bak"
        if not os.path.exists(backup):
            shutil.copy2(path, backup)
        with open(path, "w") as fh:
            writeTrackDb(items, fh)
    else:
        writeTrackDb(items, out)
    return True


def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("-t", "--tracks", required=True,
                        help="comma separated track names, each optionally "
                             "name=vis")
    parser.add_argument("-i", "--inPlace", action="store_true",
                        help="rewrite files in place, original kept as .bak")
    parser.add_argument("--hideOthers", action="store_true",
                        help="hide every track not requested (and not an "
                             "ancestor of one)")
    parser.add_argument("--strict", action="store_true",
                        help="a requested track missing from a file is an "
                             "error: that file is not written and the exit "
                             "status is 1")
    parser.add_argument("files", nargs="+",
                        help="trackDb.txt files, or - for stdin")
    args = parser.parse_args()

    if args.inPlace and "-" in args.files:
        sys.exit("-i/--inPlace cannot be used with - (stdin)")
    if not args.inPlace and len(args.files) > 1:
        sys.exit("more than one file requires -i/--inPlace")
    tracks = parseTrackList(args.tracks)
    allOk = True
    for path in args.files:
        if not processFile(path, tracks, args, sys.stdout):
            allOk = False
    if not allOk:
        sys.exit(1)


if __name__ == "__main__":
    main()