#!/usr/bin/env python3
"""Pack the BAML wrapper using an independent, explicit compiler and pack host."""

import argparse
import json
from pathlib import Path
import os
import re
import shutil
import subprocess
import tempfile

repo = Path(__file__).resolve().parents[2]
parser = argparse.ArgumentParser()
parser.add_argument("--cli", type=Path, required=True)
parser.add_argument("--pack-host", type=Path)
parser.add_argument("--target")
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--version")
parser.add_argument("--no-self-update", action="store_true")
parser.add_argument(
    "--embed-telemetry",
    metavar="ENV_VAR",
    help="embed the named environment variable's value as the BCS telemetry API key",
)
args = parser.parse_args()
# Match a Cargo package version declaration, e.g. version = "0.2.5".
version = (
    args.version
    or re.search(
        r'^version = "([^"]+)"',
        (repo / "baml_language/crates/baml/Cargo.toml").read_text(),
        re.M,
    )[1]
)
# Wrapper release identifiers, e.g. 0.2.5 or 0.2.5-canary.1.
if not re.fullmatch(r"[0-9]+\.[0-9]+\.[0-9]+(?:[-+][0-9A-Za-z.-]+)?", version):
    parser.error("invalid wrapper version")
with tempfile.TemporaryDirectory(prefix="baml-wrapper-build-") as temporary:
    project = Path(temporary) / "wrapper"
    shutil.copytree(repo / "baml_language/wrapper", project)
    post = (repo / "baml_language/crates/baml/src/telemetry.rs").read_text()

    # Named Rust string constants, e.g. const POSTHOG_HOST: &str = "https://...".
    def constant(source, name):
        return re.search(rf'const {name}: &str = "([^"]*)";', source)[1]

    constants = {
        "wrapper_version": version,
        "self_update_enabled": not args.no_self_update,
        "telemetry_api_key": constant(post, "POSTHOG_API_KEY"),
        "telemetry_default_host": constant(post, "POSTHOG_HOST"),
    }
    (project / "baml_src/build.baml").write_text(
        "\n".join(
            f"function {name}() -> {'bool' if isinstance(value, bool) else 'string'} {{\n    {json.dumps(value)}\n}}\n"
            for name, value in constants.items()
        )
    )
    command = [
        str(args.cli.resolve()),
        "pack",
        "--from",
        str(project),
        "main",
        "--output",
        str(args.output.resolve()),
    ]
    if args.embed_telemetry:
        command.extend(["--embed-telemetry", args.embed_telemetry])
    if args.target:
        command.extend(["--target", args.target])
    if args.pack_host:
        command.extend(["--host-binary", str(args.pack_host.resolve())])
    subprocess.run(
        command, check=True, cwd=repo, env=dict(os.environ, BAML_CLI_ALLOW_DIRECT="1")
    )