#!/usr/bin/env python3
"""genark - utilities to manage UCSC GenArk assembly hubs.

This is a general, subcommand-based tool. Subcommands will be added over time.

Subcommands:
  addContrib <name>   Install a contributed track collection into the GenArk
                      assembly hubs. <name> is a subdirectory of
                      /hive/data/genomes/asmHubs/contrib/ that is laid out as
                      one directory per assembly accession (GCA_*/GCF_*) plus a
                      shared docs/ directory, e.g.:

                        contrib/<name>/GCA_000000000.0/trackDb.txt
                        contrib/<name>/GCA_000000000.0/*.bb  (and/or *.bw)
                        contrib/<name>/docs/*.html

                      For each accession it:
                        - creates <assemblyHub>/contrib/<name>/ with symlinks to
                          the collection's data files (.bb/.bw) and doc pages;
                        - writes a per-assembly <name>.trackDb.txt whose
                          bigDataUrl/html paths are rewritten to be hub-root
                          relative (contrib/<name>/...);
                        - wires that trackDb block into the assembly's useOneFile
                          hub.txt, between BEGIN/END markers (idempotent).

                      --remove uninstalls: strips the hub.txt block and removes
                      the contrib/<name>/ symlink dir.

                      NOTE: a full GenArk hub rebuild regenerates hub.txt, so
                      re-run addContrib afterwards (or add <name> to the build's
                      asmHubTrackDb.sh for a durable inclusion).

  checkContrib <name> [accessions...]
                      Run hubCheck on assembly hubs that carry the collection
                      (a random --sample N by default, or --all, or the listed
                      accessions) and report problems, separating contrib-specific
                      issues from the assemblies' own pre-existing hub warnings.
                      --noTracks does a faster structure-only check.
"""

import argparse
import os
import random
import re
import shutil
import subprocess
import sys

ASMHUBS = "/hive/data/genomes/asmHubs"
CONTRIB = os.path.join(ASMHUBS, "contrib")

ACC_RE = re.compile(r"^GC[AF]_[0-9]{9}\.[0-9]+$")


def assemblyDir(acc):
    """GCA_041900255.1 -> /hive/data/genomes/asmHubs/GCA/041/900/255/GCA_041900255.1"""
    prefix = acc[0:3]            # GCA or GCF
    digits = acc[4:]            # 041900255.1
    return os.path.join(ASMHUBS, prefix, digits[0:3], digits[3:6], digits[6:9], acc)


HUBS_URL = "https://hgdownload.soe.ucsc.edu/hubs"


def hubUrl(acc):
    """Public served hub.txt URL for a GenArk assembly accession."""
    d = acc[4:]
    return "%s/%s/%s/%s/%s/%s/hub.txt" % (
        HUBS_URL, acc[0:3], d[0:3], d[3:6], d[6:9], acc)


def symlink(target, linkPath, dryRun):
    """Create/replace an absolute symlink linkPath -> target."""
    if dryRun:
        print("  ln -sf %s %s" % (target, linkPath))
        return
    if os.path.islink(linkPath) or os.path.exists(linkPath):
        os.remove(linkPath)
    os.symlink(target, linkPath)


def rewriteTrackDb(srcPath, name):
    """Return the trackDb text with local bigDataUrl/html paths made hub-root
    relative (contrib/<name>/...). Remote (http/https/ftp) bigDataUrls and
    already-prefixed paths are left untouched."""
    out = []
    prefix = "contrib/%s/" % name
    for line in open(srcPath):
        stripped = line.lstrip()
        indent = line[:len(line) - len(stripped)]
        m = re.match(r"(bigDataUrl|linkDataUrl)\s+(\S+)\s*$", stripped)
        if m:
            key, val = m.group(1), m.group(2)
            if not re.match(r"[a-z]+://", val) and not val.startswith(prefix):
                val = prefix + os.path.basename(val)
            out.append("%s%s %s\n" % (indent, key, val))
            continue
        m = re.match(r"html\s+(\S+)\s*$", stripped)
        if m:
            val = m.group(1)
            if not val.startswith(prefix):
                val = prefix + os.path.basename(val)
            out.append("%shtml %s\n" % (indent, val))
            continue
        out.append(line)
    return "".join(out)


def wireHubTxt(hubTxt, name, block, remove, dryRun):
    """Insert/replace (or remove) the marked contrib block in the assembly's
    useOneFile hub.txt. Edits the real file the hub.txt symlink points at."""
    begin = "# BEGIN genark contrib: %s" % name
    end = "# END genark contrib: %s" % name
    blockRe = re.compile(
        r"\n*" + re.escape(begin) + r".*?" + re.escape(end) + r"\n?", re.DOTALL)
    realHub = os.path.realpath(hubTxt)
    text = open(realHub).read()
    newText = blockRe.sub("\n", text).rstrip("\n") + "\n"
    if not remove:
        newText += "\n%s\n%s\n%s\n" % (begin, block.rstrip("\n"), end)
    if dryRun:
        print("  %s %s" % ("unwire hub.txt:" if remove else "wire hub.txt:", realHub))
    else:
        with open(realHub, "w") as fh:
            fh.write(newText)


def addContrib(args):
    """Install a contrib track collection into the GenArk assembly hubs:
    symlink its data files + docs into <assemblyHub>/contrib/<name>/, write a
    per-assembly <name>.trackDb.txt with hub-root-relative paths, and wire that
    block into each assembly's hub.txt. --remove undoes all of it."""
    name = args.name.rstrip("/")
    root = os.path.join(CONTRIB, name)
    if not os.path.isdir(root):
        sys.exit("error: no such contrib collection: %s" % root)
    docsDir = os.path.join(root, "docs")
    docs = []
    if os.path.isdir(docsDir):
        docs = sorted(f for f in os.listdir(docsDir) if f.endswith(".html"))

    accs = sorted(d for d in os.listdir(root) if ACC_RE.match(d)
                  and os.path.isdir(os.path.join(root, d)))
    if not accs:
        sys.exit("error: no accession directories (GCA_*/GCF_*) under %s" % root)

    done = 0
    skipped = 0
    for acc in accs:
        accDir = os.path.join(root, acc)
        asmDir = assemblyDir(acc)
        if not os.path.isdir(asmDir):
            sys.stderr.write("skip %s: no GenArk assembly hub at %s\n" % (acc, asmDir))
            skipped += 1
            continue
        dest = os.path.join(asmDir, "contrib", name)
        hubTxt = os.path.join(asmDir, "hub.txt")

        if args.remove:
            if os.path.exists(hubTxt):
                wireHubTxt(hubTxt, name, "", remove=True, dryRun=args.dry_run)
            if args.dry_run:
                print("  rm -rf %s" % dest)
            elif os.path.isdir(dest):
                shutil.rmtree(dest)
            done += 1
            continue

        if args.dry_run:
            print("# %s -> %s" % (acc, dest))
        else:
            os.makedirs(dest, exist_ok=True)

        # symlink data files (.bb / .bw) from the collection's accession dir
        for f in sorted(os.listdir(accDir)):
            if f.endswith((".bb", ".bw")):
                symlink(os.path.join(accDir, f), os.path.join(dest, f), args.dry_run)
        # symlink shared doc pages (flat, so contrib/<name>/<doc> resolves)
        for d in docs:
            symlink(os.path.join(docsDir, d), os.path.join(dest, d), args.dry_run)

        # per-assembly trackDb with hub-root-relative paths, and wire it into hub.txt
        srcTdb = os.path.join(accDir, "trackDb.txt")
        if os.path.isfile(srcTdb):
            tdb = rewriteTrackDb(srcTdb, name)
            destTdb = os.path.join(dest, "%s.trackDb.txt" % name)
            if args.dry_run:
                print("  write %s (%d bytes)" % (destTdb, len(tdb)))
            else:
                with open(destTdb, "w") as fh:
                    fh.write(tdb)
            if os.path.exists(hubTxt):
                wireHubTxt(hubTxt, name, tdb, remove=False, dryRun=args.dry_run)
            else:
                sys.stderr.write("warn %s: no hub.txt to wire at %s\n" % (acc, hubTxt))
        done += 1

    verb = "removed from" if args.remove else "installed into"
    print("addContrib %s: %s %d assemblies, skipped %d (no assembly hub)"
          % (name, verb, done, skipped))


def contribTrackNames(root, accs):
    """Track names defined by the collection (from any one accession trackDb)."""
    for acc in accs:
        tdb = os.path.join(root, acc, "trackDb.txt")
        if os.path.isfile(tdb):
            return set(re.findall(r"^track\s+(\S+)", open(tdb).read(), re.MULTILINE))
    return set()


def checkContrib(args):
    """Run hubCheck on a set of assembly hubs that carry the collection, and
    classify any reported problems as contrib-specific vs pre-existing hub
    warnings (so the collection can be signed off without wading through the
    assemblies' own tracks)."""
    name = args.name.rstrip("/")
    root = os.path.join(CONTRIB, name)
    if not os.path.isdir(root):
        sys.exit("error: no such contrib collection: %s" % root)
    accs = sorted(d for d in os.listdir(root) if ACC_RE.match(d)
                  and os.path.isdir(os.path.join(root, d)))
    tracks = contribTrackNames(root, accs)

    if args.accession:
        sel = args.accession
    elif args.all:
        sel = accs
    else:
        n = min(args.sample, len(accs))
        sel = sorted(random.sample(accs, n))

    cmd = ["hubCheck"]
    if args.noTracks:
        cmd.append("-noTracks")

    clean = other = flagged = failed = 0
    for acc in sel:
        try:
            res = subprocess.run(cmd + [hubUrl(acc)], capture_output=True,
                                 text=True, timeout=args.timeout)
            out = res.stdout + res.stderr
        except subprocess.TimeoutExpired:
            print("%s: TIMEOUT" % acc); failed += 1; continue
        problems = [l for l in out.splitlines()
                    if l.strip() and not l.startswith("Found ")]
        contribProblems = [l for l in problems
                           if any(t in l for t in tracks) or "rror" in l]
        if not problems:
            print("%s: clean" % acc); clean += 1
        elif contribProblems:
            print("%s: CONTRIB PROBLEMS (%d)" % (acc, len(contribProblems)))
            for l in contribProblems:
                print("   ! %s" % l.strip())
            flagged += 1
        else:
            print("%s: ok (%d pre-existing hub warning(s), no contrib issue)"
                  % (acc, len(problems)))
            other += 1

    print("checkContrib %s: %d checked -- %d clean, %d ok-with-hub-warnings, "
          "%d CONTRIB PROBLEMS, %d failed"
          % (name, len(sel), clean, other, flagged, failed))
    if flagged or failed:
        sys.exit(1)


def main():
    ap = argparse.ArgumentParser(
        prog="genark", description="Manage UCSC GenArk assembly hubs.")
    ap.add_argument("--dry-run", action="store_true",
                    help="show what would be done without changing anything")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("addContrib",
                       help="install a contrib track collection into the assembly hubs")
    p.add_argument("name", help="contrib subdirectory name under %s" % CONTRIB)
    p.add_argument("--remove", action="store_true",
                   help="uninstall: remove the symlinks and the hub.txt block")
    p.set_defaults(func=addContrib)

    c = sub.add_parser("checkContrib",
                       help="run hubCheck on assembly hubs carrying the collection")
    c.add_argument("name", help="contrib subdirectory name under %s" % CONTRIB)
    c.add_argument("accession", nargs="*",
                   help="specific accessions to check (default: a random sample)")
    c.add_argument("--sample", type=int, default=5,
                   help="number of random assemblies to check (default 5)")
    c.add_argument("--all", action="store_true", help="check every assembly")
    c.add_argument("--noTracks", action="store_true",
                   help="structure only, do not fetch track data (faster)")
    c.add_argument("--timeout", type=int, default=300,
                   help="per-assembly hubCheck timeout in seconds (default 300)")
    c.set_defaults(func=checkContrib)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
