#! /usr/bin/env python

# flo-srttool --- Perform operations on SubViewer (.srt) subtitles
# Copyright (c) 2003, 2004, 2008 Florent Rougon
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 dated June, 1991.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; see the file COPYING. If not, write to the
# Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
# Boston, MA  02110-1301 USA.

from __future__ import nested_scopes, division
import sys, os, locale, re, getopt, math
import flo_small_funcs

progname = os.path.basename(sys.argv[0])
progversion = "0.2"
version_blurb = """Written by Florent Rougon.

Copyright (c) 2008 Florent Rougon
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE."""

usage = """Usage: %(progname)s [OPTIONS] [INPUT_FILE]
Perform operations on a subtitles file (.srt type).

Options:
      --config-file=FILE use FILE instead of the default configuration
                       file, ~/.%(progname)s/config.py
  -d, --delay=SECONDS  shift all subtitles by SECONDS seconds
                       (float, default = 0.0)
  -a, --affine-adjustment=F1_F2-S1_S2
                       adjust the subtitle times with an affine
                       transformation; you have to identify 2 pairs
                       corresponding subtitles, one preferably
                       towards the beginning of the movie and the
                       other preferably towards the end of the movie.
                       F1 and S1 respectively indicate the times of
                       the first and second subtitles chosen in a
                       correctly-timed subtitle file, while F2 and S2
                       indicate the times present in INPUT_FILE for
                       the corresponding subtitles (wrong times with
                       respect to the movie). F1, S1, F2, S2 are all
                       expected to be in the same format as used in
                       .srt files, i.e., HH:MM:SS,mmm. The program
                       will apply an affine transformation to all
                       times of INPUT_FILE so that in the output
                       file, the first and second chosen subtitles
                       respectively happen at times F1 and S1.
  -A, --affine=A,B     apply the x -> Ax+B affine transformation to
                       every time in INPUT_FILE, where b is given in
                       seconds (float)
  -i, --index=N        add N to the subtitles' indices
                       (integer, default = 0)
  -o, --output=OFILE   output file name (default: standard output)
  -n, --newlines=STYLE newline convention for the output file ('LF'
                       for Unix, 'CR' for Macintosh, 'CRLF' for
                       Windows, or 'orig' for the same line ending
                       type as in the input file (default: 'orig')
      --help           display usage information and exit
      --version        output version information and exit""" \
  % { "progname": progname }


indexline_cre = re.compile(r"^[ \t]*(?P<index>[0-9]+)[ \t]*$")
# time_rec = re.compile(
#     r"^(?P<sh>\d{2}):(?P<sm>\d{2}):(?P<ss>\d{2}),(?P<smilli>\d{3})"
#     r"[ \t]+-+>[ \t]+"
#     r"(?P<eh>\d{2}):(?P<em>\d{2}):(?P<es>\d{2}),(?P<emilli>\d{3})\n")

time_re = r"(?P<HH>[0-9]{2}):(?P<MM>[0-9]{2}):(?P<SS>[0-9]{2})," \
          r"(?P<mmm>[0-9]{3})"
time_cre = re.compile(time_re)
anon_time_re = r"([0-9]{2}):([0-9]{2}):([0-9]{2}),([0-9]{3})"

time_span_cre = re.compile(r"%(time)s[ \t]+-+>[ \t]+%(time)s" %
                           {"time": anon_time_re})


# Exceptions raised by this module
class error(Exception):
    """Base class for exceptions in %s.""" % progname
    def __init__(self, message=None):
        self.message = message
    def __str__(self):
        return "<%s: %s>" % (self.__class__.__name__, self.message)
    def complete_message(self):
        if self.message:
            return "%s: %s" % (self.ExceptionShortDescription, self.message)
        else:
            return "%s" % self.ExceptionShortDescription
    ExceptionShortDescription = "%s generic exception" % progname

# This one is best derived directly from 'Exception' instead of 'error', so
# that we get a proper traceback when it is raised.
class ProgramError(Exception):
    """Exception raised when the program finds that it has a bug."""
    ExceptionShortDescription = "bug in %s" % progname

class NoSuchConfigurationFile(error):
    """Exception raised when the user specified a configuration file that doesn't exist."""
    ExceptionShortDescription = "no such configuration file"

class UserError(error):
    """Exception raised when the program is used incorrectly."""
    ExceptionShortDescription = "user error"

class SeveralConfigFileOptionsSupplied(UserError):
    """Exception raised when several --config-file options were supplied."""
    ExceptionShortDescription = "several --config-file options were supplied"


def write_output(params, l):
    # Set params["EOL"] to the desired string for representing end of lines
    determine_EOL_style(params)
    params["output file"].writelines(map(
        lambda s: s.replace('\n', params["EOL"]), l))


def determine_EOL_style(params):
    if params["input file"].newlines not in ('\r', '\n', '\r\n'):
        sys.exit("Unable to determine the line ending type of the input "
                 "file.")
        
    if params["newlines"] == "orig":
        params["EOL"] = params["input file"].newlines
    else:
        nl_mapping = {"LF": '\n',
                      "CR": '\r',
                      "CRLF": '\r\n'}
        try:
           params["EOL"] = nl_mapping[params["newlines"]]
        except KeyError:
            raise ProgramError()


def next_subtitle(params):
    lineno = 0
    negative_or_null_indices = 0
    
    ifile = params["input file"]
    output = []
    
    while True:
        # Skip newlines
        while True:
            l = ifile.readline()
            lineno += 1

            if l == "":
                write_output(params, output)
                return
            elif l == "\n":
                output.append(l)
            else:
                mo = indexline_cre.match(l)
                if mo is None:
                    sys.exit("Invalid format for line %u (expected to be "
                             "a subtitle index)" % lineno)
                new_index = int(mo.group("index"), 10) + params["index"]
                if new_index <= 0:
                    negative_or_null_indices += 1
                output.append("%d\n" % new_index)
                break

        # Read and process the line containing the timing stuff
        l = ifile.readline()
        lineno += 1

        mo = time_span_cre.match(l)
        if mo is None:
            sys.exit("Invalid format for line %u (expected to indicate the start "
                     "and end times of a subtitle)" % lineno)

        sh, sm, ss, smilli, eh, em, es, emilli = map(int, mo.groups())
        old_start_time = ss + 60*(sm+60*sh) + (smilli / 1000.)
        old_end_time = es + 60*(em+60*eh) + (emilli / 1000.)

        if params["delay"] is not None:
            new_start_time = old_start_time + params["delay"]
            new_end_time = old_end_time + params["delay"]
        elif params["affine"] is not None:
            a, b = params["affine"]
            new_start_time = a*old_start_time + b
            new_end_time = a*old_end_time + b
        else:
            new_start_time = old_start_time
            new_end_time = old_end_time

        new_sfrac, new_sint = math.modf(new_start_time)
        new_efrac, new_eint = math.modf(new_end_time)

        new_smilli = round(new_sfrac*1000)
        new_emilli = round(new_efrac*1000)

        smin, new_ss = divmod(new_sint, 60)
        new_sh, new_sm = divmod(smin, 60)

        emin, new_es = divmod(new_eint, 60)
        new_eh, new_em = divmod(emin, 60)

        output.append("%02u:%02u:%02u,%03u --> %02u:%02u:%02u,%03u\n" %
                    (new_sh, new_sm, new_ss, new_smilli, new_eh, new_em, new_es,
                     new_emilli))

        # Read and copy the subtitle's remaining lines
        while True:
            line = ifile.readline()
            lineno += 1
            if line == "":
                write_output(params, output)
                return
            elif line == "\n":
                output.append(line)
                break
            else:
                output.append(line)

        yield negative_or_null_indices


def process_command_line_and_config_file():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "d:i:a:A:o:n:",
                                   ["config-file=",
                                    "delay=",
                                    "index=",
                                    "affine-adjustment",
                                    "affine=",
                                    "output=",
                                    "newlines=",
                                    "help",
                                    "version"])
    except getopt.GetoptError, message:
        sys.stderr.write(usage + "\n")
        return ("exit", 1)

    # Let's start with the options that don't require any non-option argument
    # to be present
    for option, value in opts:
        if option == "--help":
            print usage
            return ("exit", 0)
        elif option == "--version":
            print "%s %s\n%s" % (progname, progversion, version_blurb)
            return ("exit", 0)

    # Now, require a correct invocation.
    if len(args) not in (0, 1):
        sys.stderr.write(usage + "\n")
        return ("exit", 1)

    params = {}

    # Get the home directory, if any, and store it in params (often useful).
    params["home_dir"] = os.path.expanduser("~")

    # Default values for options
    params["delay"] = None
    params["index"] = 0
    params["affine"] = None
    params["output file"] = sys.stdout
    params["newlines"] = "orig"

    # Check if --config-file was used
    user_cfg_file = None
    option_supplied = False
    for option, value in opts:
        if option == "--config-file":
            if option_supplied:
                raise SeveralConfigFileOptionsSupplied()
            else:
                option_supplied = True
                user_cfg_file = value

                if not os.path.exists(user_cfg_file):
                    raise NoSuchConfigurationFile(user_cfg_file)
        
    # If --config-file wasn't supplied, use the default per-user config file
    if user_cfg_file is None:
        user_cfg_file = os.path.join(params["home_dir"],
                                     ".%s" % progname, "config.py")

    # Update 'params' with those set in the config files, if they exist
    system_cfg_file = os.path.join("/etc", "%s.py" % progname)
    cfg_files = (system_cfg_file, user_cfg_file)

    recognized_params = ("newlines",)

    flo_small_funcs.import_params_from_python_cfg_files(
        namespace=params, cfg_files=cfg_files,
        recognized_params=recognized_params, prefix="")

    # Perform tilde expansion on parameters that represent files or
    # directories and were set in the configuration file (of course, for
    # arguments given on the command line, we let the shell perform the tilde
    # expansion).
#     for key in ("output_dir", "device"):
#         if params[key] is not None:
#             params[key] = os.path.expanduser(params[key])

    # General option processing
    for option, value in opts:
        if option in ("-d", "--delay"):
            try:
                params["delay"] = float(value)
            except ValueError:
                sys.stderr.write("Error:\n\ninvalid delay. Should be a "
                                 "floating point number.\n")
                return ("exit", 1)
        elif option in ("-i", "--index"):
            try:
                params["index"] = int(value)
            except ValueError:
                sys.stderr.write(
                    "Error:\n\ninvalid index constant: %s. Should be an "
                    "integer.\n" % repr(value))
                return ("exit", 1)
        elif option in ("-a", "--affine-adjustment"):
            mo = re.match(r"%(time)s_%(time)s-%(time)s_%(time)s$"
                          % {"time": anon_time_re}, value)
            if mo is None:
                sys.stderr.write(
                    "Invalid syntax for --affine-adjustment option: '%s'\n"
                    % value)
                return ("exit", 1)

            f1_HH, f1_MM, f1_SS, f1_mmm, f2_HH, f2_MM, f2_SS, f2_mmm, \
                   s1_HH, s1_MM, s1_SS, s1_mmm, s2_HH, s2_MM, s2_SS, s2_mmm \
                   = map(int, mo.groups())

            # Convert the times in milliseconds
            f1 = ((f1_HH*60 + f1_MM)*60 + f1_SS)*1000 + f1_mmm
            s1 = ((s1_HH*60 + s1_MM)*60 + s1_SS)*1000 + s1_mmm
            f2 = ((f2_HH*60 + f2_MM)*60 + f2_SS)*1000 + f2_mmm
            s2 = ((s2_HH*60 + s2_MM)*60 + s2_SS)*1000 + s2_mmm

            # Compute the coefficients of the affine transformation
            a = (s1-f1)/(s2-f2)
            # b has to be expressed in seconds
            b = (f1-a*f2)/1000

            params["affine"] = (a, b)
        elif option in ("-A", "--affine"):
            mo = re.match(r"([^,]+),([^,]+)$", value)
            if mo is None:
                sys.stderr.write("Invalid syntax for --affine option: '%s'\n"
                                 % value)
                return ("exit", 1)

            try:
                a, b = map(float, mo.groups())
            except ValueError:
                sys.stderr.write("Invalid syntax for --affine option: '%s'\n"
                                 % value)
                return ("exit", 1)

            params["affine"] = (a, b)
        elif option in ("-o", "--output"):
            params["output file"] = open(value, "wb")
        elif option in ("-n", "--newlines"):
            if not value in ("LF", "CR", "CRLF", "orig"):
                sys.stderr.write("Invalid syntax for --newlines option: '%s'\n"
                                 % value)
                return ("exit", 1)
            else:
                params["newlines"] = value
        elif option == "--config-file":
            # Special option, was handled earlier
            pass
        else:
            # The options (such as --help) that cause immediate exit
            # were already checked, and caused the function to return.
            # Therefore, if we are here, it can't be due to any of these
            # options.
            raise ProgramError("Unexpected option received from the "
                               "getopt module: '%s'" % option)

    # Check for mutually exclusive options
    if (params["delay"] is not None) and (params["affine"] is not None):
        sys.stderr.write("Incompatible options: --delay and "
                         "(--affine-adjustment or --affine)\n")
        return ("exit", 1)

    if len(args) == 0:
        params["input file"] = sys.stdin
    elif len(args) == 1:
        params["input file"] = open(args[0], "rU")
    else:
        raise ProgramError()

    return ("continue", params)
    

def main():
    if sys.hexversion < 0x02030000:
        sys.exit("This program requires a Python version greater than or "
                 "equal to 2.3, but you\nare using:\n\n%s\n\nAborting."
                 % sys.version)

    try:
        action, params = process_command_line_and_config_file()
        if action == "exit":
            sys.exit(params)

        locale.setlocale(locale.LC_ALL, '')

        for negative_or_null_indices in next_subtitle(params):
            pass

        params["input file"].close()
        params["output file"].close()

        if negative_or_null_indices > 0:
            sys.stderr.write("Warning: %u non-positive indices " \
                             "were written.\n" % negative_or_null_indices)

        sys.exit(0)

    except error, exc_instance:
        sys.stderr.write("Error: %s\n" % exc_instance.complete_message())
        sys.exit(2)
    except (IOError, OSError), exc_instance:
        sys.stderr.write("Error: %s\n" % exc_instance)
        sys.exit(2)
    # I prefer letting other exceptions generate a traceback, because they
    # would probably reveal a bug (besides being more useful that way).

if __name__ == "__main__": main()