#!/usr/bin/env python3
"""Turn a Railway environment's hand-over into a Docker Compose project.

  railway/import plan          what the hand-over should copy (used by handover)
  railway/import setup         write ~/app/compose.yaml, ~/app/services.json,
                               ~/app/crontab.jobs and ~/railway/plan.md
  railway/import list           name, compose service, kind, and volume of each service
  railway/import set SVC KEY [VALUE]
                               set a variable in compose.yaml (VALUE from stdin
                               if left out), then recreate the service

From https://exe.dev/deploy/migrate-railway. Reads ~/railway: status.json,
environment.json and variables/<service>.json, as handover wrote them.
"""
import json, os, re, secrets, shlex, socket, string, subprocess, sys

HOME = os.path.expanduser("~")
RW = os.path.join(HOME, "railway")
APP = os.path.join(HOME, "app")
COMPOSE = os.path.join(APP, "compose.yaml")

DB_KINDS = [  # (kind, image pattern)
    ("postgres", re.compile(r"postgres|postgis|pgvector|timescale", re.I)),
    ("redis", re.compile(r"(^|/)(redis|valkey)", re.I)),
    ("mysql", re.compile(r"mysql|mariadb", re.I)),
    ("mongo", re.compile(r"mongo", re.I)),
]
PASSWORD_VARS = {
    "postgres": ["POSTGRES_PASSWORD", "PGPASSWORD"],
    "redis": ["REDIS_PASSWORD", "REDISPASSWORD"],
    "mysql": ["MYSQL_ROOT_PASSWORD", "MYSQLPASSWORD", "MYSQL_PASSWORD"],
    "mongo": ["MONGO_INITDB_ROOT_PASSWORD", "MONGOPASSWORD"],
}
# Variable names that usually hold a key for someone else's API: candidates
# for an exe.dev integration. (name pattern, catalog handle or None)
KEY_HINTS = [
    (r"^OPENAI_API_KEY$", "openai"), (r"^ANTHROPIC_API_KEY$", "anthropic"),
    (r"^STRIPE_(SECRET|API|RESTRICTED)_KEY$|^STRIPE_KEY$", "stripe"),
    (r"^RESEND_API_KEY$", "resend"), (r"^SENDGRID_API_KEY$", "sendgrid"),
    (r"^POSTMARK_(SERVER_)?(API_)?TOKEN$", "postmark"), (r"^MAILGUN_API_KEY$", "mailgun"),
    (r"^TWILIO_AUTH_TOKEN$", "twilio"), (r"^GEMINI_API_KEY$|^GOOGLE_API_KEY$", "gemini"),
    (r"^GROQ_API_KEY$", "groq"), (r"^MISTRAL_API_KEY$", "mistral"),
    (r"^REPLICATE_API_TOKEN$", "replicate"), (r"^ELEVENLABS_API_KEY$", "elevenlabs"),
    (r"^GITHUB_TOKEN$", "github"), (r"^SLACK_BOT_TOKEN$", "slack-bot"),
    (r"_API_KEY$|_API_TOKEN$|_ACCESS_TOKEN$|_SECRET_KEY$", None),
]
APP_SECRET_HINT = re.compile(r"WEBHOOK|SIGNING|SESSION|COOKIE|JWT|ENCRYPT|AUTH_SECRET|SECRET_KEY_BASE|APP_KEY", re.I)


def load(path, default=None):
    try:
        with open(path) as f:
            return json.load(f)
    except FileNotFoundError:
        if default is not None:
            return default
        sys.exit(f"import: {path} is missing; run the hand-over first")


def slug(name):
    s = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
    return s or "service"


def vm_name():
    return socket.gethostname().split(".")[0]


def services():
    """Railway services of the handed-over environment, in a stable order."""
    status = load(os.path.join(RW, "status.json"))
    env = load(os.path.join(RW, "environment.json"))
    names = {e["node"]["id"]: e["node"]["name"] for e in status.get("services", {}).get("edges", [])}
    out, seen = [], set()
    for sid, inst in sorted(env.get("services", {}).items(), key=lambda kv: names.get(kv[0], kv[0]).lower()):
        if inst.get("isDeleted"):
            continue
        name = names.get(sid, sid)
        s = {"id": sid, "name": name, "slug": slug(name), "inst": inst}
        while s["slug"] in seen:
            s["slug"] += "-2"
        seen.add(s["slug"])
        src = inst.get("source") or {}
        s["image"] = src.get("image")
        s["repo"] = src.get("repo")
        s["branch"] = src.get("branch") or "main"
        s["root"] = (src.get("rootDirectory") or "").strip("/")
        s["kind"] = "code"
        if s["image"]:
            s["kind"] = "image"
            for kind, pat in DB_KINDS:
                if pat.search(s["image"]):
                    s["kind"] = kind
                    break
        mounts = [m.get("mountPath") for m in (inst.get("volumeMounts") or {}).values() if m and m.get("mountPath")]
        s["mount"] = mounts[0] if mounts else None
        s["deploy"] = inst.get("deploy") or {}
        s["build"] = inst.get("build") or {}
        net = inst.get("networking") or {}
        s["service_domains"] = net.get("serviceDomains") or {}
        s["custom_domains"] = net.get("customDomains") or {}
        s["tcp"] = bool(net.get("tcpProxies"))
        s["private"] = net.get("privateNetworkEndpoint")
        s["cron"] = (s["deploy"].get("cronSchedule") or "").strip() or None
        if s["kind"] == "code":
            s["git"] = slug(s["repo"].split("/")[-1]) if s["repo"] else s["slug"]
        out.append(s)
    return out, env


def cmd_plan():
    """Lines for handover: what to copy, tab-separated."""
    svcs, _ = services()
    for s in svcs:
        print(f"vars\t{s['name']}")
    for s in svcs:
        if s["kind"] in ("postgres", "redis"):
            print(f"{s['kind']}\t{s['name']}\t{s['mount'] or ''}")
        elif s["kind"] in ("mysql", "mongo"):
            print(f"skip\t{s['name']}\t{s['kind']}: ask Shelley to dump it")
        elif s["mount"]:
            print(f"volume\t{s['name']}\t{s['mount']}")
    for s in svcs:
        if s["kind"] in ("code", "image"):
            print(f"stop\t{s['name']}")
    for s in svcs:
        if s["kind"] == "code":
            print(f"push\t{s['name']}\t{s['git']}\t{s['repo'] or ''}\t{s['branch']}")


def cmd_list():
    svcs, _ = services()
    for s in svcs:
        print(f"{s['name']}\t{s['slug']}\t{s['kind']}\t{s['mount'] or '-'}\t{'cron' if s['cron'] else '-'}")


def new_password():
    a = string.ascii_letters + string.digits
    return "".join(secrets.choice(a) for _ in range(32))


def cmd_setup():
    svcs, env = services()
    vm = vm_name()
    host = f"{vm}.exe.xyz"
    notes = {"sealed": [], "keys": [], "unresolved": [], "warn": [], "multiline": []}

    # 1. Variables, as Railway resolved them for each service.
    for s in svcs:
        path = os.path.join(RW, "variables", s["name"] + ".json")
        raw = {k: (v or {}) for k, v in (s["inst"].get("variables") or {}).items()}
        if os.path.exists(path):
            resolved = load(path)
        else:  # fall back to the unresolved config; ${{...}} references stay
            resolved = {k: v.get("value") for k, v in raw.items()}
            notes["warn"].append(f"{s['name']}: no variables/{s['name']}.json; used unresolved values")
        s["vars"] = {}
        for k, v in resolved.items():
            if v is None or raw.get(k, {}).get("isSealed"):
                notes["sealed"].append((s["slug"], k))
                s["vars"][k] = ""
                continue
            s["vars"][k] = str(v)
            if "${{" in s["vars"][k]:
                notes["unresolved"].append((s["slug"], k))

    # 2. What changes: database passwords, public and TCP proxy addresses.
    replace = {}
    try:  # run again, setup keeps the passwords it made the first time
        with open(COMPOSE) as f:
            before = json.load(f)["services"]
    except (FileNotFoundError, ValueError, KeyError):
        before = {}
    for s in svcs:
        v = s["vars"]
        for name in PASSWORD_VARS.get(s["kind"], []):
            old = v.get(name)
            if old and len(old) >= 12 and old not in replace:
                kept = before.get(s["slug"], {}).get("environment", {}).get(name, "").replace("$$", "$")
                replace[old] = kept or new_password()
        pub = v.get("RAILWAY_PUBLIC_DOMAIN")
        if pub and pub.endswith(".up.railway.app"):
            s["railway_url"] = "https://" + pub
        for d in s["service_domains"]:
            s.setdefault("railway_url", "https://" + d)
        tcp_host, tcp_port = v.get("RAILWAY_TCP_PROXY_DOMAIN"), v.get("RAILWAY_TCP_PROXY_PORT")
        app_port = v.get("RAILWAY_TCP_APPLICATION_PORT") or {"postgres": "5432", "redis": "6379", "mysql": "3306", "mongo": "27017"}.get(s["kind"])
        private = v.get("RAILWAY_PRIVATE_DOMAIN") or f"{s['private'] or s['slug']}.railway.internal"
        s["alias"] = private
        if tcp_host and tcp_port and app_port:
            replace[f"{tcp_host}:{tcp_port}"] = f"{private}:{app_port}"

    # Public services: the first gets port 8000, the rest 8001 and up.
    web = [s for s in svcs if (s["service_domains"] or s["custom_domains"]) and not s["cron"]]
    for i, s in enumerate(web):
        s["host_port"] = 8000 + i
        ports = [c.get("port") for c in list(s["service_domains"].values()) + list(s["custom_domains"].values()) if c and c.get("port")]
        s["port"] = str(ports[0]) if ports else (s["vars"].get("PORT") or "8080")
        s["public"] = f"https://{host}/" if i == 0 else f"https://{host}:{s['host_port']}/"
        old = s["vars"].get("RAILWAY_PUBLIC_DOMAIN")
        new = host if i == 0 else f"{host}:{s['host_port']}"
        if old and old.endswith(".up.railway.app"):
            replace[old] = new
        for d in s["service_domains"]:
            replace.setdefault(d, new)
    if len(web) > 1:
        notes["warn"].append("more than one public service: " + ", ".join(f"{s['slug']} on {s['public']}" for s in web)
                             + "; custom domains reach port 8000 only")

    def rewrite(val):
        for old, new in replace.items():
            val = val.replace(old, new)
        return val

    for s in svcs:
        v = {k: rewrite(x) for k, x in s["vars"].items()}
        v.setdefault("RAILWAY_PRIVATE_DOMAIN", s["alias"])
        v.setdefault("RAILWAY_SERVICE_NAME", s["name"])
        if s.get("port"):
            v["PORT"] = s["port"]
        elif s["kind"] == "code":
            v.setdefault("PORT", "8080")
        if s["mount"]:
            v["RAILWAY_VOLUME_MOUNT_PATH"] = s["mount"]
            v.setdefault("RAILWAY_VOLUME_NAME", s["slug"] + "-volume")
        if s["kind"] in PASSWORD_VARS and "RAILWAY_TCP_PROXY_DOMAIN" in v:
            v["RAILWAY_TCP_PROXY_DOMAIN"], v["RAILWAY_TCP_PROXY_PORT"] = s["alias"], v.get("RAILWAY_TCP_APPLICATION_PORT", "")
        for k, x in v.items():
            if "\n" in x:
                notes["multiline"].append((s["slug"], k))
            for pat, handle in KEY_HINTS:
                if re.search(pat, k) and x and not APP_SECRET_HINT.search(k) and not k.startswith("RAILWAY_"):
                    notes["keys"].append((s["slug"], k, handle))
                    break
        s["vars"] = v

    # 3. compose.yaml (JSON is YAML; it keeps every value exact).
    def esc(x):
        return x.replace("$", "$$")

    compose = {"name": "app", "services": {}, "networks": {"railway": {}}}
    meta = {}
    for s in svcs:
        c = {}
        d = s["deploy"]
        start = d.get("startCommand")
        if s["kind"] == "code":
            c["image"] = f"railway/{s['slug']}:current"
            c["pull_policy"] = "never"
            builder = (s["build"].get("builder") or "RAILPACK").upper()
            dockerfile = s["vars"].get("RAILWAY_DOCKERFILE_PATH") or s["build"].get("dockerfilePath")
            if builder == "NIXPACKS":
                notes["warn"].append(f"{s['slug']}: built with Nixpacks on Railway; the VM builds it with Railpack, Railway's successor")
            meta[s["slug"]] = {"name": s["name"], "git": s["git"], "branch": s["branch"], "root": s["root"],
                               "builder": builder, "dockerfile": dockerfile, "buildCommand": s["build"].get("buildCommand"),
                               "startCommand": start, "preDeploy": d.get("preDeployCommand"),
                               "healthcheckPath": d.get("healthcheckPath"), "healthcheckTimeout": d.get("healthcheckTimeout") or 300,
                               "cron": bool(s["cron"]), "port": s.get("host_port")}
        else:
            c["image"] = s["image"]
        if start:
            if s["kind"] == "code" and meta[s["slug"]]["builder"] != "DOCKERFILE":
                c["command"] = [esc(start)]  # Railpack images run their command with bash -c
            else:
                c["entrypoint"] = [esc(a) for a in shlex.split(start)]  # Railway replaces the ENTRYPOINT
        c["environment"] = {k: esc(x) for k, x in sorted(s["vars"].items())}
        if s["mount"]:
            c["volumes"] = [f"./volumes/{s['slug']}:{s['mount']}"]
        c["networks"] = {"railway": {"aliases": [s["alias"]]}}
        if s.get("host_port"):
            c["ports"] = [f"{s['host_port']}:{s['port']}"]
        if s["cron"]:
            c["profiles"] = ["cron"]
            c["restart"] = "no"
        else:
            c["restart"] = "no" if d.get("restartPolicyType") == "NEVER" else "unless-stopped"
        drain = d.get("drainingSeconds") or s["vars"].get("RAILWAY_DEPLOYMENT_DRAINING_SECONDS")
        if drain:
            c["stop_grace_period"] = f"{int(drain)}s"
        compose["services"][s["slug"]] = c

        if (d.get("numReplicas") or 1) > 1 or len(d.get("multiRegionConfig") or {}) > 1:
            notes["warn"].append(f"{s['slug']}: ran several replicas on Railway; one container here")
        if d.get("sleepApplication"):
            notes["warn"].append(f"{s['slug']}: slept when idle on Railway; runs all the time here")
        if (d.get("registryCredentials") or {}).get("username"):
            notes["warn"].append(f"{s['slug']}: image from a private registry; run docker login on the VM first")
        if s["kind"] in ("mysql", "mongo"):
            notes["warn"].append(f"{s['slug']}: {s['kind']} data is not copied by the hand-over; dump and restore it by hand")
        if s["tcp"] and s["kind"] not in PASSWORD_VARS:
            notes["warn"].append(f"{s['slug']}: had a public TCP proxy; exe.dev serves HTTP only, so reach it over SSH")
    if env.get("buckets"):
        notes["warn"].append("the environment has Railway buckets; the app keeps using them until you move the files")
    if env.get("sharedVariables"):
        notes["warn"].append("shared variables were copied into each service that uses them")

    os.makedirs(APP, exist_ok=True)
    for s in svcs:
        if s["mount"]:
            os.makedirs(os.path.join(APP, "volumes", s["slug"]), exist_ok=True)
    old = os.umask(0o077)
    try:
        with open(COMPOSE, "w") as f:
            json.dump(compose, f, indent=2)
            f.write("\n")
    finally:
        os.umask(old)
    with open(os.path.join(APP, "services.json"), "w") as f:
        json.dump(meta, f, indent=2)
        f.write("\n")

    # 4. Cron services: crontab lines, installed by railway/move.
    with open(os.path.join(APP, "crontab.jobs"), "w") as f:
        f.write("# Railway cron services, in UTC like on Railway. A run is skipped while the last one is still going.\n")
        for s in svcs:
            if s["cron"]:
                f.write(f"{s['cron']} cd ~/app && flock -n /tmp/cron-{s['slug']}.lock docker compose run --rm -T {s['slug']} 2>&1 | systemd-cat -t cron-{s['slug']}\n")

    # The address railway/check compares with: a custom domain, or Railway's own.
    first = web[0] if web else None
    if first:
        url = ("https://" + sorted(first["custom_domains"])[0]) if first["custom_domains"] else first.get("railway_url", "")
        with open(os.path.join(RW, "railway-url"), "w") as f:
            f.write(url + "\n")

    # Code services: a bare repository each, deploying on push.
    for g in sorted({s["git"] for s in svcs if s["kind"] == "code"}):
        repo = os.path.join(HOME, "git", g + ".git")
        subprocess.run(["git", "init", "-q", "--bare", repo], check=True)
        hook = os.path.join(repo, "hooks", "pre-receive")
        if not os.path.lexists(hook):
            os.symlink(os.path.join(RW, "deploy"), hook)

    if not os.path.exists(os.path.join(APP, "crontab.base")):
        with open(os.path.join(APP, "crontab.base"), "w") as f:
            f.write("# Nightly backup of the databases and volumes into ~/backups\n"
                    "17 3 * * * ~/railway/backup 2>&1 | systemd-cat -t railway-backup\n")

    # 5. The plan, for the report.
    lines = [f"# Plan for {vm}", "", "| Railway service | Kind | On the VM |", "|---|---|---|"]
    for s in svcs:
        what = {"code": "built from " + (f"`{s['repo']}`" if s["repo"] else "pushed code") + f" into `~/git/{s.get('git')}.git`",
                "image": f"image `{s['image']}`"}.get(s["kind"], f"`{s['image']}`, data restored from the hand-over")
        where = f"compose service `{s['slug']}`"
        if s.get("public"):
            where += f", {s['public']}"
        if s["cron"]:
            where += f", cron `{s['cron']}` (UTC)"
        if s["mount"]:
            where += f", volume `~/app/volumes/{s['slug']}`"
        lines.append(f"| {s['name']} | {what} | {where} |")
    custom = sorted({d for s in svcs for d in s["custom_domains"]})
    lines += ["", "Custom domains: " + (", ".join(custom) if custom else "none"), ""]
    lines += ["Railway URLs: " + (", ".join(s["railway_url"] for s in svcs if s.get("railway_url")) or "none"), ""]
    if replace:
        lines += [f"Rewritten in variables: {sum(1 for o in replace if '.' in o)} Railway addresses, "
                  f"{sum(1 for o in replace if '.' not in o)} database passwords (new ones, made on the VM).", ""]
    def section(title, items):
        if items:
            lines.extend([f"## {title}", ""] + [f"- {i}" for i in items] + [""])
    section("Sealed variables to enter again", [f"`{k}` in `{svc}`" for svc, k in notes["sealed"]])
    section("Keys for other services, candidates for an exe.dev integration",
            [f"`{k}` in `{svc}`" + (f" (catalog: `{h}`)" if h else "") for svc, k, h in notes["keys"]])
    section("Unresolved references", [f"`{k}` in `{svc}`" for svc, k in notes["unresolved"]])
    section("Multi-line values (kept exact)", [f"`{k}` in `{svc}`" for svc, k in notes["multiline"]])
    section("Differences", notes["warn"])
    with open(os.path.join(RW, "plan.md"), "w") as f:
        f.write("\n".join(lines))
    print("\n".join(lines))


def cmd_set(svc, key, value=None):
    if value is None:
        value = sys.stdin.read()
        if value.endswith("\n") and value.count("\n") == 1:
            value = value[:-1]
    with open(COMPOSE) as f:
        text = f.read()
    try:
        compose = json.loads(text)
    except ValueError:
        import yaml
        compose = yaml.safe_load(text)
    if svc not in compose["services"]:
        sys.exit(f"set: no service {svc} in {COMPOSE}")
    compose["services"][svc].setdefault("environment", {})[key] = value.replace("$", "$$")
    tmp = COMPOSE + ".tmp"
    old = os.umask(0o077)
    with open(tmp, "w") as f:
        json.dump(compose, f, indent=2)
        f.write("\n")
    os.umask(old)
    os.replace(tmp, COMPOSE)
    print(f"set {key} in {svc}")
    if "cron" not in compose["services"][svc].get("profiles", []):
        subprocess.run(["docker", "compose", "up", "-d", svc], cwd=APP, check=True)


if __name__ == "__main__":
    a = sys.argv[1:]
    if a[:1] == ["plan"]:
        cmd_plan()
    elif a[:1] == ["list"]:
        cmd_list()
    elif a[:1] == ["setup"]:
        cmd_setup()
    elif a[:1] == ["set"] and len(a) in (3, 4):
        cmd_set(*a[1:])
    else:
        sys.exit(__doc__)