#!/opt/local/bin/python3.11
import argparse
import json
import os
import sys
import time

import pywatchman


def fieldlist(s):
    # helper for splitting a list of fields by comma
    return s.split(",")


parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    description="""
watchman-replicate-subscription can replicate an existing watchman
subscription. It queries watchman for a list of subscriptions, identifies the
source subscription (i.e., the subscription to replicate) and subscribes to watchman
using the same query.

Integrators can use this client to validate the watchman notifications their
client is receiving to localize anomalous behavior.

The source subscription is identified using any combination of the 'name',
'pid', and 'client' arguments. The provided combination must uniquely identify
a subscription. Source subscription details for a watched root can be
retrieved by running the command 'watchman-replicate-subscription PATH --list'.

By default, the replicated subscription will take the source subscription
name and prepend the substring 'replicate: ' to it. The 'qname' option can be
used to specify the replicated subscription name.

The subscription can stop after a configurable number of events are observed.
The default is a single event. You may also remove the limit and allow it to
execute continuously.

watchman-replicate-subscription will print one event per line. The event
information is determined by the fields in the identified subscription, with
each field separated by a space (or your choice of --separator).

Subscription state-enter and state-leave PDUs will be interleaved with other
events. Known subscription PDUs (currently only those generated by the
mercurial fsmonitor extension) will be enclosed in square brackets. All others will be
output in JSON format.

Events are consolidated and settled by the watchman server before they are
dispatched to watchman-replicate-subscription.

Exit Status:

The following exit status codes can be used to determine what caused
watchman-wait to exit:

0  After successfully waiting for event(s) or listing matching subscriptions
1  In case of a runtime error of some kind
2  The -t/--timeout option was used and that amount of time passed
   before an event was received
3  Execution was interrupted (Ctrl-C)

""",
)
parser.add_argument(
    "path",
    type=str,
    help="""
The path to a watched root whose subscription we'd like to replicate. The list
of watched roots can be retrieved by running 'watchman watch-list'.
""",
)
parser.add_argument(
    "-s",
    "--separator",
    type=str,
    default=" ",
    help="String to use as field separator for event output.",
)
parser.add_argument(
    "-q",
    "--qname",
    type=str,
    default=None,
    help="""
The replicated subscription name. The default will be the source subscription
with the string 'replicate: ' prepended to it.
""",
)
parser.add_argument(
    "-n",
    "--name",
    type=str,
    default=None,
    help="""
The name of the subscription to replicate.
""",
)
parser.add_argument(
    "-c",
    "--client",
    type=str,
    default=None,
    help="""
The client id of the subscription to replicate.
""",
)
parser.add_argument(
    "-p",
    "--pid",
    type=str,
    default=None,
    help="""
The process id of the subscription to replicate.
""",
)
parser.add_argument(
    "-m",
    "--max-events",
    type=int,
    default=1,
    help="""
Set the maximum number of events that will be processed.  When the limit
is reached, watchman-replicate-subscription exit.  The default is 1.  Setting
the limit to 0 removes the limit, causing watchman-wait to execute indefinitely.
""",
)
parser.add_argument(
    "-t",
    "--timeout",
    type=float,
    default=0,
    help="""
Exit if no events trigger within the specified timeout.  If timeout is
zero (the default) then keep running indefinitely.
""",
)
parser.add_argument(
    "--connect-timeout",
    type=float,
    default=100,
    help="""
Initial watchman client connection timeout. It should be sufficiently large to
prevent timeouts when watchman is busy (eg. performing a crawl). The default
value is 100 seconds.
""",
)
parser.add_argument(
    "-l",
    "--list",
    action="store_true",
    help="""
Print the matching subscription list and exit.
""",
)
parser.add_argument(
    "-f",
    "--full",
    action="store_true",
    help="""
Use with '--list' to print complete subscription information, including the
query.
""",
)
parser.add_argument(
    "--state-only",
    action="store_true",
    help="""
Print only the subscription state-enter and state-leave PDUs.
""",
)
args = parser.parse_args()

# Running total of individual file events we've seen
total_events = 0


class Subscription(object):
    root = None  # Watched root
    name = None  # Our name for this subscription
    path = None

    def __init__(self, path, name, query):
        self.name = name
        self.query = query
        self.path = os.path.abspath(path)
        if not os.path.exists(self.path):
            print("path %s (%s) does not exist." % (path, self.path), file=sys.stderr)
            sys.exit(1)

    def __repr__(self):
        return "Subscription(path=%s, name=%s, query=%s)" % (
            self.path,
            self.name,
            json.dumps(self.query),
        )

    def start(self, client):
        watch = client.query("watch-project", self.path)
        if "warning" in watch:
            print("WARNING: ", watch["warning"], file=sys.stderr)
        self.root = watch["watch"]

        # get the initial clock value so that we only get updates
        self.query["since"] = client.query("clock", self.root)["clock"]
        sub = client.query("subscribe", self.root, self.name, self.query)

    def formatField(self, fname, val):
        return str(val)

    def formatFsmonitorSubPdu(self, sub):
        out = []
        keys = ["state-enter", "state-leave"]
        for key in keys:
            if key in sub:
                out.append(key)
                out.append(sub[key])
        keys = ["abandoned"]
        for key in keys:
            if key in sub:
                out.append(key)
        if "metadata" in sub:
            metadata = sub["metadata"]
            keys = ["status", "rev", "partial", "distance"]
            for key in keys:
                if key in keys:
                    out.append(str(metadata[key]))
        return "[ %s ]" % (" ".join(out))

    def formatSubPdu(self, sub):
        # If fsmonitor metadata fields present, pretty-print
        if "metadata" in sub and all(
            key in sub["metadata"] for key in ["status", "rev", "partial", "distance"]
        ):
            return self.formatFsmonitorSubPdu(sub)
        return json.dumps(sub, indent=2)

    def emit(self, client):
        global total_events
        data = client.getSubscription(self.name)
        if data is None:
            return False
        for dat in data:
            if any(key in dat.keys() for key in ["state-enter", "state-leave"]):
                print(self.formatSubPdu(dat))
                sys.stdout.flush()
                total_events = total_events + 1
                if args.max_events > 0 and total_events >= args.max_events:
                    sys.exit(0)
            if args.state_only:
                continue
            for f in dat.get("files", []):
                out = []
                if len(repFields) == 1:
                    # When only 1 field is specified, the result is a
                    # list of just the values
                    out.append(self.formatField(repFields[0], f))
                else:
                    # Otherwise it is a list of objects
                    for fname, val in f.items():
                        out.append(self.formatField(fname, val))
                print(args.separator.join(out))
                sys.stdout.flush()
                total_events = total_events + 1
                if args.max_events > 0 and total_events >= args.max_events:
                    sys.exit(0)
        return True


def getSubIdInfo(subName, subClient, subPid):
    return "(name='%s', pid='%s' client='%s')" % (
        "Any" if subName is None else subName,
        "Any" if subPid is None else subPid,
        "Any" if subClient is None else subClient,
    )


def getSubInfo(sub, keys=None):
    rslt = {}
    if keys is None:
        keys = ["name", "pid", "client", "query"]
    subInfo = sub["info"] if "info" in sub else {}
    for key in keys:
        if key in subInfo:
            rslt[key] = subInfo[key]
        else:
            rslt[key] = ""
            keyWarn = true
    return rslt


path = args.path
subName = args.name
subPid = args.pid
subClient = args.client
client = pywatchman.client(timeout=args.connect_timeout)

# We use debug-get-subscriptions to get subscription information.  Note:
# the debug commands are not stable/supported, so we are susceptible to breakage
# if they change.  Do not copy this approach without understanding this risk.
subs = client.query("debug-get-subscriptions", path)
matchSubs = []
for sub in subs["subscribers"]:
    info = sub["info"]
    if (
        (subName is None or sub["info"]["name"] == subName)
        and (subPid is None or str(sub["info"]["pid"]) == subPid)
        and (subClient is None or str(sub["info"]["client"]) == subClient)
    ):
        matchSubs.append(sub)

if args.list:
    if len(matchSubs) == 0:
        print("No matching subscriptions")
    for sub in matchSubs:
        keys = None if args.full else ["name", "pid", "client"]
        print(json.dumps(getSubInfo(sub, keys=keys), indent=2), file=sys.stdout)
    sys.exit(0)

if len(matchSubs) == 0:
    print(
        "Error, no matching subscriptions:\n"
        "\tcriteria: %s\n"
        "\tpath: %s\n"
        "To get a list of subscriptions for a watched root, use:\n"
        "\twatchman-replicate-subscription PATH --list"
        % (getSubIdInfo(subName, subClient, subPid), path),
        file=sys.stderr,
    )
    sys.exit(1)

if len(matchSubs) > 1:
    print(
        "Error, found multiple matching subscriptions:\n"
        "\tcriteria: %s\n"
        "\tpath: %s\n"
        "Use the '--name', '--client' and '--pid' options to identify a subscription.\n"
        "To get a list of subscriptions for a watched root, use:\n"
        "\twatchman-replicate-subscription PATH --list"
        % (getSubIdInfo(subName, subClient, subPid), path),
        file=sys.stderr,
    )
    sys.exit(1)

matchSub = matchSubs[0]
repQuery = matchSubs[0]["info"]["query"]

# Subscriptions with no fields specified, will get watchman default
# and will describe themselves in the response.  If they specified
# just a single field then we'd like to know what it is because we
# won't know the name from the `files` list so we make an attempt
# to capture the field definition here and use it only in the case
# that it is a single field name,
repFields = matchSubs[0]["info"]["query"].get("fields", [])
repName = args.qname if args.qname else "replicate: " + matchSub["info"]["name"]

repSub = Subscription(name=repName, path=path, query=repQuery)
print("%s" % (repSub))

deadline = None
if args.timeout > 0:
    deadline = time.time() + args.timeout

try:
    repSub.start(client)

except pywatchman.CommandError as ex:
    print("watchman:", ex.msg, file=sys.stderr)
    sys.exit(1)

while deadline is None or time.time() < deadline:
    try:
        if deadline is not None:
            client.setTimeout(deadline - time.time())
        # wait for a unilateral response
        result = client.receive()

        # in theory we can parse just the result variable here, but
        # the client object will accumulate all subscription results
        # over time, so we ask it to remove and return those values
        # for each of the subscriptions
        repSub.emit(client)

    except pywatchman.SocketTimeout as ex:
        if deadline is not None and time.time() >= deadline:
            sys.exit(2)

        # Let's check to see if we're still functional
        try:
            vers = client.query("version")
        except Exception as ex:
            print("watchman:", str(ex), file=sys.stderr)
            sys.exit(1)

    except KeyboardInterrupt:
        # suppress ugly stack trace when they Ctrl-C
        sys.exit(3)
