#!/usr/bin/env python3
"""Build on .146 and release to native production :64192, with local SSH keys.

Requires Python 3.9+ and Paramiko on the operator computer only. Start with
`python scripts/production48.py status`; see docs/production-48.md.
"""

import argparse
import contextlib
import datetime
import hashlib
import json
import os
from pathlib import Path
import re
import shlex
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent.parent
BUILD_ROOT = "/opt/ai-chatbot-builds"
RELEASE_ROOT = "/var/www/ai-chatbot-releases"
PROD_FINGERPRINT = "SHA256:3tFAg02kzAfGNBDba+je+vgSB7TifvJPje9XIyVBgTM"
BUILD_FINGERPRINT = "SHA256:MXW5w48SFSTUaqS+dS1G9hY228hI6SLepjhCmItHL04"


def require(value, message):
    if not value:
        raise RuntimeError(message)


def git(*args):
    return subprocess.check_output(["git", *args], cwd=ROOT, text=True).strip()


def sha(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_json(path, value):
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf8")


class Connections:
    def __init__(self, args):
        import paramiko
        self.paramiko = paramiko
        self.args = args
        self.stack = contextlib.ExitStack()
        self.clients = {}

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.stack.close()

    def client(self, kind):
        import base64
        if kind in self.clients:
            return self.clients[kind]
        client = self.paramiko.SSHClient()
        self.stack.callback(client.close)
        client.load_host_keys(str(self.args.known_hosts))
        client.set_missing_host_key_policy(self.paramiko.RejectPolicy())
        options = dict(username="root", allow_agent=False, look_for_keys=False,
                       timeout=20, banner_timeout=20, auth_timeout=20)
        if kind == "jump":
            host, port, key = "192.168.58.20", 22, self.args.jump_key
        elif kind == "build":
            host, port, key = "192.168.102.146", 22, self.args.key
            # The LAN entrance has the same verified key as the existing FRP entry.
            known = client.get_host_keys().lookup("[45.192.99.48]:23322")
            require(known and "ssh-ed25519" in known, "Missing trusted builder host key")
            client.get_host_keys().add(host, "ssh-ed25519", known["ssh-ed25519"])
        else:
            host, port, key = "45.192.99.48", 64192, self.args.key
            pin = self.paramiko.Ed25519Key(data=base64.b64decode(
                "AAAAC3NzaC1lZDI1NTE5AAAAIK7afV+YwLqvUFBlxDe/q2dUrpBePcLdG+mZ0JekHYPC"))
            client.get_host_keys().add("[45.192.99.48]:64192", "ssh-ed25519", pin)
            jump = self.client("jump")
            options["sock"] = jump.get_transport().open_channel(
                "direct-tcpip", (host, port), ("127.0.0.1", 0), timeout=20)
        options.update(hostname=host, port=port, key_filename=str(key))
        client.connect(**options)
        actual = "SHA256:" + base64.b64encode(hashlib.sha256(
            client.get_transport().get_remote_server_key().asbytes()).digest()).decode().rstrip("=")
        expected = {"build": BUILD_FINGERPRINT, "prod": PROD_FINGERPRINT}.get(kind)
        require(not expected or actual == expected, f"{kind} host fingerprint mismatch")
        client.get_transport().set_keepalive(20)
        self.clients[kind] = client
        return client

    def execute(self, kind, command, stdin=None, stream=False, timeout=1800):
        channel = self.client(kind).get_transport().open_session(timeout=20)
        channel.exec_command(command)
        if stdin is not None:
            channel.sendall(stdin.encode("utf8"))
        channel.shutdown_write()
        output, errors = bytearray(), bytearray()
        deadline = time.monotonic() + timeout
        try:
            while True:
                if channel.recv_ready():
                    chunk = channel.recv(65536)
                    output.extend(chunk)
                    if stream:
                        print(chunk.decode("utf8", "replace"), end="", flush=True)
                if channel.recv_stderr_ready():
                    chunk = channel.recv_stderr(65536)
                    errors.extend(chunk)
                    if stream:
                        print(chunk.decode("utf8", "replace"), end="", file=sys.stderr, flush=True)
                if channel.exit_status_ready() and not channel.recv_ready() and not channel.recv_stderr_ready():
                    break
                require(time.monotonic() < deadline, f"{kind} operation timed out; inspect remote state")
                time.sleep(0.03)
            code = channel.recv_exit_status()
            require(code == 0, f"{kind} command failed ({code}): {errors.decode('utf8', 'replace')[-1800:]}")
            return output.decode("utf8")
        finally:
            channel.close()

    def helper(self, action, **kwargs):
        command = ["python3", "-", action]
        for key, value in kwargs.items():
            command += ["--" + key, str(value)]
        source = (ROOT / "scripts/release/host.py").read_text(encoding="utf8")
        return json.loads(self.execute("prod", shlex.join(command), stdin=source))

    def sftp(self, kind):
        return self.paramiko.SFTPClient.from_transport(
            self.client(kind).get_transport(), window_size=16 * 1024**2, max_packet_size=65536)

    def upload(self, kind, local, remote):
        last = [0]
        def progress(done, total):
            if done - last[0] >= 32 * 1024**2:
                print(f"[{kind}] transferred {done // 1024**2}/{total // 1024**2} MiB", flush=True)
                last[0] = done
        with self.sftp(kind) as transfer:
            transfer.put(str(local), remote, callback=progress)
            transfer.chmod(remote, 0o600)


def build(connections, args):
    require(not git("status", "--porcelain"), "Commit requested changes before building")
    commit = git("rev-parse", "origin/main^{commit}")
    require(git("branch", "--show-current") == "main" and git("rev-parse", "HEAD") == commit,
            "Fetch origin and switch to an up-to-date, clean main before building")
    remote = connections.execute("build", "hostname").strip()
    require(remote == "dao02-Default-string", "Unexpected builder")
    require(connections.execute("build", "node -v").strip().startswith("v20."), "Builder requires Node 20")
    release_id = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ-") + commit[:12]
    local = args.receipts / release_id
    local.mkdir(parents=True)
    source = local / "source.tar.gz"
    subprocess.run(["git", "archive", "--format=tar.gz", f"--output={source}", commit], cwd=ROOT, check=True)
    metadata = {"sourceCommit": commit, "sourceTree": git("rev-parse", f"{commit}^{{tree}}"),
                "sourceArchiveSha256": sha(source)}
    write_json(local / "input.json", metadata)
    environment = connections.helper("build-env")
    environment.update({"AUTH_SECRET": os.urandom(32).hex(), "POSTGRES_URL": "postgres://build:build@127.0.0.1:9/build",
                        "AUTH_TRUST_HOST": "true", "NODE_ENV": "production"})
    workspace = f"{BUILD_ROOT}/{release_id}"
    connections.execute("build", f"install -d -m 700 -- {shlex.quote(workspace)}")
    for name in ["source.tar.gz", "input.json"]:
        connections.upload("build", local / name, f"{workspace}/{name}")
    # Helpers are carried separately, with the reviewed operator checkout's hashes.
    for name in ["host.py", "build.sh"]:
        connections.upload("build", ROOT / "scripts/release" / name, f"{workspace}/{name}")
    setup = f"""import sys,json,pathlib,hashlib,importlib.util,os
os.umask(0o077)
workspace=pathlib.Path({workspace!r})
assert hashlib.sha256((workspace/'source.tar.gz').read_bytes()).hexdigest()=={metadata['sourceArchiveSha256']!r}
spec=importlib.util.spec_from_file_location('release_host',workspace/'host.py'); module=importlib.util.module_from_spec(spec);spec.loader.exec_module(module)
module.safe_extract(workspace/'source.tar.gz',workspace/'source')
environment=json.load(sys.stdin)
with open(workspace/'source/.env.production.local','x') as stream:
 for key,value in environment.items():stream.write(key+'='+json.dumps(value)+'\\n')
"""
    # stdin belongs to the program, so the program itself travels as a quoted -c argument.
    connections.execute("build", shlex.join(["python3", "-c", setup]), stdin=json.dumps(environment))
    print(f"[build] {release_id} from main {commit}", flush=True)
    log = connections.execute("build", shlex.join(["nice", "-n", "10", "bash", f"{workspace}/build.sh", workspace]),
                              stream=True, timeout=2400)
    (local / "build.log").write_text(log, encoding="utf8")
    smoke = json.loads(connections.execute("build", shlex.join(
        ["python3", f"{workspace}/host.py", "smoke", "--workspace", workspace]), timeout=240))
    write_json(local / "smoke.json", smoke)
    print(json.dumps(smoke), flush=True)
    print("[pack] packaging source, build, pinned dependencies and public assets", flush=True)
    result = json.loads(connections.execute("build", shlex.join(
        ["python3", f"{workspace}/host.py", "pack", "--workspace", workspace]), timeout=600))
    result.update({"builder": remote, "mainRef": commit,
                   "helperHashes": {name: sha(ROOT / "scripts/release" / name) for name in ["host.py", "build.sh"]}})
    write_json(local / "build.json", result)
    print(json.dumps(result), flush=True)


def stage(connections, args):
    local = args.receipts / args.release
    receipt = json.loads((local / "build.json").read_text(encoding="utf8"))
    artifact = local / "artifact.tar.gz"
    if not artifact.exists():
        with connections.sftp("build") as transfer:
            transfer.get(f"{BUILD_ROOT}/{args.release}/artifact.tar.gz", str(artifact),
                         max_concurrent_prefetch_requests=32)
    require(sha(artifact) == receipt["artifactSha256"], "Downloaded artifact checksum mismatch")
    state = connections.helper("status")
    require(state["freeBytes"] > artifact.stat().st_size + 1024**3 + 64 * 1024**2,
            "Insufficient space to upload artifact while retaining 1 GiB reserve")
    connections.execute("prod", f"install -d -m 700 -- {RELEASE_ROOT}")
    connections.upload("prod", artifact, f"{RELEASE_ROOT}/{args.release}.tar.gz")
    result = connections.helper("stage", release=args.release, digest=receipt["artifactSha256"])
    write_json(local / "stage.json", result)
    print(json.dumps(result))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["status", "jobs", "build", "stage", "activate", "rollback"])
    parser.add_argument("--release")
    parser.add_argument("--key", type=Path, default=Path.home() / ".ssh/45.192.99.48.pem")
    parser.add_argument("--jump-key", type=Path, default=Path.home() / ".ssh/youyidian_prod_ed25519")
    parser.add_argument("--known-hosts", type=Path, default=Path.home() / ".ssh/known_hosts")
    parser.add_argument("--receipts", type=Path, default=ROOT / ".git/releases")
    args = parser.parse_args()
    if args.action in {"stage", "activate", "rollback"}:
        require(args.release and re.fullmatch(r"\d{8}T\d{6}Z-[0-9a-f]{12}", args.release), "--release is required")
    args.receipts = args.receipts.resolve()
    with Connections(args) as connections:
        if args.action == "build":
            build(connections, args)
        elif args.action == "stage":
            stage(connections, args)
        else:
            kwargs = {"release": args.release} if args.release else {}
            result = connections.helper(args.action, **kwargs)
            if args.release:
                write_json(args.receipts / args.release / (args.action + ".json"), result)
            print(json.dumps(result))


if __name__ == "__main__":
    main()