mutate

#!/usr/bin/env python3

# Setup:
#
# pip install mutagen
# pip install pandas
# pip install prompt_toolkit

from mutagen import File, Metadata
from mutagen.mp3 import MP3
from mutagen.id3 import ID3, ID3TimeStamp
from mutagen.mp4 import MP4, MP4Tags
from mutagen.flac import FLAC
from mutagen._vorbis import VCommentDict

from prompt_toolkit import PromptSession, prompt
from prompt_toolkit.completion import NestedCompleter, FuzzyWordCompleter
from prompt_toolkit.auto_suggest import AutoSuggest, Suggestion
from prompt_toolkit.shortcuts import clear

import pandas as pd
from pandas.api.types import is_numeric_dtype

import os, os.path as path
from functools import partial
import re
import sys
import tempfile
import shutil as sh
from importlib import import_module

f_num          = 'num'
f_artist       = 'artist'
f_title        = 'title'
f_album_artist = 'album_artist'
f_album        = 'album'
f_year         = 'year'
f_comment      = 'comment'
f_genre        = 'genre'
f_disc         = 'disc'

fields = [f_num, f_artist, f_title, f_album_artist, f_album, f_year, f_comment, f_genre, f_disc]

f_filename = 'filename'
f_file     = 'file'
f_changed  = 'changed'

dataCols = [ f_num         , f_artist , f_title , f_album_artist , f_album , f_year , f_comment     , f_genre , f_disc       ]
dataExtr = {
  MP3:     [ 'TRCK'        , 'TPE1'   , 'TIT2'  , 'TPE2'         , 'TALB'  , 'TDRC' , 'COMM'        , 'TCON'  , 'TPOS'       ],
  MP4:     [ 'trkn'        , '©ART'   , '©nam'  , 'aART'         , '©alb'  , '©day' , '©cmt'        , '©gen'  , 'disk'       ],
  FLAC:    [ 'TRACKNUMBER' , 'ARTIST' , 'TITLE' , 'ALBUMARTIST'  , 'ALBUM' , 'DATE' , 'DESCRIPTION' , 'GENRE' , 'DISCNUMBER' ]
}

# TODO apply title_case <field>
# TODO help
cmd_print  = "print"
cmd_tags   = "tags"
cmd_remove = "remove"
cmd_edit   = "edit"
cmd_write  = "write"
cmd_quit   = "quit"

def completer():
  cmds = [cmd_print, cmd_tags, cmd_remove, cmd_edit, cmd_write, cmd_quit]
  d = dict.fromkeys(cmds, None)
  d[cmd_edit] = d['e'] = FuzzyWordCompleter(fields + [f_filename])
  d[cmd_remove] = d['rm'] = d['r'] = FuzzyWordCompleter(fields)
  return NestedCompleter.from_nested_dict(d)

class CommandSuggest(AutoSuggest):
  edit_remove_re = re.compile(f"({cmd_edit}|e|{cmd_remove}|rm|r) ")
  def get_suggestion(self, buffer, document):
    if re.fullmatch(self.edit_remove_re, document.text):
      return Suggestion('0 1-3 4,5 album_artist comment')

s = PromptSession(completer=completer(), auto_suggest=CommandSuggest())

def write(df):
  cols = data_columns(df)

  def writeOne(row):
    _, filepath_temp = tempfile.mkstemp()
    try:
      filepath_orig = row[f_file]
      sh.copy2(filepath_orig, filepath_temp)
      ext = path.splitext(filepath_orig)[1].lower()
      f = File(filepath_temp)
      print("\nBEFORE:\n" + f.pprint())
      ft = type(f)
      tags = None
      if ft == MP3:
        tags = ID3()
      elif ft == MP4:
        tags = MP4Tags()
      elif ft == FLAC:
        f.delete()
        f = File(filepath_temp)
        f.add_tags()
        tags = f.tags
      de = dataExtr[ft]

      abort = False
      for c in cols:
        tag = de[dataCols.index(c)]
        val = row[c]
        if pd.isna(val):
          continue
        try:
          if ft == MP3:
            Tag = getattr(import_module("mutagen.id3"), tag)
            if c == f_year:
              val = [ID3TimeStamp(str(val))]
            else:
              val = [val]
            tags.setall(tag, [Tag(3, text=val)])
          elif ft == MP4 and (c in [f_num, f_disc]):
            tags[tag] = [(val, 0)]
          else:
            tags[tag] = str(val)
        except:
          print("Error while setting tag '%s' to '%s': %s" % (tag, val, sys.exc_info()[1]))
          abort = True

      if not abort:
        if (ft != FLAC):
          f.tags = tags
        print("\nAFTER:\n" + f.pprint())
        f.save(filepath_temp)

        fname = row[f_filename]
        if filename(filepath_orig) == fname:
          sh.move(filepath_temp, filepath_orig)
        else:
          d = path.dirname(filepath_orig)
          d = (d + '/') if d else ""
          new_filepath = d + fname + ext
          sh.move(filepath_temp, new_filepath)
          print("\nRENAMED: %s -> %s" % (filepath_orig, new_filepath))
          os.remove(filepath_orig)
    finally:
      try:
        os.remove(filepath_temp)
      except:
        pass

  for _, row in df[df.changed].iterrows():
    writeOne(row)

  print("\nDone!")

def filename(filepath):
  return path.splitext(path.basename(filepath))[0]

def get_raw(file, col):
  tpe = type(file)
  key = dataExtr[tpe][dataCols.index(col)]
  val = file.get(key)
  # print(col, key, val)
  if not val and tpe == MP3 and col == f_comment:
    val = file.tags.getall(key)
    val = val[0].text[0] if val and val[0].text else None
  elif val:
    if tpe == MP4 or tpe == FLAC:
      if isinstance(val, list):
        val = val[0]
      # Yep, it should be done twice
      if isinstance(val, tuple):
        val = val[0]
  return str(val).strip() if val else None

def artist_title(s):
  s = s.replace('_', ' ') if s and s.count('_') and s.count('-') and not s.count(' ') else s
  arr = s.split(" - ") if s else None
  return arr if arr and len(arr) == 2 else (None, None)

non_alpha_re = r'[^a-zA-Z]+'
# TODO add "club mix"
tc_lower = ['a', 'an', 'and', 'by', 'for', 'of', 'or', 'to', 'the', 'remix',
  'rework', 'extended', 'version', 'mix']
tc_upper = ['DJ']
def apply_title_case(str):
  def f(x):
    if re.sub(non_alpha_re, '', x).lower() in tc_lower:
      return x.lower()
    # FIXME won't work for "^Dj"
    elif re.sub(non_alpha_re, '', x).upper() in tc_upper:
      return x.upper()
    else:
      return x
  arr = str.split(' ')
  arr = [arr[0]] + list(map(f, arr[1:]))
  str = ' '.join(arr)
  return str

def post_process(val, col, file, filepath):
  if col == f_comment:
    # Skip first 2 lines
    #val = '\n'.join(val.splitlines()[2:])
    None
  elif col == f_artist:
    arr = [
      lambda: val,
      lambda: artist_title(get_raw(file, f_title))[0],
      lambda: artist_title(filename(filepath))[0]
    ]
    val = next(filter(None, map(lambda x: x(), arr)))
    val = apply_title_case(val)
    # TODO fix "vs" & "feat"
  elif col == f_title:
    val = val if val else artist_title(filename(filepath))[1]
    if val:
      # Remove artist
      val = val.split(" - ")[-1]
      val = val.replace('-', ' ')
      arr = val.split(' ')
      val = apply_title_case(val)
      # Remove "(original mix)"
      val = re.sub(" ?\\(original mix\\)", "", val, flags=re.IGNORECASE)
  elif col == f_album and val:
    # Remove date
    val = re.sub(r"\(?(\d{4}-..-..|\d{4})\)?", "", val)
    val = apply_title_case(val)
  elif col in [f_num, f_disc] and val:
    val = re.sub(r"(\d+)/?.*", r"\1", val)
    val = int(val) if val else None
  elif col == f_year and val:
    val = re.sub(r".*(\d{4}).*", r"\1", val)
    val = int(val) if val else None
  return val

def get(filepath, file, col):
  try:
    val = get_raw(file, col)
    val = post_process(val, col, file, filepath)
    # print(col, val)
    return val
  except:
    print("Error while processing '%s': %s" % (col, repr(sys.exc_info()[1])))

def remove(what, df):
  what = [w for token in what for w in token.split(',')]
  idxs = []
  cols = []
  for token in what:
    if token.isdigit():
      idxs.append(int(token))
    elif '-' in token:
      t = token.split('-')
      idxs += range(int(t[0]), int(t[1]) + 1)
    else:
      cols.append(token)

  for c in cols:
    query = df[c].notnull()
    if idxs:
      query &= idxs
    df.loc[query, f_changed] = True
  # TODO if both idxs & cols specified, then only cells should be emptied
  df.drop(index=idxs, columns=cols, inplace=True)

def edit(df, idx, col):
  try:
    if idx == None:
      if col == f_filename:
        single_artist = df[f_artist].nunique() == 1
        val = prompt(
          "New template for 'filename': ",
          default = "%%n %s%%t" % ('' if single_artist else "%a - ")
        )
        zero_pad = None
        if '%n' in val and f_num in df.columns:
          zero_pad = str(df[f_num].map(lambda x: len(str(x))).max())
          zero_pad = '%0' + zero_pad + 'd'
        def f(r):
          result = val
          if '%n' in val:
            result = result.replace('%n', (zero_pad % (r[f_num])) if r[f_num] else "None")
          if '%a' in val:
            result = result.replace('%a', str(r[f_artist]))
          if '%t' in val:
            result = result.replace('%t', str(r[f_title]))
          return result
        updated = df.apply(f, axis=1)
        df.loc[df[col] != updated, f_changed] = True
        df[col] = updated
      else:
        default = df[col].mode()
        multiline = False
        expected_size = 1
        if df[col].nunique() > 1:
          expected_size = df[col].size
          default = df[col].str.cat(sep='\n')
          multiline = True
        else:
          if default.empty and col == f_album_artist:
            default = df[f_artist].mode()
          default = '' if default.empty else default[0]
        val = prompt(
          "\nNew values for '%s'%s" % (
            col,
            ' (Alt+Enter):\n\n' if multiline else ': '
          ),
          default=default,
          multiline=multiline,
        )
        if multiline and '\n' in val:
          val = pd.Series(val.split('\n'))
          if val.size != expected_size:
            print("\nInput size %d doesn't match expected %d" % (val.size, expected_size))
            return False
          df.loc[df[col] != val, f_changed] = True
          df[col] = val
        else:
          val = val if val else None
          df.loc[df[col] != val, [col, f_changed]] = [val, True]
    else:
      old = str(df.loc[idx, col])
      val = prompt("New value for '%s': " % (col), default=old)
      val = val if val else None
      if val != old:
        val = pd.to_numeric(val) if is_numeric_dtype(df[col]) else val
        df.loc[idx, [col, f_changed]] = [val, True]
  except KeyboardInterrupt:
    return False
  return True

def data_columns(df, plus=[]):
  return df.columns.intersection(dataCols + plus)

fs = sys.argv[1:]
if not fs:
  fs = os.listdir()
fs = list(filter(path.isfile, fs))
if not fs:
  print("Nothing interesting at", os.getcwd())
else:
  def row(filepath):
    file = File(filepath)
    tpe = type(file)
    if file and tpe in [MP3, MP4, FLAC]:
      print("Loading %s: %s" % (tpe.__name__, filepath))
    else:
      print("Skipped: %s" % (filepath))
      return None
    extr = partial(get, filepath, file)
    return list(map(extr, dataCols)) + [filename(filepath), filepath, False]

  df = pd.DataFrame(
    data = filter(None, map(row, fs)),
    columns = dataCols + [f_filename, f_file, f_changed]
  ).astype({
    f_num:  'Int32',
    f_year: 'Int32'
  })
  if df.empty:
    print("Nothing to do here!")
    sys.exit(0)
  df.sort_values(f_num, inplace=True)
  df.reset_index(drop=True, inplace=True)

  pd.set_option('display.max_columns', None)
  pd.set_option('display.max_colwidth', 32)

  def printData(what=[]):
    cols = None

    if "all" in what:
      cols = data_columns(df, plus=[f_filename, f_file, f_changed])
    else:
      cols = data_columns(df, plus=[f_filename])

    print()
    pd.set_option('display.width', sh.get_terminal_size().columns)
    if f_changed in what:
      print(df.loc[df[f_changed], cols])
    else:
      print(df[cols])

  printData()

  def handle(string):
    cmd, *opts = string.split(' ')
    if cmd in [cmd_quit, 'q']:
      sys.exit(0)
    if cmd in [cmd_remove, 'rm', 'r']:
      if not opts:
        print("No options provided!")
        return

      remove(opts, df)
      clear()
      printData()
    elif cmd in [cmd_edit, 'e']:
      if not opts:
        print("No options provided!")
        return

      idx = None
      col = None

      for i in opts:
        if i.isdigit():
          idx = int(i)
        elif i in data_columns(df, plus=[f_filename]):
          col = i

      if edit(df, idx, col):
        clear()
        printData()
    elif cmd in [cmd_print, 'p']:
      clear()
      printData(opts)
    elif cmd in [cmd_tags, 't']:
      if not opts:
        print("No options provided!")
        return

      idx = int(opts[0])
      tags = File(df.at[idx, f_file]).pprint()
      print(tags)
    elif cmd in [cmd_write, 'w']:
      write(df)
      sys.exit(0)
    elif cmd:
      print()
      cmd = s.prompt("%s what? " % cmd)
      handle(cmd)

  while True:
    try:
      print()
      cmd = s.prompt("What's next? ")
      handle(cmd)
    except (EOFError, KeyboardInterrupt):
      sys.exit(1)

# vim: et sw=2 ts=2 sts=2