#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11,<3.12"
# dependencies = [
#     "modal==1.4.1",
#     "click>=8.0",
#     "dockerfile-parse>=2.0.0",
# ]
# ///
"""Modal sandbox management for Offload.

Unified CLI for creating, executing commands on, and destroying Modal sandboxes.
"""

import sys

sys.dont_write_bytecode = True

import io
import json
import logging
import math
import os
import tarfile
import tempfile
import threading
import time
from pathlib import Path

import click
import modal
from dockerfile_parse import DockerfileParser

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
handler = logging.StreamHandler(sys.stderr)
handler.setFormatter(logging.Formatter("%(message)s"))
logger.addHandler(handler)


def copy_dir_to_sandbox(sandbox, local_dir: str, remote_dir: str) -> None:
    """Recursively copy a local directory to the sandbox using tar."""
    logger.info("Creating tar archive from %s...", local_dir)

    # Create tar archive in memory
    tar_buffer = io.BytesIO()

    with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
        for root, dirs, files in os.walk(local_dir):
            # Filter directories in-place
            dirs[:] = [
                d
                for d in dirs
                if not d.startswith(".")
                and d not in ("__pycache__", "node_modules", "target", ".venv", "venv")
            ]

            for fname in files:
                if fname.startswith(".") or fname.endswith(".pyc"):
                    continue
                local_path = os.path.join(root, fname)
                rel_path = os.path.relpath(local_path, local_dir)
                tar.add(local_path, arcname=rel_path)

    tar_buffer.seek(0)
    tar_data = tar_buffer.getvalue()

    logger.info("Transferring tar archive (%d bytes) to sandbox...", len(tar_data))

    # Create remote directory and transfer tar
    sandbox.mkdir(remote_dir, parents=True)
    tar_remote_path = f"{remote_dir}/.transfer.tar"
    with sandbox.open(tar_remote_path, "wb") as f:
        f.write(tar_data)

    logger.info("Extracting tar archive in %s...", remote_dir)

    # Extract on sandbox
    sandbox.exec("tar", "-xf", tar_remote_path, "-C", remote_dir).wait()

    # Clean up tar file
    sandbox.exec("rm", "-f", tar_remote_path).wait()

    logger.info("Tar-based transfer complete")


def copy_from_sandbox(sandbox, remote_path: str, local_path: str) -> None:
    """Copy a file from the sandbox to local filesystem."""
    logger.info("Downloading %s to %s...", remote_path, local_path)

    # Read file content directly from sandbox
    with sandbox.open(remote_path, "rb") as f:
        data = f.read()

    logger.info("Received %d bytes", len(data))

    # Create parent directory if needed
    local_parent = os.path.dirname(local_path.rstrip("/")) or "."
    os.makedirs(local_parent, exist_ok=True)

    # Write to local file
    with open(local_path, "wb") as f:
        f.write(data)

    logger.info("Download complete: %s -> %s", remote_path, local_path)


@click.group()
def cli():
    """Modal sandbox management for Offload."""
    pass


DOCKERIGNORE_FILE = ".dockerignore"


def read_dockerignore_patterns() -> list[str]:
    """Read patterns from .dockerignore file."""
    if not os.path.isfile(DOCKERIGNORE_FILE):
        return []
    patterns = []
    with open(DOCKERIGNORE_FILE) as f:
        for line in f:
            line = line.strip()
            # Skip empty lines and comments
            if line and not line.startswith("#"):
                patterns.append(line)
    return patterns


_LAYER_BOUNDARY_INSTRUCTIONS = frozenset({"RUN", "COPY", "ADD"})


def _build_image_from_dockerfile(
    dockerfile_path: str,
    context_dir: str = ".",
) -> modal.Image:
    """Parse a Dockerfile and build a Modal image with per-layer caching.

    Instead of building the entire Dockerfile as a single monolithic image via
    modal.Image.from_dockerfile(), this function parses the Dockerfile into its
    constituent instructions and applies them one batch at a time using
    modal.Image.dockerfile_commands(). Each batch ends at a RUN, COPY, or ADD
    instruction (or at end-of-file), so that Modal caches each layer
    independently. Non-filesystem instructions (ENV, ARG, WORKDIR, etc.) are
    batched together with the next layer-boundary instruction to reduce API
    round-trips without sacrificing cacheability.

    This approach prevents Modal's image builder from OOMing during the
    post-build save/materialize phase for large Dockerfiles, since each layer
    is materialized separately rather than as one giant blob.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        tmpfile = Path(tmpdir) / "Dockerfile"
        # When context_dir is provided, resolve the Dockerfile relative to it
        # so that checkpoint builds use the Dockerfile from the exported tree.
        resolved_dockerfile = Path(context_dir) / dockerfile_path
        with open(resolved_dockerfile) as f:
            dockerfile_contents = f.read()
        tmpfile.write_text(dockerfile_contents)

        dfp = DockerfileParser(str(tmpfile))

        if dfp.is_multistage:
            logger.warning(
                "Multistage Dockerfiles are not supported for per-layer "
                "decomposition; falling back to monolithic from_dockerfile()"
            )
            return modal.Image.from_dockerfile(dockerfile_path, context_dir=context_dir)

        # Find the last FROM instruction
        last_from_index = None
        for i, instr in enumerate(dfp.structure):
            if instr["instruction"] == "FROM":
                last_from_index = i

        if last_from_index is None:
            logger.error("Dockerfile must contain a FROM instruction")
            sys.exit(1)

        base_image_ref = dfp.baseimage
        logger.info("Base image: %s", base_image_ref)
        image = modal.Image.from_registry(base_image_ref)

        instructions = dfp.structure[last_from_index + 1 :]

        # Batch instructions: accumulate until we hit a layer-boundary
        # instruction (RUN, COPY, ADD) or end-of-file, then flush as a single
        # dockerfile_commands() call. This reduces Modal API round-trips while
        # preserving per-layer caching at meaningful boundaries.
        batch: list[str] = []
        layer_count = 0
        for instr in instructions:
            if instr["instruction"] == "COMMENT":
                continue
            batch.append(instr["content"])
            if instr["instruction"] in _LAYER_BOUNDARY_INSTRUCTIONS:
                logger.info(
                    "Building layer %d (%d instruction(s), ending with %s)...",
                    layer_count,
                    len(batch),
                    instr["instruction"],
                )
                image = image.dockerfile_commands(
                    batch,
                    context_dir=context_dir,
                )
                batch = []
                layer_count += 1

        # Flush any trailing non-boundary instructions (e.g. ENV, EXPOSE, CMD)
        if batch:
            logger.info(
                "Building final layer %d (%d trailing instruction(s))...",
                layer_count,
                len(batch),
            )
            image = image.dockerfile_commands(
                batch,
                context_dir=context_dir,
            )
            layer_count += 1

        logger.info("Dockerfile decomposed into %d cached layers", layer_count)
        return image


def _build_fresh_base_image(
    app, dockerfile_path: str | None, context_dir: str = "."
) -> tuple[modal.Image, str]:
    """Build a fresh base image (no caching)."""
    if dockerfile_path is None:
        logger.info("Building default base image...")
        base_img = modal.Image.debian_slim(python_version="3.11").pip_install("pytest")
    else:
        logger.info("Building base image from %s with context_dir=%s", dockerfile_path, context_dir)
        base_img = _build_image_from_dockerfile(dockerfile_path, context_dir=context_dir)

    base_img.build(app)
    # Materialize to get base image_id
    temp_sandbox = modal.Sandbox.create(app=app, image=base_img, timeout=10)
    temp_sandbox.terminate()
    base_img_id = base_img.object_id
    return base_img, base_img_id


def _build_final_image(
    app,
    base_img: modal.Image,
    base_img_id: str,
    include_cwd: bool,
    copy_dirs: tuple[str, ...],
    ignore_patterns: list[str],
    sandbox_init_cmd: str | None = None,
) -> str:
    """Build final image with cwd/copy-dirs on top of base. Returns image_id."""
    final_img = base_img

    if include_cwd:
        logger.info("Adding current directory as /app...")
        final_img = final_img.add_local_dir(
            ".", "/app", copy=True, ignore=ignore_patterns
        )

    # Add user-specified directories
    for copy_spec in copy_dirs:
        if ":" not in copy_spec:
            logger.warning(
                "Invalid copy-dir format '%s', expected 'local:remote'",
                copy_spec,
            )
            continue
        local_path, remote_path = copy_spec.split(":", 1)
        if not os.path.isdir(local_path):
            logger.warning("Local directory '%s' not found, skipping", local_path)
            continue
        logger.info("Adding %s -> %s to image", local_path, remote_path)
        final_img = final_img.add_local_dir(
            local_path, remote_path, copy=True, ignore=ignore_patterns
        )

    if sandbox_init_cmd:
        logger.info("Running sandbox_init_cmd: %s", sandbox_init_cmd)
        final_img = final_img.run_commands(sandbox_init_cmd)

    # Build and materialize the final image if we added anything
    if final_img is not base_img:
        final_img.build(app)
        temp_sandbox = modal.Sandbox.create(app=app, image=final_img, timeout=10)
        temp_sandbox.terminate()
        return final_img.object_id
    else:
        return base_img_id


def _derive_image_from_base(
    app,
    base_img: modal.Image,
    patch_file: str | None,
    project_root: str,
    post_patch_cmd: str | None = None,
) -> str:
    """Apply a binary patch to a base image and return the new image ID."""
    img = base_img

    if patch_file is not None:
        img = img.add_local_file(patch_file, "/tmp/offload.patch", copy=True)
        img = img.run_commands(
            f"offload apply-diff /tmp/offload.patch --project-root {project_root}"
        )

        if post_patch_cmd:
            logger.info("Running post_patch_cmd: %s", post_patch_cmd)
            img = img.run_commands(
                post_patch_cmd,
                env={"OFFLOAD_PATCH_FILE": "/tmp/offload.patch"},
            )

        img = img.run_commands("rm /tmp/offload.patch")
    elif post_patch_cmd:
        logger.info("Running post_patch_cmd (no patch): %s", post_patch_cmd)
        img = img.run_commands(post_patch_cmd)

    img.build(app)
    temp_sandbox = modal.Sandbox.create(app=app, image=img, timeout=10)
    temp_sandbox.terminate()
    return img.object_id


@cli.command("prepare")
@click.argument("dockerfile_path", required=False, default=None)
@click.option(
    "--include-cwd",
    is_flag=True,
    help="Include current directory in the image (added after cache lookup)",
)
@click.option(
    "--copy-dir",
    "copy_dirs",
    multiple=True,
    help="Copy local dir into image (format: local_path:remote_path)",
)
@click.option(
    "--sandbox-init-cmd",
    default=None,
    help="Command to run during image build after cwd/copy-dirs are applied",
)
@click.option(
    "--context-dir",
    default=".",
    help="Docker build context directory",
)
@click.option(
    "--cached",
    is_flag=True,
    hidden=True,
    help="Deprecated no-op kept for backward compatibility",
)
@click.option(
    "--from-base-image",
    default=None,
    help="Modal image ID to build incrementally on top of",
)
@click.option(
    "--patch-file",
    default=None,
    help="Path to a binary patch file to apply on top of the base image",
)
@click.option(
    "--sandbox-project-root",
    default="/app",
    help="Project root path in sandbox",
)
@click.option(
    "--post-patch-cmd",
    default=None,
    help="Command to run after thin-diff patch is applied",
)
def prepare(
    dockerfile_path: str | None,
    include_cwd: bool,
    copy_dirs: tuple[str, ...],
    sandbox_init_cmd: str | None,
    context_dir: str,
    cached: bool,
    from_base_image: str | None,
    patch_file: str | None,
    sandbox_project_root: str,
    post_patch_cmd: str | None,
):
    """Prepare a Modal image (build only, no sandbox creation).

    DOCKERFILE_PATH: Optional path to a Dockerfile. If provided, builds from
    that Dockerfile. If omitted, builds the default pytest image.

    The --include-cwd and --copy-dir options add source code to the image.

    When --from-base-image is set, builds incrementally from the given base
    image by applying the patch file (if provided) on top of it.

    Prints the image_id to stdout for use with 'create'.
    """
    # Incremental build mode
    if from_base_image is not None:
        with modal.enable_output():
            app_name = "offload-checkpoint-sandbox"
            app = modal.App.lookup(app_name, create_if_missing=True)

            base_img = modal.Image.from_id(from_base_image)

            logger.info("Building incremental image from base")

            image_id = _derive_image_from_base(
                app,
                base_img,
                patch_file,
                sandbox_project_root,
                post_patch_cmd=post_patch_cmd,
            )

        sys.stdout.write("%s\n" % image_id)
        return

    # Standard (full build) mode
    # Read ignore patterns from .dockerignore
    ignore_patterns = read_dockerignore_patterns()
    if ignore_patterns:
        logger.debug(
            "Using %d ignore patterns from %s", len(ignore_patterns), DOCKERIGNORE_FILE
        )

    # Determine app name based on whether we have a Dockerfile
    if dockerfile_path is None:
        app_name = "offload-sandbox"
    else:
        if not os.path.isfile(dockerfile_path):
            logger.error("Error: Dockerfile not found: %s", dockerfile_path)
            sys.exit(1)
        app_name = "offload-dockerfile-sandbox"

    with modal.enable_output():
        app = modal.App.lookup(app_name, create_if_missing=True)

        base_image, base_image_id = _build_fresh_base_image(app, dockerfile_path, context_dir)

        image_id = _build_final_image(
            app,
            base_image,
            base_image_id,
            include_cwd,
            copy_dirs,
            ignore_patterns,
            sandbox_init_cmd=sandbox_init_cmd,
        )

    sys.stdout.write("%s\n" % image_id)


@cli.command()
@click.argument("sandbox_id")
def destroy(sandbox_id: str):
    """Terminate a Modal sandbox."""
    sandbox = modal.Sandbox.from_id(sandbox_id)
    sandbox.terminate()
    logger.info("Terminated sandbox %s", sandbox_id)


@cli.command("download")
@click.argument("sandbox_id")
@click.argument("paths", nargs=-1, required=True)
def download(sandbox_id: str, paths: tuple[str, ...]):
    """Download files or directories from a Modal sandbox.

    SANDBOX_ID is the Modal sandbox ID to download from.

    PATHS are one or more path specifications in the format "remote_path:local_path".
    Each specification downloads the remote file to the local path.

    Examples:

        modal_sandbox.py download sb-abc123 "/tmp/junit.xml:./results/junit.xml"

        modal_sandbox.py download sb-abc123 "/app/out:./out" "/app/logs:./logs"
    """
    sandbox = modal.Sandbox.from_id(sandbox_id)

    for path_spec in paths:
        if ":" not in path_spec:
            logger.error(
                "Invalid path format '%s', expected 'remote_path:local_path'", path_spec
            )
            sys.exit(1)

        remote_path, local_path = path_spec.split(":", 1)
        if not remote_path:
            logger.error("Empty remote path in '%s'", path_spec)
            sys.exit(1)
        if not local_path:
            logger.error("Empty local path in '%s'", path_spec)
            sys.exit(1)

        try:
            copy_from_sandbox(sandbox, remote_path, local_path)
        except Exception as e:
            logger.error("Failed to download %s: %s", remote_path, e)
            sys.exit(1)

    logger.info("Download complete")


def stream_output(source, dest):
    """Stream lines from source to dest, flushing after each line."""
    for line in source:
        dest.write(line)
        dest.flush()


@cli.command("exec")
@click.argument("sandbox_id")
@click.argument("command")
def exec_command(sandbox_id: str, command: str):
    """Execute a command on an existing Modal sandbox."""
    sandbox = modal.Sandbox.from_id(sandbox_id)

    # Execute command
    process = sandbox.exec("bash", "-c", command)

    # Stream stdout and stderr concurrently using threads
    stdout_thread = threading.Thread(
        target=stream_output, args=(process.stdout, sys.stdout)
    )
    stderr_thread = threading.Thread(
        target=stream_output, args=(process.stderr, sys.stderr)
    )
    stdout_thread.start()
    stderr_thread.start()
    stdout_thread.join()
    stderr_thread.join()

    process.wait()
    sys.exit(process.returncode)


# App and function for the 'run' subcommand
run_app = modal.App("offload-test")
run_image = modal.Image.debian_slim(python_version="3.11").pip_install("pytest")


@run_app.function(image=run_image, timeout=600)
def _run_test(cmd: str) -> dict:
    import subprocess

    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    # Print output for streaming visibility
    if result.stdout:
        print(result.stdout, end="")
    if result.stderr:
        print(result.stderr, end="", file=sys.stderr)
    return {
        "exit_code": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr,
    }


@cli.command()
@click.argument("command")
def run(command: str):
    """Run a test command on Modal (ephemeral function execution)."""
    with run_app.run():
        result = _run_test.remote(command)

    # Output JSON for Offload to parse
    print(json.dumps(result))
    sys.exit(result["exit_code"])


@cli.command("create")
@click.argument("image_id")
@click.option(
    "--copy-dir",
    "copy_dirs",
    multiple=True,
    help="Copy local dir to sandbox (format: local_path:remote_path)",
)
@click.option(
    "--env",
    "env_vars",
    multiple=True,
    help="Environment variable (format: KEY=VALUE)",
)
@click.option(
    "--cpu",
    type=float,
    default=None,
    help="CPU cores per sandbox",
)
@click.option(
    "--memory-gb",
    "memory_gb",
    type=float,
    default=None,
    help="Memory request per sandbox, in GiB (converted to MiB via "
    "ceil(value * 1024), passed to modal.Sandbox.create(memory=...)). "
    "Modal's default when unset is 128 MiB. Example: --memory-gb 8",
)
@click.option(
    "--experimental-options",
    "experimental_options",
    default=None,
    help="JSON string of experimental options to pass to Sandbox.create()",
)
def create_from_image(
    image_id: str,
    copy_dirs: tuple[str, ...] = (),
    env_vars: tuple[str, ...] = (),
    cpu: float | None = None,
    memory_gb: float | None = None,
    experimental_options: str | None = None,
):
    """Create sandbox using existing image_id.

    IMAGE_ID is the Modal image ID to use.
    """
    t0 = time.time()

    # Log received arguments
    logger.debug("[%.2fs] create_from_image called with:", time.time() - t0)
    logger.debug("[%.2fs]   image_id: %s", time.time() - t0, image_id)
    logger.debug("[%.2fs]   copy_dirs: %s", time.time() - t0, copy_dirs)
    logger.debug("[%.2fs]   env_vars, %d total", time.time() - t0, len(env_vars))

    # Parse environment variables
    env_dict = {}
    for env_spec in env_vars:
        if "=" not in env_spec:
            logger.warning("Invalid env format '%s', expected 'KEY=VALUE'", env_spec)
            continue
        key, value = env_spec.split("=", 1)
        env_dict[key] = value

    app_name = "offload-sandbox"
    app = modal.App.lookup(app_name, create_if_missing=True)

    # Load image from ID and verify it exists
    logger.debug("[%.2fs] Loading image %s...", time.time() - t0, image_id)
    try:
        image = modal.Image.from_id(image_id)
    except Exception as e:
        logger.error("Failed to load image %s: %s", image_id, e)
        logger.error(
            "The image may have been garbage collected. "
            "Try running 'prepare' again to rebuild the image."
        )
        sys.exit(1)
    logger.debug("[%.2fs] Image loaded", time.time() - t0)

    # Create secrets from env dict if any
    secrets = []
    if env_dict:
        secrets = [modal.Secret.from_dict(env_dict)]

    logger.debug("[%.2fs] Creating sandbox...", time.time() - t0)
    try:
        create_kwargs = dict(
            app=app,
            image=image,
            workdir="/app",
            timeout=3600,
            secrets=secrets,
        )
        if cpu is not None:
            create_kwargs["cpu"] = cpu
        if memory_gb is not None:
            create_kwargs["memory"] = math.ceil(memory_gb * 1024)
        if experimental_options is not None:
            exp_opts = json.loads(experimental_options)
            create_kwargs["experimental_options"] = exp_opts
            logger.debug(
                "[%.2fs]   experimental_options: %s",
                time.time() - t0,
                experimental_options,
            )
        sandbox = modal.Sandbox.create(**create_kwargs)
    except Exception as e:
        logger.error("Failed to create sandbox with image %s: %s", image_id, e)
        logger.error(
            "The image may have been garbage collected. "
            "Run 'prepare' again to rebuild."
        )
        sys.exit(1)
    logger.debug("[%.2fs] Sandbox created", time.time() - t0)

    # Copy user-specified directories
    logger.debug(
        "[%.2fs] Processing %d user-specified copy-dir(s)",
        time.time() - t0,
        len(copy_dirs),
    )
    for i, copy_spec in enumerate(copy_dirs):
        logger.info("[%.2fs] copy_dirs[%d]: '%s'", time.time() - t0, i, copy_spec)
        if ":" not in copy_spec:
            logger.warning(
                "Invalid copy-dir format '%s', expected 'local:remote'", copy_spec
            )
            continue
        local_path, remote_path = copy_spec.split(":", 1)
        if not os.path.isdir(local_path):
            logger.warning("Local directory '%s' not found, skipping", local_path)
            continue
        logger.info(
            "[%.2fs] Copying %s to %s...", time.time() - t0, local_path, remote_path
        )
        copy_dir_to_sandbox(sandbox, local_path, remote_path)
        logger.info("[%.2fs] Copy complete", time.time() - t0)

    logger.info("[%.2fs] Sandbox ready: %s", time.time() - t0, sandbox.object_id)
    sys.stdout.write("%s\n" % sandbox.object_id)


if __name__ == "__main__":
    cli()
