#!/usr/bin/env python3
# BCAO - BandCamp Automatic Organiser
# copyright 2018-2019 @LynnearSoftware@fedi.lynnesbian.space
# Licensed under the GPLv3: https://www.gnu.org/licenses/gpl-3.0.html#content
# input: a .zip from bandcamp
# output: it organises it, adds cover art, puts it in the right place...

import subprocess, argparse, sys, os, re, base64

def log(message: str, importance: int = 0):
	if not args.quiet or importance > 0:
		print(message)

def die(message: str, code: int = 1):
	log(message)
	exit(code)

try:
	from mutagen.oggvorbis import OggVorbis
	from mutagen.mp3 import MP3
	from mutagen.flac import FLAC
	from mutagen.flac import Picture
	from mutagen.aac import AAC
except ImportError:
	die("Please install python3-mutagen (pip install mutagen)")

try:
	subprocess.check_output(["unzip", "--help"])
except:
	die("Please install unzip (apt install unzip)")

try:
	subprocess.check_output(["which", "convert"])
	subprocess.check_output(["which", "identify"])
except:
	die("Please install imagemagick, and ensure convert and identify are in your $PATH (apt install imagemagick")

parser = argparse.ArgumentParser(description="Extracts the given zip file downloaded from Bandcamp and organises it.")
parser.add_argument('zip', help='The zip file to use')
parser.add_argument('-d', '--destination', dest='destination', default='/home/lynne/Music/Music/', help="The directory to organise the music into. Default: /home/lynne/Music/Music/")
parser.add_argument('-q', '--quiet', dest='quiet', action='store_true', help='Disable non-error output and assume default artist name.')
parser.add_argument('-t', '--threshold', dest='threshold', nargs=1, default=300, help="Maximum acceptable cover art file size in kilobytes. Default: 300")

args = parser.parse_args()

if not os.path.exists(args.zip):
	die(f"Couldn't find {args.zip}.", 2)

log("Extracting...")
zipname = os.path.splitext(os.path.basename(args.zip))[0]
tmp = "/tmp/bcao/{0}".format(zipname)
subprocess.check_output(["rm", "-rf", tmp])
subprocess.check_output(['mkdir', '-p', tmp])
subprocess.check_output(["unzip", args.zip, "-d", tmp])
files = []
songs = []

for root, dirs, filez in os.walk(tmp):
	for file in filez: # for every file
		files.append(root + os.sep + file)

cover = ""
artists = []
album = ""
fileExtensionRegex = re.compile(r"[^.]+$")
probablyASongRegex = re.compile(r"^.+ - .+ - \d{2,} .+\.[^.]+$") # matches "artist - album - 01 track.xyz" but not "some weird bonus thing.mp3"
musicExts = ["ogg", "flac", "alac", "aiff", "wav", "mp3", "opus", "m4a", "aac", "oga"]
bannedCharacters = ["?", "\\", "/", ":", "|", "*", "\"", "<", ">"] # characters disallowed in NTFS filenames
trackCount = 0

log("Processing... please wait.")
# use "01 Song.ogg" instead of "1 Song.ogg". Also works for albums with more than 99 tracks if that ever happens
trackNumberLength = len(str(len(files)))
if trackNumberLength < 2:
	trackNumberLength = 2

for file in files:
	if os.path.basename(file).lower() in ["cover.png", "cover.jpg", "cover.gif"]:
		# we've found our cover art
		cover = file
	ext = re.search(fileExtensionRegex, file.lower()).group(0)
	if ext in musicExts:
		if re.search(probablyASongRegex, file) != None:
			# this is definitely a song and probably not a bonus sound file or whatever
			if ext == "ogg":
				f = OggVorbis(file)
				name = "{0} {1}.{2}".format(f["TRACKNUMBER"][0].zfill(trackNumberLength), f["TITLE"][0], ext)
				for bc in bannedCharacters:
					if bc in name:
						name = name.replace(bc, "-") # replace banned characters with dashes
				songs.append(name)
				if album == "":
					album = f["ALBUM"][0]
				if f["ARTIST"][0] not in artists:
					artists.append(f["ARTIST"][0])
				trackCount = trackCount + 1
				subprocess.check_output(["mv", file, os.path.dirname(file) + os.sep + name])

			else:
				log("UNSUPPORTED FORMAT BECAUSE LYNNE IS A LAZY MORON")
if len(artists) > 1:
	artists.append("Various Artists")

if cover == "":
	# TODO: HANDLE THIS PROPERLY
	die("couldn't find the cover art :)))))")

while os.path.getsize(cover) / 1024 > args.threshold:
	if os.path.basename(cover) != "cover-lq.jpg":
		nucover = os.path.dirname(cover) + os.sep + "cover-lq.jpg"
		subprocess.check_output(["convert", cover, "-quality", "85", "-strip", nucover]) # convert the file to a jpeg
		cover = nucover
	else:
		subprocess.check_output(["convert", cover, "-resize", "90%", cover]) # shrink it slightly

with open(cover, "rb") as cvr:
	data = cvr.read()

imginfo = subprocess.check_output(["identify", "-ping", "-format", r"%w,%h,%m,%[bit-depth]", cover]).decode("utf-8").split(",")

picture = Picture()
picture.data = data
picture.type = 3
picture.mime = "image/{0}".format(imginfo[2].lower())
picture.width = int(imginfo[0])
picture.height = int(imginfo[1])
picture.depth = int(imginfo[3])

picture_data = picture.write()
encoded_data = base64.b64encode(picture_data)
vcomment_value = encoded_data.decode("ascii")

for song in songs:
	f = OggVorbis(tmp + os.sep + song)
	f["metadata_block_picture"] = [vcomment_value]
	f.save()

artist = artists[0]
artists.append("Custom...")

# log("Please choose the artist name to use when creating the folder.")
choice = 0
while True:
	log("Artist directory:")
	for i in range(len(artists)):
		log("{0}) {1}".format(i + 1, artists[i]))
	choice = 1 if args.quiet else input("> ")
	try:
		choice = artists[int(choice) - 1]
	except KeyError:
		log(f"Please choose an option from 1 to {len(artists)}.")
		continue
	break
if choice == "Custom...":
	log("Enter the name to use.")
	choice = input("> ")
# log("Setting artist to {0}".format(choice))
artist = choice

mPath = os.path.join(args.destination, artist, album)
subprocess.check_output(["mkdir", "-p", mPath])

for root, dirs, filez in os.walk(tmp):
	for file in filez: # for every file
		subprocess.check_output(["mv", root + os.sep + file, mPath + os.sep + file])

log("Deleting {}...".format(tmp))
subprocess.check_output(["rm", "-rf", tmp])

log("Done!")