#!/usr/bin/perl

# This script uses a MythTV database to transcode recorded TV programs to
# MPEG-4 video using the H.264 and AAC codecs, cutting the video at certain
# commercial points selected by the user within MythTV's frontend or by an
# automatic commercial flagger, complete with SRT subtitles extracted from
# MythTV's embedded VBI closed-caption data and iTunes-style metadata about
# the program, if it can be found.

# In short, this script calls a bunch of programs to convert a file like
# 1041_20100523000000.mpg into a file like
# /usr/video/Program Name/Program Name - Episode Name.mp4, including captions,
# file tags, and with commercials clipped.

# Requirements:
# - FFmpeg
# - x264 (ffmpeg will be called with -vcodec libx264)
# - neroAacEnc (http://www.nero.com/enu/technologies-aac-codec.html)
# - MP4Box (http://gpac.sourceforge.net)
# - ccextractor (http://ccextractor.sourceforge.net)
# - AtomicParsley (http://atomicparsley.sourceforge.net - the tarball labelled
#                  'debian' contains a generic Linux binary)

# Usage: <script> %CHANID% %STARTTIME%

# Script adapted from a script by Christopher Neufeld, released GPLv2.
# May not work correctly on recordings that span the autumn time change.

use strict;
use warnings;
use POSIX "mkfifo";
use Cwd;
use DBI;
use File::Copy;

# Make sure the executables mentioned above can be found in
# one of these directories.
$ENV{'PATH'} = '/bin:/usr/bin:/usr/local/bin';

# Enter MythTV database information below. Below are the default values.
my $mysql_host = "localhost";
my $mysql_db = "mythconverg";
my $mysql_user = "mythtv";
my $mysql_passwd = "mythtv";

# This directory will be used to store intermediate files while encoding, and
# should be large enough to hold at least the H.264-encoded version of a
# reasonably sized recording without any commercial cuts.
# Around 4-8GB free should work well, but 2GB might be okay.
my $tmp_dir = "/tmp";

# This directory will store the generated MP4 files. Subdirectories will be
# created for each series.
my $final_dir = "/usr/video";

# When cutting commercials at the beginning or end of a video, sometimes
# MythTV's automatic commercial flagger will mark the cut-point a few seconds
# ahead or behind the actual start or end point of the video. Therefore, any
# clips at the beginning or end of the video that are shorter than this
# threshold (in seconds) will be ignored.
my $thresh = 5;

# Average bitrate of the H.264 video, in kb/s.
my $vid_br = 1000;

# Average quality of the AAC audio, from 0 to 1.
my $audio_q = 0.3;

# Whether to make the resulting files iPod Touch 3G-compatible.
my $iPod = 1;

# ISO 639-2 three-letter language code to use for audio and subtitle tracks.
my $lang_code = "eng";

sub short_time_to_long {
  my $st = shift;
  my $t = "";
  $t .= substr($st, 0, 4) . "-" . substr($st, 4, 2) . "-" . substr($st, 6, 2);
  $t .= " ";
  $t .= substr($st, 8, 2) . ":" . substr($st, 10, 2) . ":" . substr($st, 12, 2);
  return $t; }

sub long_time_to_short {
  my $lt = shift;
  $lt =~ tr/ :-//d;
  return $lt; }

sub seconds_to_time {
  my $seconds = shift;
  my $hours = int($seconds / 3600);
  $seconds -= $hours * 3600;
  my $minutes = int($seconds / 60);
  $seconds -= $minutes * 60;
  return sprintf "%.2d:%.2d:%.2d", $hours, $minutes, $seconds; }

sub seconds_to_time_frac {
  my $seconds = shift;
  my $hours = int($seconds / 3600);
  $seconds -= $hours * 3600;
  my $minutes = int($seconds / 60);
  $seconds -= $minutes * 60;
  my $split = $seconds - int($seconds);
  my $s = "";
  if($split != 0) {
    $s = sprintf "%.4f", $split;
    $s =~ s/^[^.]*\./\./; }
  return sprintf "%.2d:%.2d:%.2d$s", $hours, $minutes, $seconds; }

sub add_trail_slash {
  my $str = shift;
  if($$str !~ m/\/$/) {
    $$str .= "/"; } }

sub get_path {
  my ($channel, $long_time) = @_;
  my $path;
  my $db = DBI->connect("DBI:mysql:database=$mysql_db;host=$mysql_host",
                        $mysql_user, $mysql_passwd);
  die "Could not connect to database." unless defined $db;
  my $query = $db->prepare("select endtime, basename, storagegroup from " .
                           "recorded where chanid = ? and starttime = ?;");
  $query->execute($channel, $long_time);
  my ($endtime, $basename, $storagegroup) = $query->fetchrow_array();
  die "Unable to locate recording basename for channel $channel " .
      "and time $long_time." unless defined $basename;
  $query = $db->prepare("select dirname from storagegroup " .
                        "where groupname = ?;");
  $query->execute($storagegroup);
  while(my @dir = $query->fetchrow_array()) {
    if(-e $dir[0] . $basename) {
      $path = $dir[0] . $basename; } }
  die "Unable to find recording." unless defined $path;
  return $path; }

sub get_video_param {
  my $video = shift;
  my ($fps, $duration, $res);
  open my $pipe, "ffmpeg -i $video 2>&1 |"
      or die "Unable to determine FPS / duration.";
  foreach my $line (<$pipe>) {
    if($line =~ m/([0-9]*\.?[0-9]*) fps/) {
      $fps = $1; }
    elsif($line =~ m/Duration: ([0-9]+):([0-9]+):([0-9]+)\.([0-9]+),/) {
      $duration = $4 / 100 + $3 + $2 * 60 + $1 * 3600; }
    if($line =~ m/Video:.* ([0-9]+)x([0-9]+) \[/) {
      $res = [$1, $2]; } }
  close $pipe;
  die "Unable to determine FPS of video." unless defined $fps;
  die "Unable to determine duration of video." unless defined $duration;
  die "Unable to determine resolution of video." unless defined $res;
  return ($fps, $duration, $res); }

sub get_cutlist {
  my ($channel, $time) = @_;
  my @cutlist;
  print "*** Obtaining cut list ***\n";
  open my $pipe, "mythcommflag --getcutlist -c $channel -s $time 2>/dev/null |"
      or die "Unable to obtain cut list.";
  foreach my $line (<$pipe>) {
    if($line =~ m/^Cutlist: /) {
      my @afterCut = split / /, $line;
      if(@afterCut > 1 and $afterCut[1] ne "\n") {
        @cutlist = split /,/, $afterCut[1]; } } }
  close $pipe;
  return @cutlist; }

sub get_cut {
  my $cut = shift;
  my @spl = split /-/, $cut;
  return ($spl[0], $spl[1]); }

sub extract_captions {
  my ($start, $end, $delay, $subs, $fps, $files) = @_;
  my $start_s = seconds_to_time($start / $fps);
  my $end_s = seconds_to_time($end / $fps);
  my $delay_s = int($delay / $fps * 1000);
  my $blank = 1;
  print "*** Extracting captions to $files->{'srt'} ***\n";
  system "ccextractor -o $files->{'fifo'} -startat $start_s -endat $end_s " .
         "-delay -$delay_s -utf8 --no_progress_bar $files->{'video'} " .
         "> /dev/null 2>&1 &";
  open my $pipe, '<', $files->{'fifo'};
  open my $file, '>>', $files->{'srt'};
  foreach my $line (<$pipe>) {
    if($blank == 1) {
      print $file ++$$subs . "\n";
      $blank = 0; }
    else {
      if($line =~ m/^\s*$/) {
        $blank = 1; }
      print $file $line; } }
  close $pipe;
  close $file; }

sub adjust_res {
  my $res = shift;
  my $fmt = "mp4";
  my $prof = "high";
  my $size = "";
  if($iPod) {
    my $aspect = $res->[0] * 1.0 / $res->[1];
    if($aspect > (640.0 / 480.0)) {
      my $hres = int(640.0 / $aspect + 0.5);
      if($hres % 2 == 1) {
        $hres += 1; }
      my $pad = (480 - $hres) / 2;
      $size = "-s 640x" . $hres . " -padtop $pad -padbottom $pad"; }
    else {
      my $vres = int(480.0 * $aspect + 0.5);
      if($vres % 2 == 1) {
        $vres += 1; }
      my $pad = (640 - $vres) / 2;
      $size = "-s " . $vres . "x480 -padleft $pad -padright $pad"; }
    $fmt = "ipod";
    $prof = "ipod640"; }
  return ($fmt, $prof, $size); }

sub encode_video {
  my ($start, $end, $fps, $res, $suffix, $files) = @_;
  my $start_s = $start / $fps;
  my $dur_s = $end / $fps - $start_s;
  my $file = $files->{'h264'} . $suffix;
  my $cwd = getcwd();
  unlink $file if -e $file;
  chdir $files->{'dir'};
  my ($fmt, $prof, $size) = adjust_res($res);
  print "*** Encoding video to $files->{'dir'} - first pass ***\n";
  system "ffmpeg -y -i $files->{'video'} -ss $start_s -t $dur_s -pass 1 " .
         "-an -vcodec libx264 -vb " . $vid_br . "k -bt " . $vid_br . "k " .
         "-vpre slow_firstpass -vpre $prof $size -threads 0 -f h264 " .
         "/dev/null > /dev/null 2>&1";
  print "*** Encoding video to $file - second pass ***\n";
  system "ffmpeg -y -i $files->{'video'} -ss $start_s -t $dur_s -pass 2 " .
         "-an -vcodec libx264 -vb " . $vid_br . "k -bt " . $vid_br . "k " .
         "-vpre slow -vpre $prof $size -threads 0 -f $fmt $file " .
         "> /dev/null 2>&1";
  die "Could not encode video." unless -e $file and ! -z $file;
  chdir $cwd; }

sub encode_audio {
  my ($start, $end, $fps, $suffix, $files) = @_;
  my $start_s = $start / $fps;
  my $dur_s = $end / $fps - $start_s;
  my $file = $files->{'aac'} . $suffix;
  unlink $file if -e $file;
  print "*** Encoding audio to $file ***\n";
  system "ffmpeg -y -i $files->{'video'} -ss $start_s -t $dur_s -vn " .
         "-acodec pcm_s16le -f wav - 2>/dev/null | " .
         "neroAacEnc -ignorelength -q $audio_q -if - -of $file " .
         "> /dev/null 2>&1";
  die "Could not encode audio." unless -e $file and ! -z $file; }

sub segment {
  my ($start, $end, $delay, $seg, $subs, $fps, $res, $files) = @_;
  my $suffix = "-" . $$seg++;
  my $start_s = seconds_to_time($start / $fps);
  my $end_s = seconds_to_time($end / $fps);
  print "*** Segment $$seg: [$start_s - $end_s] ***\n";
  extract_captions($start, $end, $delay, $subs, $fps, $files);
  encode_audio($start, $end, $fps, $suffix, $files);
  encode_video($start, $end, $fps, $res, $suffix, $files); }

sub clean_tmp {
  my ($files, $fifo) = @_;
  my @log2pass = ("ffmpeg2pass-0.log", "x264_2pass.log",
                  "x264_2pass.log.mbtree");
  print "*** Removing temporary files ***\n";
  unlink $files->{'srt'} if -e $files->{'srt'};
  if($fifo) {
    mkfifo $files->{'fifo'}, 0700 unless -p $files->{'fifo'}; }
  else {
    unlink $files->{'fifo'} if -e $files->{'fifo'}; }
  for my $file (@log2pass) {
    $file = $files->{'dir'} . $file;
    unlink $file if -e $file; } }

sub cut_and_transcode {
  my ($duration, $fps, $r, $cutlist, $files) = @_;
  my $pos = 0;
  my $delay = 0;
  my $seg = 0;
  my $subs = 0;
  clean_tmp($files, 1);
  for my $cut (@$cutlist) {
    my ($start, $end) = get_cut($cut);
    if($start > $thresh * $fps && $start > $pos) {
      segment($pos, $start, $delay, \$seg, \$subs, $fps, $r, $files); }
    $delay += $end - $start;
    $pos = $end; }
  if($pos < ($duration - $thresh) * $fps) {
    segment($pos, $duration * $fps, $delay, \$seg, \$subs, $fps, $r, $files); }
  unlink $files->{'fifo'} if -p $files->{'fifo'};
  return $seg; }

sub remux {
  my ($seg, $files) = @_;
  print "*** Remuxing to $files->{'mp4'} ***\n";
  unlink $files->{'mp4'} if -e $files->{'mp4'};
  unlink $files->{'mp4tmp'} if -e $files->{'mp4tmp'};
  my $tmp = $files->{'mp4'};
  $tmp =~ s/\/[^\/]*$//;
  for(my $s = 0; $s < $seg; $s++) {
    my $h264 = $files->{'h264'} . "-" . $s;
    my $aac = $files->{'aac'} . "-" . $s;
    if($s == 0) {
      system "MP4Box -tmp $tmp -add $h264 -add $aac $files->{'mp4'}" .
             " > /dev/null 2>&1";
      die "Could not remux video." unless -e $files->{'mp4'}
          and ! -z $files->{'mp4'}; }
    else {
      system "MP4Box -tmp $tmp -add $h264 -add $aac $files->{'mp4tmp'}" .
             " > /dev/null 2>&1";
      die "Could not remux video." unless -e $files->{'mp4tmp'}
          and ! -z $files->{'mp4tmp'};
      system "MP4Box -tmp $tmp -cat $files->{'mp4tmp'} $files->{'mp4'}" .
             " > /dev/null 2>&1";
      unlink $files->{'mp4tmp'}; }
    unlink $h264;
    unlink $aac; }
  system "MP4Box -tmp $tmp -add $files->{'srt'} $files->{'mp4'}" .
         " > /dev/null 2>&1";
  system "MP4Box -tmp $tmp -lang $lang_code $files->{'mp4'}" .
         " > /dev/null 2>&1";
  unlink $files->{'srt'}; }

sub get_metadata {
  my ($channel, $long_time) = @_;
  my ($path, @credits);
  my $count = 0;
  my $db = DBI->connect("DBI:mysql:database=$mysql_db;host=$mysql_host",
                        $mysql_user, $mysql_passwd);
  die "Could not connect to database." unless defined $db;
  my $query = $db->prepare("select title, subtitle, description, category," .
                           "originalairdate, syndicatedepisodenumber from " .
                           "recordedprogram where chanid = ? and " .
                           "starttime = ?;");
  $query->execute($channel, $long_time);
  my $hash = $query->fetchrow_hashref();
  die "Unable to obtain metadata for channel $channel and time $long_time."
      unless defined $hash;
  $query = $db->prepare("select rating from recordedrating " .
                        "where chanid = ? and starttime = ?;");
  $query->execute($channel, $long_time);
  $hash->{'rating'} = $query->fetchrow_array();
  $query = $db->prepare("select name from channel where chanid = ?;");
  $query->execute($channel);
  $hash->{'channel'} = $query->fetchrow_array();
  $query = $db->prepare("select name, role from recordedcredits, people " .
                        "where chanid = ? and starttime = ? and " .
                        "people.person = recordedcredits.person order by " .
                        "role, substring_index(name, ' ', -1) asc;");
  $query->execute($channel, $long_time);
  while(my @person = $query->fetchrow_array()) {
    my $role = lc $person[1];
    $role =~ s/_/ /g;
    push @credits, [$person[0], $role]; }
  $hash->{'credits'} = \@credits;
  return $hash; }

sub print_metadata {
  my $meta = shift;
  print "*** Printing file metadata ***\n";
  for my $key (keys %$meta) {
    if($key ne "credits") {
      my $val = $meta->{$key};
      print "$key: $val\n"; } }
  my $credits = $meta->{'credits'};
  if(@$credits == 0) {
    print "No credits\n"; }
  else {
    print "Credits:\n";
    for my $credit (@$credits) {
      print "$credit->[0] ($credit->[1])\n"; } } }

sub get_versions {
  my ($ffmpeg, $x264, $nero, $pipe);
  open $pipe, "ffmpeg -version 2> /dev/null |";
  foreach my $line (<$pipe>) {
    if($line =~ m/^(FFmpeg.*)$/) {
      $ffmpeg = $1; } }
  close $pipe;
  open $pipe, "x264 --version 2> /dev/null |";
  foreach my $line (<$pipe>) {
    if($line =~ m/^(x264.*)$/) {
      $x264 = $1; } }
  close $pipe;
  open $pipe, "neroAacEnc 2>&1 |";
  foreach my $line (<$pipe>) {
    if($line =~ m/Package version:\s*([0-9.]+)/) {
      $nero = "Nero AAC Encoder $1"; } }
  close $pipe;
  return "$ffmpeg, $x264, $nero"; }

sub xml_credits_section {
  my ($people, $key, $xml) = @_;
  if(@$people > 0) {
    $$xml .= "    <key>$key</key>\n";
    $$xml .= '    <array>' . "\n";
    for my $member (@$people) {
      $$xml .= '      <dict>' . "\n";
      $$xml .= '        <key>name</key>' . "\n";
      $$xml .= "        <string>$member</string>\n";
      $$xml .= '      </dict>' . "\n"; }
    $$xml .= '    </array>' . "\n"; } }

sub xml_credits {
  my $meta = shift;
  my $credits = $meta->{'credits'};
  my $xml;
  if(@$credits == 0) {
    return $xml; }
  my (@cast, @directors, @producers);
  for my $credit (@$credits) {
    if($credit->[1] eq 'actor' or $credit->[1] eq 'host') {
      push @cast, $credit->[0]; }
    elsif($credit->[1] eq 'director') {
      push @directors, $credit->[0]; }
    elsif($credit->[1] eq 'executive producer') {
      push @producers, $credit->[0]; } }
  for my $credit (@$credits) {
    if($credit->[1] eq 'guest star') {
      push @cast, $credit->[0]; }
    elsif($credit->[1] eq 'producer') {
      push @producers, $credit->[0]; } }
  for my $credit (@$credits) {
    if($credit->[1] eq 'writer') {
      push @producers, $credit->[0]; } }
  if(@cast == 0 and @directors == 0 and @producers == 0) {
    return $xml; }
  $xml .= '<?xml version="1.0" encoding="UTF-8"?>' . "\n";
  $xml .= '<!DOCTYPE plist PUBLIC "-//Apple Computer//DTD PLIST 1.0//EN"' .
          ' "http://www.apple.com/DTDs/PropertyList-1.0.dtd">' . "\n";
  $xml .= '<plist version="1.0">'. "\n";
  $xml .= '  <dict>' . "\n";
  xml_credits_section(\@cast, 'cast', \$xml);
  xml_credits_section(\@directors, 'directors', \$xml);
  xml_credits_section(\@producers, 'producers', \$xml);
  $xml .= '  </dict>' . "\n";
  $xml .= '</plist>' . "\n\n";
  $xml =~ s/\"/\\\"/g;
  return $xml; }

sub write_metadata {
  my ($long_time, $meta, $files) = @_;
  my $str = "AtomicParsley $files->{'mp4'}";
  my $versions = get_versions();
  my $file = $files->{'mp4'};
  $file =~ s/\.mp4$//;
  my $final;
  if(defined $meta->{'title'} and defined $meta->{'subtitle'}) {
    my $title_dir = $files->{'finaldir'} . $meta->{'title'};
    mkdir $title_dir unless -d $title_dir;
    $final = $title_dir . "/" . $meta->{'title'} . " - " .
             $meta->{'subtitle'} . ".mp4"; }
  else {
    $final = $files->{'mp4'}; }
  my @tmp = glob "$file-temp-*.mp4";
  for my $f (@tmp) {
    unlink $f; }
  print "*** Adding metadata to $final ***\n";
  my $xml = xml_credits($meta);
  if(defined $meta->{'title'}) {
    my $title = $meta->{'title'};
    $str .= " --artist \"$title\" --album \"$title\"" .
            " --albumArtist \"$title\" --TVShowName \"$title\""; }
  if(defined $meta->{'subtitle'}) {
    $str .= " --title \"$meta->{'subtitle'}\""; }
  if(defined $meta->{'category'}) {
    $str .= " --genre \"$meta->{'category'}\""; }
  if(defined $meta->{'originalairdate'}) {
    $str .= " --year $meta->{'originalairdate'}"; }
  if(defined $meta->{'description'}) {
    $str .= " --description \"$meta->{'description'}\""; }
  if(defined $meta->{'channel'}) {
    $str .= " --TVNetwork \"$meta->{'channel'}\""; }
  if(defined $meta->{'syndicatedepisodenumber'}) {
    $str .= " --TVEpisode \"$meta->{'syndicatedepisodenumber'}\""; }
  if(defined $meta->{'rating'}) {
    $str .= " --contentRating \"$meta->{'rating'}\""; }
  if(-e $files->{'video'} . ".png") {
    my $art = $files->{'video'} . ".png";
    $str .= " --artwork $art"; }
  $long_time =~ s/ /T/;
  $long_time .= "Z";
  $str .= " --stik \"TV Show\"";
  $str .= " --purchaseDate $long_time";
  $str .= " --encodingTool \"$versions\"";
  $str .= " --grouping \"MythTV Recording\"";
  if(defined $xml) {
    $str .= " --rDNSatom \"$xml\" name=iTunMOVI domain=com.apple.iTunes"; }
  system $str . " > /dev/null 2>&1";
  @tmp = glob "$file-temp-*.mp4";
  if(@tmp == 1) {
    move $tmp[0], $final;
    unlink $files->{'mp4'}; }
  else {
    die "Could not add metadata tags."; } }

sub generate_edl {
  my ($fps, $cutlist, $files) = @_;
  open my $file, '>', $files->{'edl'} or die "Could not write to EDL file.";
  for my $cut (@$cutlist) {
    my ($start, $end) = get_cut($cut);
    $start /= $fps;
    $end /= $fps;
    print $file "$start $end 0\n"; }
  close $file; }

die "Usage: $0 <channel> <start-time>" if @ARGV != 2;
add_trail_slash(\$tmp_dir);
add_trail_slash(\$final_dir);
my $channel = $ARGV[0];
my $time = $ARGV[1];
my $long_time = short_time_to_long($time);
my $video = get_path($channel, $long_time);
my ($fps, $duration, $res) = get_video_param($video);
my $base = $tmp_dir . $channel . "_" . $time;
my $fifo = $base . ".tmp";
my $srt = $base . ".srt";
my $edl = $base . ".edl";
my $h264 = $base . ".h264";
my $aac = $base . ".aac";
my $xml = $base . ".xml";
my $mp4 = $video;
$mp4 =~ s/\.[^.]*$//;
$mp4 .= ".mp4";
my $mp4tmp = $mp4 . "-tmp";
my %files = (video => $video, fifo => $fifo, srt => $srt, edl => $edl,
             h264 => $h264, aac => $aac, xml => $xml, mp4 => $mp4,
             mp4tmp => $mp4tmp, dir => $tmp_dir, finaldir => $final_dir);
print "*** Transcoding $video ***\n";
my @cutlist = get_cutlist($channel, $time);
my $seg = cut_and_transcode($duration, $fps, $res, \@cutlist, \%files);
remux($seg, \%files);
my $meta = get_metadata($channel, $time);
print_metadata($meta);
write_metadata($long_time, $meta, \%files);
clean_tmp(\%files, 0);