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

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)

# --- interactive prompts ---
input_dir = input("Enter directory containing .m4b files: ").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]

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

# collect .m4b files robustly (avoid glob issues with [])
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")]

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

# sort (simple lexicographic; change if you want natural sort)
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
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 -> select 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)

# final output name
final_name = input("\nEnter final audiobook name (used as filename and Title metadata): ").strip()
if not final_name:
    print("❌ No output name provided.")
    sys.exit(1)
final_filename = os.path.join(input_dir, final_name if final_name.lower().endswith(".m4b") else final_name + ".m4b")

# Extract durations and chapter titles; also capture global tags & cover art from first file
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 {}
    # prefer 'title' tag as chapter name; fall back to filename
    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
        # extract cover art from first file (if exists)
        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 file (global tags + chapters)
meta_path = "chapters_ffmetadata.txt"
with open(meta_path, "w", encoding="utf-8") as mf:
    mf.write(";FFMETADATA1\n")
    # copy global tags except title/name (we'll set our own)
    for k, v in (global_tags or {}).items():
        if k.lower() not in ("title", "name", "©nam"):
            mf.write(f"{k}={v}\n")
    mf.write(f"title={final_name}\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 file list (absolute paths)
file_list = "concat_list.txt"
with open(file_list, "w", encoding="utf-8") as fl:
    for p in selected_paths:
        fl.write(f"file '{os.path.abspath(p)}'\n")

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

# Mux metadata (chapters + global tags [+ cover art]) into final file
cmd = [
    "ffmpeg", "-y", "-hide_banner", "-loglevel", "error",
    "-i", merged, "-i", meta_path
]

if have_cover:
    cmd += ["-i", cover_temp]

cmd += [
    "-map", "0",          # audio
    "-map_metadata", "1", # metadata
    "-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.")
print("Original files were not modified.")

# --- 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}")

# --- optional deletion of source files ---
del_prompt = input("\nDo you want to delete the original files that were merged? [y/N]: ").strip().lower()
if del_prompt == "y":
    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.")