#!/usr/bin/env python3
import sys, re, os, subprocess, pathlib, xattr, langcodes

# Static config
DEBUG        = False
mux_vers     = 1
mux_xattr    = 'zib.mux'
base_path    = '/Volumes/Shared/TV/Willow'
min_size     = 0.9
lang_default = 'eng'

# CLI config
if len(sys.argv) > 1:
  base_path = sys.argv[1]

# Init
root = pathlib.Path(base_path)
re_seaep = re.compile(r'\bS(\d{2,3})E(\d{2,3})\b', re.IGNORECASE)
re_media = re.compile(r'^(.*?)(?:\.\w\w)?(?:\.sdh)?\.(srt|mp4|m4v|mkv)$', re.IGNORECASE)
re_subs  = re.compile(r'(?:\.(\w\w))?(?:\.(sdh))?\.srt$', re.IGNORECASE)
if DEBUG:
  print(f"Root: {root}")

for (parent, dirs, files) in root.walk():
  if DEBUG:
    print(f"Processing: {parent}", file=sys.stderr)

  # Reset for each directory
  media = []

  # Scan this directory for media
  for file in files:
    info = {
      'name': file,
      'path': parent.joinpath(file),
    }

    # Parse media paths
    try:
      m = re_media.search(info['name'])
      if m:
        info['base']   = m.group(1)
        info['ext']    = m.group(2).lower()
        info['g_base'] = info['base'].lower()
    except BaseException as e:
      print(f"Parsing error (media): {e}", file=sys.stderr)
      sys.exit(-1)

    # Skip non-media files
    if 'ext' not in info:
      if DEBUG:
        print(f"Skipping non-media file: {info['name']}", file=sys.stderr)
      continue

    # Parse season/episode data
    try:
      m = re_seaep.search(info['name'])
      if m:
        info['season']  = int(m.group(1))
        info['episode'] = int(m.group(2))
        info['g_seaep'] = 'S{:02d}E{:02d}'.format(info['season'], info['episode'])
    except BaseException as e:
      print(f"Parsing error (seaep): {e}", file=sys.stderr)
      sys.exit(-1)

    # Parse xattrs
    try:
      if (
        'ext' in info and info['ext'] == 'mkv'
         and mux_xattr in xattr.listxattr(info['path'])
      ):
        info['mux'] = int.from_bytes(xattr.getxattr(info['path'], mux_xattr))
    except BaseException as e:
      print(f"Parsing error (xattr): {e}", file=sys.stderr)
      sys.exit(-1)

    if DEBUG:
      print(f"Found media: {info}", file=sys.stderr)
    media.append(info)

  # Find groups
  groups  = {}
  for file in media:
    if 'g_base' in file:
      g_base = file['g_base']
      if file['g_base'] not in groups:
        groups[g_base] = []
      if file not in groups[g_base]:
        groups[g_base].append(file)

  # Merge groups
  for (g_base, group) in groups.items():
    if len(group) == 1:
      if 'mux' in group[0] and group[0]['mux'] >= mux_vers:
        if DEBUG:
          print(f"Skipping non-grouped file with current mux_vers: {group[0]['path']}", file=sys.stderr)
        continue

    if DEBUG:
      print(f"Merging group: {group[0]['base']}", file=sys.stderr)

    # Construct args and file lists
    size    = 0
    sources = []
    subs    = []
    video   = []
    temp    = group[0]['path'].parent.joinpath(f'{group[0]['base']}-remux.mkv')
    final   = group[0]['path'].parent.joinpath(f'{group[0]['base']}.mkv')
    args    = [ 'mkvmerge', '--default-language', lang_default, '-o', temp ]
    for file in group:

      # Track all source files
      size += file['path'].stat().st_size
      sources.append(file['path'])

      # Find the video source
      if file['ext'] in [ 'mkv', 'mp4', 'm4v' ]:
        if DEBUG:
          print(f"Found video: {file['path'].name}", file=sys.stderr)
        video.append(file['path'])

      # Parse subtitle languages
      else:
        if DEBUG:
          print(f"Found subtitle: {file['path'].name}", file=sys.stderr)
        sub = { 'path': file['path'], 'lang': lang_default, 'sdh': False }
        m = re_subs.search(file['path'].name)
        if m:
          for part in m.groups():
            if part.lower() == 'sdh':
              sub['sdh'] = True
            else:
              code = langcodes.Language.get(part.lower())
              if code:
                sub['lang'] = code.to_alpha3()
              else:
                print(f"Cound not match language ({file['path']}): {part}", file=sys.stderr)
              if DEBUG:
                print(f"Using language: {sub['lang']}", file=sys.stderr)
        subs.append(sub)

    # We should find exactly one video source
    if len(video) != 1:
      print(f"Did not find exactly 1 video: {group[0]['base']}", file=sys.stderr)
      continue
    args.extend(video)

    # Subtitles get language and SDH metadata
    for sub in subs:
      args.extend([ '--language', f"0:{sub['lang']}" ])
      if sub['sdh']:
        args.extend([ '--hearing-impaired-flag', '0' ])
      args.append(sub['path'])
    if DEBUG:
      print(f"Mux command: {args}", file=sys.stderr)

    # mkvmerge
    if temp.exists():
      print(f"Mux output exists: {temp}", file=sys.stderr)
      continue
    merge = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    if DEBUG:
      print(f"Muxing result: {merge}", file=sys.stderr)

    # Verify
    if merge.returncode != 0 or not temp.exists() or temp.stat().st_size < (min_size * size):
      print(f"Mux error: {merge.stdout.decode('utf-8')}", file=sys.stderr)
      continue

    # Set xattr
    try:
      xattr.setxattr(temp, mux_xattr, int(1).to_bytes())
    except BaseException as e:
      print(f"xattrs error ({temp}): {e}", file=sys.stderr)
      continue

    # Remove sources
    try:
      for source in sources:
        if DEBUG:
          print(f"Removing: {source}", file=sys.stderr)
        source.unlink()
    except BaseException as e:
      print(f"Unlink error: {e}", file=sys.stderr)
      continue

    # Rename
    err = None
    try:
      if DEBUG:
        print(f"Renaming: {temp} -> {final}", file=sys.stderr)
      temp.rename(final)
    except BaseException as e:
      err = e
    if not final.exists() or final.stat().st_size < (min_size * size):
      err = 'Missing output'
    if err is not None:
      print(f"Rename error ({final}): {e}", file=sys.stderr)
      continue