#!/usr/bin/env python3
import os, sys, subprocess, json, argparse

def ffprobe_json(path):
    out = subprocess.check_output([
        "ffprobe", "-v", "quiet", "-print_format", "json",
        "-show_format", "-show_streams", path
    ])
    return json.loads(out)

def run(cmd):
    print("running:", " ".join(cmd))
    subprocess.run(cmd, check=True)

# --- argument parsing ---
parser = argparse.ArgumentParser(
    description="Merge multiple .m4b audiobook files into one with chapters."
)
parser.add_argument(
    "input_dir",
    nargs="?",
    help="Directory containing .m4b files (can drag-and-drop). If omitted, you will be prompted."
)
parser.add_argument(
    "-a", "--all",
    action="store_true",
    help="Automatically select all .m4b files (skips range prompt)."
)
parser.add_argument(
   "-d","--delete",
    action="store_true",
    help="Delete the original files that were merged (no prompt)."
)
parser.add_argument(
    "-o", "--output",
    help="Final output filename (with or without .m4b extension)."
)

args = parser.parse_args()

# --- input dir handling ---
if args.input_dir:
    input_dir = args.input_dir.strip()
else:
    input_dir = input("Enter directory containing .m4b files (or drag folder here): ").strip()

# strip accidental surrounding quotes
if (input_dir.startswith("'") and input_dir.endswith("'")) or \
   (input_dir.startswith('"') and input_dir.endswith('"')):
    input_dir = input_dir[1:-1]

# replace escaped spaces only
input_dir = input_dir.replace(r"\ ", " ")

# normalize
input_dir = os.path.normpath(input_dir)

if not os.path.isdir(input_dir):
    print("❌ Path not found or not a directory:", input_dir)
    sys.exit(1)

# collect .m4b files robustly (skip hidden / ._ files)
all_files = [
    f for f in os.listdir(input_dir)
    if os.path.isfile(os.path.join(input_dir, f))
    and f.lower().endswith(".m4b")
    and not f.startswith(".")  # skip hidden macOS resource files
]

if not all_files:
    print("❌ No .m4b files found in that directory.")
    sys.exit(1)

all_files = sorted(all_files, key=lambda s: s.lower())

print("\nFound the following .m4b files:")
for idx, fname in enumerate(all_files, start=1):
    print(f"[{idx}] {fname}")

# range selection
if args.all:
    start, end = 1, len(all_files)
else:
    rng = input("\nEnter range to merge (e.g. 1-10) or single number (e.g. 5) [Enter = all]: ").strip()
    try:
        if not rng:  # empty input = all
            start, end = 1, len(all_files)
        elif "-" in rng:
            a, b = rng.split("-", 1)
            start, end = int(a), int(b)
        else:
            start = int(rng)
            end = start
        assert 1 <= start <= end <= len(all_files)
    except Exception:
        print("❌ Invalid range. Use start-end where numbers are within the shown list.")
        sys.exit(1)

selected = all_files[start-1:end]
selected_paths = [os.path.join(input_dir, s) for s in selected]
print(f"\nSelected {len(selected)} file(s):")
for s in selected:
    print("  ", s)

# enforce minimum 2 files if --all was used
if args.all and len(selected) < 2:
    print("⚠️ --all requires at least 2 .m4b files to merge.")
    sys.exit(1)
# final output name
if args.output:
    final_name = args.output.strip()
else:
    # derive from first selected file
    first_file = os.path.splitext(selected[0])[0]  # filename w/o extension
    # cut off at " - " if present
    base_name = first_file.split(" - ")[0].strip()
    final_name = base_name

# if not final_name:
#     print("❌ No output name provided.")
#     sys.exit(1)

# normalize extension
if not final_name.lower().endswith(".m4b"):
    final_name += ".m4b"

final_filename = os.path.join(input_dir, final_name)


# Extract durations, chapters, metadata
chapters = []
current_start = 0.0
global_tags = {}
cover_temp = "cover.jpg"
have_cover = False

for idx, p in enumerate(selected_paths):
    meta = ffprobe_json(p)
    dur = float(meta["format"]["duration"])
    tags = meta["format"].get("tags", {}) or {}
    chapter_title = tags.get("title") or tags.get("©nam") or os.path.splitext(os.path.basename(p))[0]
    chapters.append({"start": current_start, "end": current_start+dur, "title": chapter_title})
    current_start += dur
    if idx == 0:
        global_tags = tags
        try:
            run([
                "ffmpeg", "-hide_banner", "-loglevel", "error",
                "-i", p, "-an", "-vcodec", "copy", cover_temp
            ])
            have_cover = True
            print("✅ Extracted cover art from first file.")
        except Exception:
            print("⚠️ No cover art found in first file.")

# Write FFmetadata
meta_path = "chapters_ffmetadata.txt"
with open(meta_path, "w", encoding="utf-8") as mf:
    mf.write(";FFMETADATA1\n")
    for k, v in (global_tags or {}).items():
        if k.lower() not in ("title", "name", "©nam"):
            mf.write(f"{k}={v}\n")
    title_tag = final_name[:-4] if final_name.lower().endswith(".m4b") else final_name #remove .m4b from title
    mf.write(f"title={title_tag}\n\n")
    for ch in chapters:
        mf.write("[CHAPTER]\n")
        mf.write("TIMEBASE=1/1000\n")
        mf.write(f"START={int(ch['start']*1000)}\n")
        mf.write(f"END={int(ch['end']*1000)}\n")
        safe_title = str(ch['title']).replace("\n", " ").replace("\r"," ")
        mf.write(f"title={safe_title}\n\n")

# Build concat list
file_list = os.path.join(os.getcwd(), "concat_list.txt")
with open(file_list, "w", encoding="utf-8") as fl:
    for p in selected_paths:
        safe_path = os.path.abspath(p).replace("'", r"'\''")
        fl.write(f"file '{safe_path}'\n")

# Concatenate audio only
merged = "merged.m4b"
try:
    run([
        "ffmpeg", "-hide_banner", "-loglevel", "error",
        "-f", "concat", "-safe", "0", "-i", file_list,
        "-map", "0:a",
        "-c:a", "copy",
        merged
    ])
except subprocess.CalledProcessError:
    print("❌ ffmpeg failed while concatenating.")
    sys.exit(1)

# Mux metadata + cover
cmd = [
    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
    "-i", merged, "-i", meta_path
]
if have_cover:
    cmd += ["-i", cover_temp]
cmd += [
    "-map", "0",
    "-map_metadata", "1",
    "-c", "copy"
]
if have_cover:
    cmd += ["-map", "2", "-disposition:v:0", "attached_pic"]
cmd += [final_filename]

try:
    run(cmd)
except subprocess.CalledProcessError:
    print("❌ ffmpeg failed while writing metadata.")
    sys.exit(1)

print(f"\n✅ Done. Created '{final_filename}' with {len(chapters)} chapters.")

# cleanup temp files
for f in [merged, meta_path, file_list, cover_temp]:
    if os.path.exists(f):
        try:
            os.remove(f)
        except Exception as e:
            print(f"⚠️ Failed to delete temp file {f}: {e}")

# delete originals?
# if args.delete:
#     delete_confirmed = True
# else:
#     del_prompt = input("\nDo you want to delete the original files that were merged? [y/N]: ").strip().lower()
#     delete_confirmed = del_prompt == "y"

if args.delete:
    for f in selected_paths:
        try:
            os.remove(f)
            print(f"Deleted: {f}")
        except Exception as e:
            print(f"⚠️ Could not delete {f}: {e}")
else:
    print("Original files kept.")