#!/usr/bin/env python3
"""
Shared logic for the B2G CI/CD pipeline.

Both ci/buildspec-ci.yml (pull-request validation) and ci/buildspec-cd.yml (deploy
on merge to main) call into this, so the "which services are affected" rules exist
in exactly one place. Everything is driven by ci/service-map.json.

Subcommands
  changed   --base REF [--head REF]     print affected service names, one per line
  baseline  --service NAME              print the commit sha currently deployed, or nothing
  build     --service NAME --tag SHA [--push]
  smoke     --service NAME --tag SHA    start the image and assert it is serving
  deploy    --service NAME --tag SHA    register a task def revision and roll the service

Exit codes are meaningful: any non-zero aborts the CodeBuild phase.
"""

import argparse
import json
import os
import re
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
MAP = json.load(open(os.path.join(HERE, "service-map.json")))
SERVICES = {s["name"]: s for s in MAP["services"]}
REGION = MAP["region"]

# The deployment tier. Resolved in main() from --env and held here rather than threaded
# through every function, because it is genuinely global to one invocation: a single run
# of this script only ever addresses one cluster.
#
# It is NOT read from the branch. CD passes --env explicitly, so which tier a build
# deploys to is a property of the CodeBuild project, visible in its configuration, rather
# than something inferred from git state at run time.
ENV = "prod"
CLUSTER = MAP["environments"][ENV]["cluster"]


def use_env(name):
    """Point this invocation at one tier. Called once, from main()."""
    global ENV, CLUSTER
    if name not in MAP["environments"]:
        raise SystemExit("unknown --env %r; known: %s"
                         % (name, ", ".join(sorted(MAP["environments"]))))
    ENV = name
    CLUSTER = MAP["environments"][name]["cluster"]


def family(svc):
    """Task-definition family for the current tier.

    Prod uses b2g-<name>, demo b2g-demo-<name>. They must not share: register-task-definition
    creates a new REVISION of whatever family it is given, so a demo deploy against the prod
    family would silently push a revision onto production's family - and a rollback there
    would land on an image demo built.
    """
    prefix = MAP["environments"][ENV].get("familyPrefix", "b2g-")
    return prefix + svc["name"]


def moves_latest():
    """Only the prod tier moves the :latest tag.

    Both tiers push to the same ECR repositories, so if demo also moved :latest it would
    silently repoint the tag production last pushed. Nothing in ci/ reads :latest to
    decide what to deploy - task definitions pin an explicit sha - so this only affects
    what a human sees when they list the repository.
    """
    return MAP["environments"][ENV].get("movesLatest", False)
SHA_RE = re.compile(r"^[0-9a-f]{40}$")


def log(msg):
    """Diagnostics go to stderr so a subcommand's stdout stays machine-readable."""
    print(msg, file=sys.stderr, flush=True)


def run(cmd, check=True, capture=True, cwd=ROOT):
    """
    Run a command, echoing it for the build log.

    Everything we emit here goes to STDERR on purpose: subcommands like `baseline` and
    `plan` are consumed by the caller, so their stdout must contain the result and
    nothing else. CodeBuild interleaves both streams in the log, so nothing is lost.
    """
    printable = " ".join(cmd) if isinstance(cmd, list) else cmd
    log("+ %s" % printable)
    r = subprocess.run(cmd, shell=isinstance(cmd, str), cwd=cwd,
                       capture_output=capture, text=True)
    if capture and r.stdout:
        log(r.stdout.rstrip())
    if capture and r.stderr and r.returncode != 0:
        log(r.stderr.rstrip())
    if check and r.returncode != 0:
        raise SystemExit("command failed (%d): %s" % (r.returncode, printable))
    return r


def registry():
    acct = run(["aws", "sts", "get-caller-identity", "--query", "Account",
                "--output", "text"]).stdout.strip()
    return "%s.dkr.ecr.%s.amazonaws.com" % (acct, REGION)


# --------------------------------------------------------------------------- changed

def changed_paths(base, head):
    r = run(["git", "diff", "--name-only", "%s..%s" % (base, head)])
    return [p for p in r.stdout.splitlines() if p.strip()]


def base_image_dirty(paths):
    return any(p.startswith(pref) for p in paths for pref in MAP["baseImage"]["paths"])


def affected(paths):
    """Service names made dirty by these paths, in service-map order."""
    dirty = set()
    # The 9 backends are all built FROM the shared php-nginx image, so a change there
    # invalidates every one of them even though their own directories are untouched.
    if base_image_dirty(paths):
        dirty.update(n for n, s in SERVICES.items() if s.get("needsBase"))
    for name, svc in SERVICES.items():
        if any(p.startswith(pref) for p in paths for pref in svc["paths"]):
            dirty.add(name)
    return [s["name"] for s in MAP["services"] if s["name"] in dirty]


def cmd_changed(args):
    paths = changed_paths(args.base, args.head)
    print("--- %d changed path(s) between %s and %s"
          % (len(paths), args.base[:12], args.head[:12]), file=sys.stderr)
    for p in paths[:40]:
        print("      %s" % p, file=sys.stderr)
    if len(paths) > 40:
        print("      ... and %d more" % (len(paths) - 40), file=sys.stderr)
    if base_image_dirty(paths):
        print("--- base image touched: cascading to all backends", file=sys.stderr)
    for name in affected(paths):
        print(name)


# -------------------------------------------------------------------------- baseline

def deployed_baseline(name):
    """
    The deploy baseline is what is ACTUALLY RUNNING, not HEAD~1.

    Reading it back from the live task definition makes the pipeline self-healing: a
    deploy that failed or was skipped is picked up on the next run, and several merges
    in a row behave correctly. Returns None when the running tag is not a commit sha we
    recognise, which callers treat as "must deploy".
    """
    svc = SERVICES[name]
    if not svc.get("ecsService"):
        return None
    r = run(["aws", "ecs", "describe-services", "--cluster", CLUSTER,
             "--services", svc["ecsService"], "--region", REGION,
             "--query", "services[0].taskDefinition", "--output", "text"], check=False)
    td = r.stdout.strip()
    if not td or td == "None":
        return None
    r = run(["aws", "ecs", "describe-task-definition", "--task-definition", td,
             "--region", REGION,
             "--query", "taskDefinition.containerDefinitions[0].image",
             "--output", "text"], check=False)
    image = r.stdout.strip()
    if ":" not in image:
        return None
    tag = image.rsplit(":", 1)[1]
    if not SHA_RE.match(tag):
        log("--- %s runs tag '%s' (not a commit sha) -> treat as needing deploy"
            % (name, tag))
        return None
    # Guard against a rewritten history referencing a commit we no longer have.
    if run(["git", "cat-file", "-e", tag + "^{commit}"], check=False).returncode != 0:
        log("--- %s runs sha %s which is not in this history -> redeploy"
            % (name, tag[:12]))
        return None
    return tag


def cmd_baseline(args):
    tag = deployed_baseline(args.service)
    if tag:
        print(tag)


# ------------------------------------------------------------------------------ plan

def owns_path(svc, path):
    """
    Does `path` belong to this service?

    A plain prefix test is not enough once one service's tree contains another's. The nibe
    frontend lives at nibe_datamanager/nibedatamanager_fe/, inside the nibe backend's
    nibe_datamanager/ prefix, so editing a single Angular component would otherwise mark the
    PHP backend dirty and rebuild both images on every frontend change.

    excludePaths carves those subtrees back out. Checked first, so an exclusion always wins
    over the prefix that would have claimed the file.
    """
    for ex in svc.get("excludePaths", []):
        if path.startswith(ex):
            return False
    return any(path.startswith(pref) for pref in svc["paths"])


def cmd_plan(args):
    """
    Decide what the CD run must do, as JSON:  {"base": bool, "services": [...]}

    Each service is judged against its OWN deployed sha rather than a single global
    baseline, so a service whose last deploy failed is retried while its neighbours
    are left alone.
    """
    head = args.head
    plan = {"base": False, "services": []}
    for svc in MAP["services"]:
        name = svc["name"]
        baseline = deployed_baseline(name)
        if not baseline:
            plan["services"].append({"name": name, "reason": "no commit sha deployed"})
            if svc.get("needsBase"):
                plan["base"] = True
            continue
        if baseline == head:
            continue
        paths = changed_paths(baseline, head)
        own = any(owns_path(svc, p) for p in paths)
        viabase = svc.get("needsBase") and base_image_dirty(paths)
        if own or viabase:
            plan["services"].append({
                "name": name,
                "reason": "base image changed" if (viabase and not own) else "sources changed",
                "from": baseline[:12],
            })
            if viabase:
                plan["base"] = True
    print(json.dumps(plan, indent=1))


# ----------------------------------------------------------------------------- build

def build_base(tag, push):
    b = MAP["baseImage"]
    uri = "%s/%s" % (registry(), b["ecr"])
    tags = [b["tag"], tag] + (["latest"] if moves_latest() else [])
    cmd = ["docker", "build"]
    for t in tags:
        cmd += ["-t", "%s:%s" % (uri, t)]
    run(cmd + ["-f", b["dockerfile"], b["context"]], capture=False)
    # The 9 backends inherit this image, so validate the shared nginx config once here
    # rather than discovering a typo in nine downstream builds. Also log the PHP patch
    # level and security-relevant ini values, so the build log shows what shipped.
    run(["docker", "run", "--rm", "%s:%s" % (uri, b["tag"]),
         "sh", "-c", "mkdir -p /run/nginx && nginx -t && php -v && "
                     "php -i | grep -E "
                     "'^(memory_limit|post_max_size|session.upload_progress.enabled) '"])
    if push:
        for t in tags:
            run(["docker", "push", "%s:%s" % (uri, t)], capture=False)
    return "%s:%s" % (uri, b["tag"])


def cmd_build(args):
    svc = SERVICES[args.service]
    reg = registry()
    uri = "%s/%s" % (reg, svc["ecr"])
    tags = [args.tag] + (["latest"] if moves_latest() else [])
    cmd = ["docker", "build"]
    for t in tags:
        cmd += ["-t", "%s:%s" % (uri, t)]
    cmd += ["-f", svc["dockerfile"], svc["context"]]
    if svc.get("needsBase"):
        base = "%s/%s:%s" % (reg, MAP["baseImage"]["ecr"], MAP["baseImage"]["tag"])
        cmd[2:2] = ["--build-arg", "BASE_IMAGE=" + base]
    run(cmd, capture=False)
    if args.push:
        for t in tags:
            run(["docker", "push", "%s:%s" % (uri, t)], capture=False)


# ----------------------------------------------------------------------------- smoke

def cmd_smoke(args):
    svc = SERVICES[args.service]
    uri = "%s/%s:%s" % (registry(), svc["ecr"], args.tag)
    cid = run(["docker", "run", "-d", uri]).stdout.strip()
    try:
        run(["sh", "-c", "sleep 10"])
        if svc.get("healthCmd"):
            # kong:3.9 ships neither wget nor curl, so an HTTP probe is impossible.
            run(["docker", "exec", cid, "sh", "-c", svc["healthCmd"]])
        else:
            # 127.0.0.1 deliberately: nginx binds IPv4 only, and /etc/hosts also maps
            # localhost to ::1, which busybox wget tries first and gets refused.
            port = svc["healthPort"]
            run(["docker", "exec", cid, "wget", "-q", "-O", "/dev/null",
                 "http://127.0.0.1:%d/healthz" % port])
            print("healthz OK on %d" % port)

            # /healthz is served by nginx and never touches PHP, so it stays green even
            # when the application cannot boot at all. filemst shipped for days that way:
            # every request returned a PHP fatal, and because display_errors defaulted to
            # On the fatal came back inside a 200 response. Fetch the application root and
            # fail on a fatal in the BODY, independent of status - several services
            # legitimately answer 404/400/500 on a bare root because they have no default
            # route, so the status alone proves nothing either way.
            body = run(["docker", "exec", cid, "sh", "-c",
                        "wget -q -O - http://127.0.0.1:%d/ 2>/dev/null | head -c 4000 || true"
                        % port], check=False).stdout
            for marker in ("Fatal error", "Parse error", "Uncaught "):
                if marker in body:
                    raise SystemExit(
                        "FAIL %s: application root returned a PHP %s\n%s"
                        % (svc["name"], marker.strip(), body[:600]))
            print("application root boots (no PHP fatal in body)")
        if svc["name"].startswith("fe-"):
            run(["docker", "exec", cid, "wget", "-q", "-O", "/dev/null",
                 "http://127.0.0.1:80/"])
            # SPA bundles must stay host-agnostic; a baked absolute gateway URL pins the
            # image to one hostname and silently breaks when the public URL changes.
            bad = run(["docker", "exec", cid, "sh", "-c",
                       "grep -rl 'localhost:8000\\|10\\.0\\.11\\.142' "
                       "/usr/share/nginx/html || true"], check=False).stdout.strip()
            if bad:
                raise SystemExit("FAIL absolute gateway URL baked into bundle:\n" + bad)
            print("index OK, bundle is host-agnostic")
    finally:
        run(["docker", "rm", "-f", cid], check=False)


# ---------------------------------------------------------------------------- deploy

def cmd_deploy(args):
    svc = SERVICES[args.service]
    if not svc.get("ecsService"):
        print("%s has no ECS service; image pushed only" % args.service)
        return
    uri = "%s/%s:%s" % (registry(), svc["ecr"], args.tag)

    # Take the live task definition and change only the image, so anything configured
    # out of band (env, secrets, sizing) is preserved rather than reset to a template.
    td = json.loads(run(["aws", "ecs", "describe-task-definition", "--task-definition",
                         family(svc), "--region", REGION,
                         "--output", "json"]).stdout)["taskDefinition"]
    for k in ("taskDefinitionArn", "revision", "status", "requiresAttributes",
              "compatibilities", "registeredAt", "registeredBy", "deregisteredAt"):
        td.pop(k, None)
    hit = False
    for c in td["containerDefinitions"]:
        if c["name"] == svc["container"]:
            c["image"] = uri
            hit = True
    if not hit:
        raise SystemExit("container %r not found in %s" % (svc["container"], svc["taskFamily"]))

    path = "/tmp/td-%s.json" % args.service
    json.dump(td, open(path, "w"))
    rev = run(["aws", "ecs", "register-task-definition", "--cli-input-json",
               "file://" + path, "--region", REGION,
               "--query", "taskDefinition.taskDefinitionArn",
               "--output", "text"]).stdout.strip()
    print("registered %s" % rev)

    run(["aws", "ecs", "update-service", "--cluster", CLUSTER,
         "--service", svc["ecsService"], "--task-definition", rev,
         "--region", REGION, "--query", "service.serviceName", "--output", "text"])
    if getattr(args, "no_wait", False):
        print("%s rollout started %s (not waiting)" % (svc["ecsService"], args.tag[:12]))
        return
    wait_stable([svc["ecsService"]])
    print("%s deployed %s" % (svc["ecsService"], args.tag[:12]))


def wait_stable(names):
    """Block until every named ECS service has finished rolling.

    One rollout takes ~4.5 min end to end (new task healthy + old task drained through
    the ALB deregistration delay), and the CLI waiter gives up after 40 x 15s = 10 min.
    Waiting per service serialised a 16-service deploy into 70+ min - longer than the
    build timeout - so the buildspec now starts every rollout first and calls this once.
    The waiter accepts at most 10 services per call; each batch is retried a few times
    because rollouts running concurrently can legitimately take longer than one waiter.
    """
    names = [n for n in names if n]
    for i in range(0, len(names), 10):
        batch = names[i:i + 10]
        print("waiting for %s to stabilise..." % ", ".join(batch))
        for attempt in range(1, 5):
            r = run(["aws", "ecs", "wait", "services-stable", "--cluster", CLUSTER,
                     "--services", *batch, "--region", REGION], check=False)
            if r.returncode == 0:
                break
            log("waiter attempt %d/4 did not converge, retrying" % attempt)
        else:
            raise SystemExit("services did not stabilise: %s" % ", ".join(batch))
        print("stable: %s" % ", ".join(batch))


def cmd_wait(args):
    names = [SERVICES[n]["ecsService"] for n in args.services
             if n in SERVICES and SERVICES[n].get("ecsService")]
    if not names:
        print("nothing to wait for")
        return
    wait_stable(names)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--env", default="prod",
                    help="deployment tier (default: prod). Selects the ECS cluster and "
                         "whether this run is allowed to move the :latest tag.")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("changed"); p.set_defaults(fn=cmd_changed)
    p.add_argument("--base", required=True); p.add_argument("--head", default="HEAD")

    p = sub.add_parser("baseline"); p.set_defaults(fn=cmd_baseline)
    p.add_argument("--service", required=True)

    p = sub.add_parser("build"); p.set_defaults(fn=cmd_build)
    p.add_argument("--service", required=True); p.add_argument("--tag", required=True)
    p.add_argument("--push", action="store_true")

    p = sub.add_parser("smoke"); p.set_defaults(fn=cmd_smoke)
    p.add_argument("--service", required=True); p.add_argument("--tag", required=True)

    p = sub.add_parser("deploy"); p.set_defaults(fn=cmd_deploy)
    p.add_argument("--service", required=True); p.add_argument("--tag", required=True)
    p.add_argument("--no-wait", action="store_true",
                   help="start the rollout and return; follow with `wait`")

    p = sub.add_parser("wait"); p.set_defaults(fn=cmd_wait)
    p.add_argument("--services", nargs="+", required=True,
                   help="pipeline service names (as in service-map.json) to wait for")

    p = sub.add_parser("build-base"); p.set_defaults(
        fn=lambda a: build_base(a.tag, a.push))
    p.add_argument("--tag", required=True); p.add_argument("--push", action="store_true")

    p = sub.add_parser("plan"); p.set_defaults(fn=cmd_plan)
    p.add_argument("--head", required=True)

    args = ap.parse_args()
    use_env(args.env)
    args.fn(args)


if __name__ == "__main__":
    main()