#!/usr/bin/env python3

import atexit
import os
import subprocess
import sys
import tempfile

SENSITIVE_NAMES = {".env", ".envrc", "secret.json", "secrets.json"}

SENSITIVE_DIRS = [
    os.path.expanduser("~/.ssh"),
    os.path.expanduser("~/.gnupg"),
]

OVERRIDE_HEADER = """\
services:
  claude:
    volumes:
"""


def find_sensitive_files(workspace):
    """Walk workspace and return paths to sensitive files."""
    matches = []
    for dirpath, _dirnames, filenames in os.walk(workspace):
        for name in filenames:
            if name in SENSITIVE_NAMES:
                matches.append(os.path.join(dirpath, name))
    return matches


def find_sensitive_dirs(workspace):
    """Return SENSITIVE_DIRS entries that exist and fall under workspace."""
    workspace = os.path.realpath(workspace)
    matches = []
    for d in SENSITIVE_DIRS:
        real = os.path.realpath(d)
        if real.startswith(workspace + os.sep) and os.path.isdir(real):
            matches.append(real)
    return matches


def generate_override(workspace, sensitive_files, sensitive_dirs=None):
    """Write a compose override YAML that mounts /dev/null over each file
    and tmpfs over each sensitive directory."""
    fd, path = tempfile.mkstemp(prefix="claude-env-override-", suffix=".yml")
    with os.fdopen(fd, "w") as f:
        f.write(OVERRIDE_HEADER)
        if sensitive_files:
            for host_path in sensitive_files:
                rel = os.path.relpath(host_path, workspace)
                f.write(f"      - /dev/null:/workspace/{rel}:ro\n")
        else:
            f.write("      []\n")
        if sensitive_dirs:
            f.write("    tmpfs:\n")
            for dir_path in sensitive_dirs:
                rel = os.path.relpath(dir_path, workspace)
                f.write(f"      - /workspace/{rel}:ro\n")
    return path


def main():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    workspace = os.getcwd()

    sensitive_files = find_sensitive_files(workspace)
    sensitive_dirs = find_sensitive_dirs(workspace)
    override_file = generate_override(workspace, sensitive_files, sensitive_dirs)
    atexit.register(lambda: os.unlink(override_file))

    compose_file = os.path.join(script_dir, "docker-compose.yml")
    uid = os.getuid()
    gid = os.getgid()

    print(f"++ loading {workspace}")

    cmd = [
        "docker", "compose",
        "-f", compose_file,
        "-f", override_file,
        "run", "--rm",
    ]

    # Allocate TTY only when a human is at the terminal
    if not sys.stdin.isatty():
        cmd.append("-T")

    cmd.extend([
        "-v", f"{workspace}:/workspace",
        "-v", f"{workspace}/.container-abra:/home/claude/.abra",
        "claude",
        *sys.argv[1:],
    ])

    env = {**os.environ, "HOST_UID": str(uid), "HOST_GID": str(gid)}

    result = subprocess.run(cmd, env=env)
    sys.exit(result.returncode)


if __name__ == "__main__":
    main()