headquarter/apps/api/src/services/build/manifest_compiler.py

"""Manifest compiler: transforms ToolDefinitionManifest into Dockerfile + Compose."""

import hashlib
import json
import shlex
from copy import deepcopy
from typing import Any

import yaml

from src.services.config.config_profile_resolver import expand_container_path
from src.services.docker import sort_volumes_by_specificity


def resolve_base(manifest: dict) -> dict:
    """Merge a base definition into a tool manifest.

    If the manifest has base_definition_id, the base manifest is loaded
    and merged. Tool-specific values override base values.

    Args:
        manifest: The tool manifest JSON (may reference a base)

    Returns:
        A fully resolved manifest with base values merged in.
    """
    result = deepcopy(manifest)

    base_definition_id = result.pop("base_definition_id", None)
    result.pop("base_version", None)

    if base_definition_id:
        # This will be provided by the caller (they have the DB session)
        # For now, we assume the manifest has been pre-resolved
        # or the caller provides the base manifest separately.
        pass

    return result


def deep_merge(base: dict, override: dict) -> dict:
    """Deep merge two manifests. Arrays are concatenated; dicts are merged.

    Args:
        base: The base manifest.
        override: The tool-specific overrides.

    Returns:
        Merged manifest.
    """
    merged = deepcopy(base)

    for key, value in override.items():
        if key == "mounts" and isinstance(value, list):
            # Concatenate mount arrays
            existing = merged.get("mounts", [])
            merged["mounts"] = existing + deepcopy(value)
        elif key == "scripts" and isinstance(value, dict):
            # Merge script categories
            if "scripts" not in merged:
                merged["scripts"] = {}
            for script_key, script_value in value.items():
                existing = merged["scripts"].get(script_key, [])
                merged["scripts"][script_key] = existing + deepcopy(script_value)
        elif key == "packages" and isinstance(value, dict):
            # Union package arrays
            if "packages" not in merged:
                merged["packages"] = {}
            for pkg_key, pkg_value in value.items():
                if (
                    pkg_key in merged["packages"]
                    and isinstance(merged["packages"][pkg_key], list)
                    and isinstance(pkg_value, list)
                ):
                    merged["packages"][pkg_key] = merged["packages"][
                        pkg_key
                    ] + deepcopy(pkg_value)
                else:
                    merged["packages"][pkg_key] = deepcopy(pkg_value)
        elif key == "env" and isinstance(value, dict):
            # Dict merge: override wins on key conflict
            if "env" not in merged:
                merged["env"] = {}
            merged["env"].update(deepcopy(value))
        elif (
            isinstance(value, dict) and key in merged and isinstance(merged[key], dict)
        ):
            # Generic dict merge
            merged[key] = {**merged[key], **deepcopy(value)}
        else:
            # Override entirely
            merged[key] = deepcopy(value)

    return merged


def compile_dockerfile(manifest: dict) -> str:
    """Compile a resolved manifest into a Dockerfile string.

    Args:
        manifest: Fully resolved manifest JSON.

    Returns:
        Dockerfile content.
    """
    lines: list[str] = []

    # FROM
    base_image = manifest.get("base_image", "ubuntu:24.04")
    lines.append(f"FROM {base_image}")
    lines.append("")

    # Build-time environment
    env = manifest.get("env", {})
    for key, value in env.items():
        lines.append(f"ENV {key}={shlex.quote(value)}")
    if env:
        lines.append("")

    # System packages (apt)
    apt_packages = manifest.get("packages", {}).get("apt", [])
    if manifest.get("user"):
        # Ensure sudo is available for permission-fixing startup scripts
        apt_packages = list(apt_packages)
        if "sudo" not in apt_packages:
            apt_packages.append("sudo")
    if apt_packages:
        lines.append("RUN apt-get update && apt-get install -y \\")
        for pkg in apt_packages[:-1]:
            lines.append(f"    {pkg} \\")
        lines.append(f"    {apt_packages[-1]} \\")
        lines.append("    && rm -rf /var/lib/apt/lists/*")
        lines.append("")

    # Node.js
    node = manifest.get("packages", {}).get("node")
    if node:
        version = node.get("version", "20")
        lines.append(
            f"RUN curl -fsSL https://deb.nodesource.com/setup_{version}.x | bash - && \\"
        )
        lines.append("    apt-get install -y nodejs && \\")
        lines.append("    rm -rf /var/lib/apt/lists/*")
        lines.append("")

    # NPM global packages
    npm_packages = manifest.get("packages", {}).get("npm_global", [])
    if npm_packages:
        pkg_list = " ".join(shlex.quote(p) for p in npm_packages)
        lines.append(f"RUN npm install -g {pkg_list}")
        lines.append("")

    # Pip packages
    pip_packages = manifest.get("packages", {}).get("pip", [])
    if pip_packages:
        pkg_list = " ".join(shlex.quote(p) for p in pip_packages)
        lines.append(f"RUN pip install {pkg_list}")
        lines.append("")

    # User creation
    user = manifest.get("user")
    home_dir = get_manifest_home_dir(manifest)
    workspace_name = manifest.get("workspace_name", "{{WORKSPACE_NAME}}")
    if user:
        name = user["name"]
        uid = user["uid"]
        gid = user["gid"]
        create_home = "-m " if user.get("create_home", True) else ""
        shell = user.get("shell", "/bin/bash")
        lines.append(f"RUN groupadd -g {gid} {name} && \\")
        lines.append(f"    useradd -u {uid} -g {gid} {create_home}-s {shell} {name}")
        lines.append("")
        # Set HOME and USER for runtime compatibility
        lines.append(f"ENV HOME={home_dir}")
        lines.append(f"ENV USER={name}")
        lines.append("")
        # Ensure home directory exists and is writable by the user.
        # Recursively chown so any files copied from /etc/skel by useradd -m
        # (e.g. .bashrc, .config) are owned by the container user.
        lines.append(
            f"RUN mkdir -p {home_dir} && chown -R {name}:{name} {home_dir} && chmod 755 {home_dir}"
        )
        # Pre-create common config directories so apps like ranger can write
        # their configs on first run without permission errors.
        common_dirs = [".config", ".local/share", ".cache"]
        for d in common_dirs:
            lines.append(
                f"RUN mkdir -p {home_dir}/{d} && chown -R {name}:{name} {home_dir}/{d}"
            )
        lines.append("")

        # Configure passwordless sudo so startup scripts can fix permissions
        lines.append(
            f'RUN echo "{name} ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/{name} && chmod 0440 /etc/sudoers.d/{name}'
        )
        lines.append("")

    # Build scripts
    build_scripts = manifest.get("scripts", {}).get("build", [])
    for script in build_scripts:
        # Normalize multi-line scripts into single RUN command
        stripped_lines = [
            line.strip() for line in script.strip().split("\n") if line.strip()
        ]
        if stripped_lines:
            normalized = " && ".join(stripped_lines)
            lines.append(f"RUN {normalized}")
    if build_scripts:
        lines.append("")

    # After build scripts, ensure everything in home is owned by the user
    if user and build_scripts:
        lines.append(f"RUN chown -R {name}:{name} {home_dir}")
        lines.append("")

    # Create mount target directories and /workspace compatibility symlink.
    # The symlink target includes the workspace/repo name so legacy scripts
    # that cd into /workspace still land on the right project.
    mounts = manifest.get("mounts", [])
    if mounts:
        dirs = [mount["target"] for mount in mounts]
        dir_str = " ".join(dirs)
        lines.append(f"RUN mkdir -p {dir_str}")
        if user:
            lines.append(f"RUN chown -R {user['name']}:{user['name']} {dir_str}")
        lines.append("")

    workspace_target = f"{home_dir}/{workspace_name}"
    lines.append(f"RUN mkdir -p {workspace_target}")
    if user:
        lines.append(
            f"RUN ln -sfn {workspace_target} /workspace && chown -R {user['name']}:{user['name']} {home_dir}"
        )
    else:
        lines.append(f"RUN ln -sfn {workspace_target} /workspace")
    lines.append("")

    # Entrypoint for startup scripts
    startup_scripts = manifest.get("scripts", {}).get("startup", [])
    if startup_scripts:
        lines.append(
            "COPY .headquarter/entrypoint.sh /usr/local/bin/headquarter-entrypoint"
        )
        lines.append("RUN chmod +x /usr/local/bin/headquarter-entrypoint")
        lines.append("")

    # Switch to runtime user
    if user:
        lines.append(f"USER {user['name']}")

    # Set WORKDIR to the configured home directory unless runtime.working_dir
    # explicitly overrides it.
    runtime = manifest.get("runtime", {})
    working_dir = runtime.get("working_dir")
    if working_dir:
        lines.append(f"WORKDIR {expand_container_path(working_dir, home_dir)}")
    else:
        lines.append(f"WORKDIR {home_dir}")
    lines.append("")

    # Entrypoint and CMD
    if startup_scripts:
        lines.append('ENTRYPOINT ["/usr/local/bin/headquarter-entrypoint"]')

    command = runtime.get("command", ["/bin/bash"])
    cmd_json = json.dumps(command)
    lines.append(f"CMD {cmd_json}")

    return "\n".join(lines)


def compile_entrypoint(manifest: dict) -> str:
    """Generate the startup entrypoint script from startup scripts.

    Injects a permission-fixer preamble that runs as root (or via sudo) before
    any user-defined startup script. It chowns the home directory and a safe
    subset of mount parents to the container user, creates the /workspace
    compatibility symlink, and avoids recursive chown of large repo subtrees.

    Args:
        manifest: Fully resolved manifest JSON.

    Returns:
        Shell script content.
    """
    lines = ["#!/bin/bash", "set -e", ""]

    user = manifest.get("user")
    home_dir = get_manifest_home_dir(manifest)
    workspace_name = manifest.get("workspace_name", "{{WORKSPACE_NAME}}")
    workspace_target = f"{home_dir}/{workspace_name}"

    # Permission fixer preamble: run as root when possible, else fall back to
    # passwordless sudo configured in the Dockerfile.
    lines.append("# Permission fixer preamble")
    lines.append("CONTAINER_USER=''")
    lines.append('if [ "$(id -u)" = '"'"'0'"'"' ]; then')
    if user:
        lines.append(f"    CONTAINER_USER='{user['name']}'")
    lines.append("else")
    lines.append("    # Try passwordless sudo; ignore failure so the container still starts")
    lines.append("    if sudo -n true 2>/dev/null; then")
    lines.append("        SUDO='sudo'")
    lines.append("    else")
    lines.append("        SUDO=''")
    lines.append("    fi")
    lines.append("fi")
    lines.append("")

    if user:
        name = user["name"]
        uid = user["uid"]
        gid = user["gid"]
        lines.append(f"USER_NAME='{name}'")
        lines.append(f"USER_UID='{uid}'")
        lines.append(f"USER_GID='{gid}'")
        lines.append(f"HOME_DIR='{home_dir}'")
        lines.append(f"WORKSPACE_TARGET='{workspace_target}'")
        lines.append("")
        lines.append("fix_owner() {")
        lines.append("    local path=\"$1\"")
        lines.append('    [ -e "$path" ] || return 0')
        lines.append('    if [ -n "$SUDO" ]; then')
        lines.append('        sudo chown "$USER_UID:$USER_GID" "$path" 2>/dev/null || true')
        lines.append('    elif [ "$(id -u)" = "0" ]; then')
        lines.append('        chown "$USER_UID:$USER_GID" "$path" 2>/dev/null || true')
        lines.append('    fi')
        lines.append("}")
        lines.append("")
        lines.append("# Ensure home directory exists and is owned by the container user")
        lines.append('mkdir -p "$HOME_DIR"')
        lines.append('fix_owner "$HOME_DIR"')
        lines.append("")
        lines.append("# Ensure workspace target exists and is owned by the container user")
        lines.append('mkdir -p "$WORKSPACE_TARGET"')
        lines.append('fix_owner "$WORKSPACE_TARGET"')
        lines.append("")
        lines.append("# Create /workspace compatibility symlink")
        lines.append('ln -sfn "$WORKSPACE_TARGET" /workspace')
        lines.append("")
        lines.append("# Fix ownership of declared mount targets (top-level only)")
        for mount in manifest.get("mounts", []):
            target = mount.get("target")
            if not target:
                continue
            # Expand any ~/$HOME placeholders in the mount target.
            expanded = target.replace("~", home_dir).replace("$HOME", home_dir)
            if expanded.startswith(home_dir) and not mount.get("readonly", False):
                lines.append(f'fix_owner "{expanded}"')
        lines.append("")

    startup_scripts = manifest.get("scripts", {}).get("startup", [])
    for script in startup_scripts:
        lines.append(script)
        lines.append("")

    lines.append('exec "$@"')
    return "\n".join(lines)


def compile_compose(manifest: dict, variables: dict[str, Any]) -> str:
    """Compile a resolved manifest into a Docker Compose string.

    Args:
        manifest: Fully resolved manifest JSON.
        variables: Resolved values: IMAGE_TAG, INSTANCE_NAME, REPO_PATH, etc.

    Returns:
        Docker Compose YAML content.
    """
    runtime = manifest.get("runtime", {})
    user = manifest.get("user")
    interface_type = manifest["interface_type"]
    home_dir = get_manifest_home_dir(manifest)

    # Determine the workspace/repo name from variables when available.
    workspace_name = variables.get(
        "WORKSPACE_NAME",
        variables.get("REPO_NAME", "workspace"),
    )

    service: dict[str, Any] = {
        "image": variables["IMAGE_TAG"],
        "container_name": variables["INSTANCE_NAME"],
        "restart": "unless-stopped",
    }

    # Terminal-specific fields
    if runtime.get("stdin_open", False):
        service["stdin_open"] = True
    if runtime.get("tty", False):
        service["tty"] = True
    if runtime.get("working_dir"):
        service["working_dir"] = expand_container_path(
            runtime["working_dir"], home_dir
        )

    # User override
    if user:
        service["user"] = f"{user['uid']}:{user['gid']}"

    # Ports for web tools
    default_port = manifest.get("default_port")
    if interface_type == "web" and default_port:
        service["ports"] = [f"{variables['TOOL_PORT']}:{default_port}"]

    # Environment
    env = manifest.get("env", {})
    if env:
        service["environment"] = dict(env)

    # Merge extra env from config
    extra_env = variables.get("EXTRA_ENV", {})
    if extra_env:
        if "environment" not in service:
            service["environment"] = {}
        service["environment"].update(extra_env)

    # Volumes from mount schema
    volumes = []
    has_explicit_repo_mount = False
    for mount in manifest.get("mounts", []):
        source = resolve_mount_source(mount, variables)
        if not source:
            continue
        if mount.get("source_type") == "repo":
            has_explicit_repo_mount = True
        target = expand_container_path(mount["target"], home_dir)
        readonly = ":ro" if mount.get("readonly", False) else ""
        volumes.append(f"{source}:{target}{readonly}")

    # Synthesize a default repo/workspace mount when the manifest does not
    # declare an explicit repo mount. This preserves the repo root directory
    # name under the configured home directory.
    if not has_explicit_repo_mount and variables.get("REPO_PATH"):
        target = f"{home_dir}/{workspace_name}"
        volumes.append(f"{variables['REPO_PATH']}:{target}")

    # Append extra volumes from tool config / config profile
    for vol in variables.get("EXTRA_VOLUMES", []):
        vol_str = f"{vol['source']}:{vol['target']}"
        if vol.get("readonly"):
            vol_str += ":ro"
        volumes.append(vol_str)

    if volumes:
        service["volumes"] = sort_volumes_by_specificity(volumes)

    compose = {"services": {"app": service}}
    result = yaml.dump(compose, default_flow_style=False)

    # Debug: log mount resolution so we can diagnose missing mounts
    import logging
    logger = logging.getLogger(__name__)
    logger.debug(
        "compile_compose: REPO_PATH=%s SSH_PATH=%s EXTRA_VOLUMES=%s mounts=%s volumes=%s",
        variables.get("REPO_PATH", "<empty>"),
        variables.get("SSH_PATH", "<empty>"),
        variables.get("EXTRA_VOLUMES", []),
        manifest.get("mounts", []),
        volumes,
    )

    return result


def resolve_mount_source(mount: dict, variables: dict[str, Any]) -> str:
    """Resolve a mount's source_type to an actual host path.

    Args:
        mount: Mount definition from manifest.
        variables: Resolved variables dict.

    Returns:
        Host path string, or empty string if unresolved.
    """
    source_type = mount.get("source_type", "host_path")

    if source_type == "repo":
        return variables.get("REPO_PATH", "")
    elif source_type == "ssh_key":
        return variables.get("SSH_PATH", "")
    elif source_type == "instance":
        instance_dir = variables.get("INSTANCE_DIR", "")
        mount_name = mount.get("name", "unknown")
        return f"{instance_dir}/mounts/{mount_name}"
    elif source_type == "git_mount":
        ref = mount.get("git_mount_ref", "default")
        return variables.get(f"GIT_MOUNT_{ref}", "")
    elif source_type == "host_path":
        return mount.get("source", "")

    return ""


def get_manifest_home_dir(manifest: dict) -> str:
    """Get the home directory for a container based on manifest config.

    Precedence:
      1. manifest["home_directory"] if present and non-empty.
      2. /home/{user.name} if manifest.user.name is present.
      3. /root otherwise.

    Args:
        manifest: Fully resolved manifest JSON.

    Returns:
        Home directory path (e.g., /home/user or /root).
    """
    home_directory = manifest.get("home_directory")
    if home_directory and isinstance(home_directory, str) and home_directory.strip():
        return home_directory.strip()

    user = manifest.get("user")
    if user and user.get("name"):
        return f"/home/{user['name']}"
    return "/root"


def compute_image_tag(tool_name: str, manifest: dict) -> str:
    """Compute a deterministic image tag from manifest content.

    Args:
        tool_name: Human-readable tool name.
        manifest: Fully resolved manifest JSON.

    Returns:
        Docker image tag string.
    """
    # Canonicalize: sort keys, stable JSON
    canonical = json.dumps(manifest, sort_keys=True, separators=(",", ":"))
    hash_suffix = hashlib.sha256(canonical.encode()).hexdigest()[:8]
    safe_name = tool_name.lower().replace(" ", "-").replace("_", "-")
    return f"headquarter/{safe_name}-{hash_suffix}:latest"


def merge_with_config(manifest: dict, profile: dict | None = None) -> dict:
    """Merge ConfigProfile overrides into a manifest.

    Args:
        manifest: Base manifest from tool definition.
        profile: Resolved ConfigProfile (optional).

    Returns:
        Manifest with overrides applied.
    """
    result = deepcopy(manifest)

    extra_env: dict[str, str] = {}
    extra_volumes: list[dict] = []

    # Apply ConfigProfile
    if profile:
        if profile.get("environment_variables"):
            extra_env.update(profile["environment_variables"])
        if profile.get("mounts"):
            extra_volumes.extend(profile["mounts"])
        # Profile hints override everything
        hints = profile.get("hints", {})
        if hints.get("start_command"):
            result["runtime"] = result.get("runtime", {})
            result["runtime"]["command"] = hints["start_command"].split()
        if hints.get("working_directory"):
            result["runtime"] = result.get("runtime", {})
            result["runtime"]["working_dir"] = hints["working_directory"]
        if hints.get("port_override"):
            result["default_port"] = hints["port_override"]

    # Store merged extras for the compose compiler
    result["_extra_env"] = extra_env
    result["_extra_volumes"] = extra_volumes

    return result