#!/usr/bin/env python3

import os
import sys
import subprocess
import json
import argparse
import glob

class config:
  # command paths
  ffmpeg = '/usr/local/bin/ffmpeg'
  ffprobe = '/usr/local/bin/ffprobe'

  # default values (will be ignored if copyable)
  ## video
  video_copy_if_possible = True
  video_codec = 'libx264'
  video_crf = 18
  ## audio
  audio_copy_if_possible = True
  audio_codec = 'libfdk_aac'
  audio_bitrate = '256k'
  ## other
  output_format = 'mp4'
  options = '-movflags faststart'

class color:
  red = '\033[31m'

def encode(in_file, verbose):
  print(f'Encoding: {os.path.basename(in_file)} to {config.output_format}')
  video_copy = False
  audio_copy = False
  proc = subprocess.run(f'{config.ffprobe} -loglevel quiet -show_streams -print_format json "{in_file}"', shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  streams = json.loads(proc.stdout.decode('utf8')).get('streams')
  if streams == None:
    print(color.red + '[Error] Source file is not a valid video file.')
    return

  video = None
  videos = [x for x in streams if x['codec_type'] == 'video']
  if len(videos) == 0:
    print(color.red + '[Error] Source file does not contain any video stream.')
  else:
    video = videos[0]

  audio = None
  audios = [x for x in streams if x['codec_type'] == 'audio']
  if len(audios) == 0:
    print(color.red + '[Error] Source file does not contain any audio stream.')
  else:
    jpn_audios = [x for x in audios if x['tags']['language'] == 'jpn']
    if len(jpn_audios) == 0:
      print('[Info] Source file does not contain Japanese audio stream. Using the first one...')
      audio = audios[0]
    else:
      audio = jpn_audios[0]
  
  if video == None or audio == None:
    return
  
  if verbose:
    print(f'video: {json.dumps(video, indent=4)}')
    print(f'audio: {json.dumps(audio, indent=4)}')

  if video['codec_name'] == 'h264':
    video_copy = config.video_copy_if_possible
  if audio['codec_name'] == 'aac':
    audio_copy = config.audio_copy_if_possible

  video_codec = 'copy' if video_copy else config.video_codec
  audio_codec = 'copy' if audio_copy else config.audio_codec
  options = config.options
  options += '' if video_copy else f' -crf {config.video_crf}'
  options += '' if audio_copy else f' -b:a {config.audio_bitrate}'
  options += ' -loglevel verbose' if verbose else ' -loglevel warning -stats'
  out_file = f'{os.path.splitext(in_file)[0]}.{config.output_format}'

  if os.path.exists(out_file):
    print('[Info] Output file already exists. Skipping.')
    return

  if verbose:
    print(f'{config.ffmpeg} -i "{in_file}" -c:v {video_codec} -c:a {audio_codec} {options} "{out_file}"')
  try:
    subprocess.run(f'{config.ffmpeg} -i "{in_file}" -c:v {video_codec} -c:a {audio_codec} {options} "{out_file}"', shell=True)
  except KeyboardInterrupt:
    print(f'\n[Info] Keyboard Interrupt: Removed temporary file.')
    os.remove(out_file)
    exit(1)
    

if __name__ == '__main__':
  argparser = argparse.ArgumentParser()
  argparser.add_argument('path', type=str, help='input video file or container directory')
  argparser.add_argument('-v', '--verbose', action='store_true', help='show verbose message')
  argparser.add_argument('-r', '--recursive', action='store_true', help='process files recursively in directory')
  argparser.add_argument('-c:v', '--video-codec', type=str, dest='video_codec', help=f'set video codec (default: {config.video_codec})')
  argparser.add_argument('-crf', '--video-crf', type=int, dest='video_crf', help=f'set video crf value (default: {config.video_crf})')
  argparser.add_argument('-c:a', '--audio-codec', type=str, dest='audio_codec', help=f'set audio codec (default: {config.audio_codec})')
  argparser.add_argument('-b:a', '--audio-bitrate', type=str, dest='audio_bitrate', help=f'set audio bitrate value: (default: {config.audio_bitrate})')
  argparser.add_argument('-if', '--input-format', type=str, dest='input_format', help='set input format to indicate which file to process in the directory (ex: mp4)')
  argparser.add_argument('-of', '--output-format', type=str, dest='output_format', help=f'set output file format (default: {config.output_format})')
  args = argparser.parse_args()

  if args.video_codec:
    config.video_codec = args.video_codec
  if args.video_crf:
    config.video_crf = args.video_crf
  if args.audio_codec:
    config.audio_codec = args.audio_codec
  if args.audio_bitrate:
    config.audio_bitrate = args.audio_bitrate
  if args.output_format:
    config.output_format = args.output_format

  if os.path.isdir(args.path):
    if args.input_format:
      if args.input_format == config.output_format:
        print(color.red + '[Error] Output format is the same as input format.')
        exit()
      else:
        if args.recursive: 
          for in_file in glob.iglob(f'{args.path}/**/*.{args.input_format}', recursive=True):
            encode(in_file, args.verbose)
        else:
          for in_file in glob.iglob(f'{args.path}/*.{args.input_format}'):
            encode(in_file, args.verbose)
    else:
      print(color.red + '[Error] To encode videos recursively, you need to specify input format with --input-format (-if) option.')
      exit()
  elif os.path.splitext(args.path)[1] == config.output_format:
    print(color.red + '[Error] Output format is the same as input format.')
    exit()
  else:
    encode(args.path, args.verbose)
