#!/usr/bin/env python3
import argparse
import plistlib
import struct
import subprocess
from pathlib import Path

COLORS = {
    "gray":   1,
    "green":  2,
    "purple": 3,
    "blue":   4,
    "yellow": 5,
    "red":    6,
    "orange": 7,
}


def tag_file(path: Path, color: str) -> None:
    idx = COLORS[color]
    new_tag = f"{color.capitalize()}\n{idx}"

    # Read existing tags, then append if not already present
    result = subprocess.run(["xattr", "-px", "com.apple.metadata:_kMDItemUserTags", path], capture_output=True, text=True)
    if result.returncode == 0:
        raw = bytes.fromhex(result.stdout.replace("\n", "").replace(" ", ""))
        existing = plistlib.loads(raw)
    else:
        existing = []

    if new_tag not in existing:
        existing.append(new_tag)

    plist_data = plistlib.dumps(existing, fmt=plistlib.FMT_BINARY)
    subprocess.run(["xattr", "-wx", "com.apple.metadata:_kMDItemUserTags", plist_data.hex(), path], check=True)

    # FinderInfo: color in bits 1-3 of fdFlags (offset 8, big-endian uint16)
    result = subprocess.run(["xattr", "-px", "com.apple.FinderInfo", path], capture_output=True, text=True)
    if result.returncode == 0:
        raw = bytes.fromhex(result.stdout.replace("\n", "").replace(" ", ""))
        fi = bytearray(raw) if len(raw) >= 32 else bytearray(32)
    else:
        fi = bytearray(32)

    flags = struct.unpack_from(">H", fi, 8)[0]
    flags = (flags & ~0x000E) | (idx << 1)
    struct.pack_into(">H", fi, 8, flags)
    subprocess.run(["xattr", "-wx", "com.apple.FinderInfo", bytes(fi).hex(), path], check=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Tag files with Finder color labels")
    parser.add_argument(
        "--tag", required=True, choices=list(COLORS),
        metavar="COLOR", help=f"One of: {', '.join(COLORS)}",
    )
    parser.add_argument("files", nargs="+", type=Path)
    args = parser.parse_args()

    for path in args.files:
        if not path.exists():
            print(f"skipping {path}: not found")
            continue
        tag_file(path, args.tag)
        print(f"tagged {path} → {args.tag}")


if __name__ == "__main__":
    main()