#!/usr/bin/python3
#
# Python wrapper script for cptutils, combines the component
# programs to allow for arbitrary conversion (if possible)
#
# Copyright (c) J.J. Green 2012, 2014, 2016, 2020

from __future__ import annotations
from typing import (
    List, Dict, Tuple, Any, Generator, Union, Sequence, Optional
)

import os
import sys
import tempfile
import subprocess
import atexit

from zipfile import ZipFile
from argparse import ArgumentParser, Action, Namespace

# version string, set by sed on build (see Makefile)

version: str = '1.82'

# list of files & directories to be deleted at exit

delfiles: List[str] = []
deldirs: List[str] = []

# data on gradient types, this is not used directly but later
# sliced into per-column dicts

gradient_data: Dict[
    str,
    Tuple[str, List[str], str, bool, Optional[List[int]]]
] = {
    'c3g': (
        'CSS3 gradient',
        ['css3'],
        'c3g',
        False,
        None
    ),
    'cpt': (
        'GMT colour palette table',
        [],
        'cpt',
        False,
        None
    ),
    'ggr': (
        'GIMP gradient',
        [],
        'ggr',
        False,
        [c for c in b'GIMP Grad']
    ),
    'gpf': (
        'Gnuplot palette',
        [],
        'gpf',
        False,
        None
    ),
    'gpl': (
        'GIMP palette',
        [],
        'gpl',
        False,
        [c for c in b'GIMP Pale']
    ),
    'grd': (
        'Photoshop gradient',
        ['grd5'],
        'grd',
        True,
        [c for c in b'8BGR\x00\x05']
    ),
    'inc': (
        'POV-Ray header',
        ['pov'],
        'inc',
        False,
        None
    ),
    'lut': (
        'Medcon lookup table',
        [],
        'lut',
        False,
        None
    ),
    'map': (
        'Tecplot map',
        [],
        'map',
        False,
        [c for c in b'#!MC 1410']
    ),
    'pg': (
        'PostGIS colour map',
        [],
        'pg',
        False,
        None
    ),
    'png': (
        'PNG image',
        [],
        'png',
        False,
        [c for c in b'\x89PNG']
    ),
    'psp': (
        'PaintShop Pro gradient',
        ['grd3', 'jgd'],
        'PspGradient',
        False,
        [c for c in b'8BGR\x00\x03']
    ),
    'qgs': (
        'QGIS style colour-ramp',
        [],
        'qgs',
        True,
        [c for c in b'<!DOCTYPE qgis']
    ),
    'sao': (
        'DS9/SAO colour table',
        [],
        'sao',
        False,
        None
    ),
    'svg': (
        'SVG gradient',
        [],
        'svg',
        True,
        [c for c in b'<?xml ']
    )
}

# conversion adjacency matrix, implemented as a dict-of-dicts
# with the entries being the conversion programs

gradient_matrix: Dict[str, Dict[str, str]] = {
    'gpl': {
        'cpt': 'gplcpt'
    },
    'cpt': {
        'svg': 'cptsvg',
        'pg': 'cptpg'
    },
    'svg': {
        'cpt': 'svgcpt',
        'c3g': 'svgcss3',
        'ggr': 'svggimp',
        'gpf': 'svggpt',
        'inc': 'svgpov',
        'map': 'svgmap',
        'pg': 'svgpg',
        'png': 'svgpng',
        'psp': 'svgpsp',
        'qgs': 'svgqgs',
        'sao': 'svgsao'

    },
    'ggr': {
        'svg': 'gimpsvg',
        'lut': 'gimplut'
    },
    'grd': {
        'svg': 'pssvg'
    },
    'qgs': {
        'svg': 'qgssvg'
    },
    'psp': {
        'svg': 'pspsvg'
    }
}


def data_slice(i: int) -> Dict[str, Any]:
    '''
    dicts for type -> names, aliases, and so on from the
    main data-dict
    '''
    return {
        k: v[i]
        for k, v in gradient_data.items()
    }


gradient_names: Dict[str, str] = data_slice(0)
gradient_aliases: Dict[str, List[str]] = data_slice(1)
gradient_exts: Dict[str, str] = data_slice(2)
gradient_burstable: Dict[str, bool] = data_slice(3)
gradient_magic: Dict[str, Optional[List[int]]] = data_slice(4)

# gradient_types dict from alias list, this used to specify types
# (eg -i psp) and to guess file-types from extensions.

gradient_types: Dict[str, str] = {
    k: gtype
    for gtype, galiases in gradient_aliases.items()
    for k in (gtype, gradient_exts[gtype], *galiases)
}

# conversion (di)graph from the adjacency matrix

gradient_graph: Dict[str, List[str]] = {
    k: list(v.keys())
    for k, v in gradient_matrix.items()
}

# list of used programs

programs: List[str] = [
    program
    for d in gradient_matrix.values()
    for program in d.values()
]

programs.append('svgsvg')


def rwformats(
        M: Dict[str, Any],
        N: Dict[str, Any]
) -> Tuple[Dict[str, bool], Dict[str, bool]]:
    '''
    Dicts of type -> readable, writeable
    '''

    rfmt: Dict[str, bool] = {}
    wfmt: Dict[str, bool] = {}

    for name in N.keys():
        rfmt[name] = False
        wfmt[name] = False

    for nr in M.keys():
        rfmt[nr] = True
        for nw in M[nr].keys():
            wfmt[nw] = True

    return (rfmt, wfmt)


def formats_supported(M: Dict[str, Any], N: Dict[str, Any]) -> None:
    '''
    Print the formats supported to stdout from the adjacency matrix
    M and the name -> description dict N
    '''

    rfmt, wfmt = rwformats(M, N)

    print('supported formats:')

    for name in sorted(N.keys()):
        rc = 'R' if rfmt[name] else '-'
        wc = 'W' if wfmt[name] else '-'
        print(f'| {name:3} | {N[name]:25} | {rc}{wc} |')


def capabilities() -> None:
    '''
    Print the formats supported to YAML from the adjacency matrix
    M and the name -> description dict N
    '''

    rfmt, wfmt = rwformats(gradient_matrix, gradient_names)

    def quoted(s: str) -> str:
        return '"' + s + '"'

    def boolstring(val: bool) -> str:
        return 'true' if val else 'false'

    print(f'# gradient-convert {version} capabilities')
    for name in sorted(gradient_names.keys()):
        print(f'{name}:')
        print(f'  read: {boolstring(rfmt[name])}')
        print(f'  write: {boolstring(wfmt[name])}')
        print(f'  burst: {boolstring(gradient_burstable[name])}')
        print(f'  desc: {quoted(gradient_names[name])}')
        print(f'  alias: [{", ".join(map(quoted, gradient_aliases[name]))}]')
        print(f'  ext: {quoted(gradient_exts[name])}')
        magic = gradient_magic[name]
        if magic is None:
            print(f'  magic: null')
        else:
            print(f'  magic: [{", ".join(map(str, magic))}]')

def graphviz() -> None:
    '''
    Print the conversion graph in GraphViz dot format
    '''
    print('# gradient-convert %s output' % version)
    print('')
    print('digraph {')
    for src, dsts in gradient_matrix.items():
        print('  %s -> {%s}' % (src, ' '.join(dsts)))
    print('}')

def shortest_path(
        graph: Dict[str, List[str]],
        start: str,
        end: str,
        path: List[str] | None = None
) -> List[str] | None:
    '''
    A simple shortest path code determines the call sequence, taken from
    https://legacy.python.org/doc/essays/graphs.html
    '''
    if path is None:
        path = []
    path = path + [start]
    if start == end:
        return path
    if start not in graph:
        return None
    shortest = None
    for node in graph[start]:
        if node not in path:
            newpath = shortest_path(graph, node, end, path)
            if newpath and (not shortest or len(newpath) < len(shortest)):
                shortest = newpath
    return shortest


def type_from_path(path: str) -> str:
    '''
    Gradient type from path
    '''
    ext: str | None = os.path.splitext(path)[1][1:]
    if not ext:
        print(f'cannot determine file extension for {path}')
        sys.exit(1)

    gradient_type: str | None = gradient_types.get(ext)
    if gradient_type is None:
        print(f'unknown gradient extension {ext}')
        formats_supported(gradient_matrix, gradient_names)
        sys.exit(1)

    return gradient_type


def run_clist(clist: List[str], topath: str | None, verbose: bool) -> bool:
    '''
    Run a list of commands which implement the conversion
    '''
    if verbose:
        print(f'  {" ".join(clist)}')
    if subprocess.call(clist) != 0:
        print(f'failed call to {clist[0]}: aborting')
        return False
    if topath and not os.path.exists(topath):
        print(f'failed to create {topath}: aborting')
        return False
    return True


def convert(
        ipath: str,
        opath: str,
        opt: Dict[str, Any]
) -> bool:
    '''
    The main conversion routine
    '''

    global delfiles
    global deldirs

    # for the intermediate filenames; use the basenam of the
    # input file, but make the file location in a tmpname()
    # directory (so that we won't stomp on users' local data)

    tempdir: str = tempfile.mkdtemp()
    deldirs.append(tempdir)

    # the call-list

    clist: List[str]

    # Here we handle the multiple-gradient files, although this
    # is a bit convoluted it is much less messy than my attempt
    # to do this in a generic fashion; we'd not expect to support
    # many other multi-gradient formats in any case

    if opt['zipped']:

        # first call with burst to create the files in tempdir

        opt['burst'] = True
        opt['zipped'] = False

        if not convert(ipath, tempdir, opt):
            return False

        # number of file created by burst

        nburst: int = len([f for f in os.listdir(tempdir)])

        if nburst == 0:
            print('nothing to zip')
            return False

        # now zip the results

        if opt['verbose']:
            print('creating zipfile')

        arcdir: str = os.path.splitext(os.path.split(opath)[1])[0]
        zf: ZipFile = ZipFile(opath, mode='w')

        try:
            for f in os.listdir(tempdir):
                input_path: str = f'{tempdir}/{f}'
                arc_path: str = f'{arcdir}/{f}'
                zf.write(input_path, arcname=arc_path)
                delfiles.append(input_path)
        finally:
            zf.close()

        return True

    elif opt['burst']:

        # basename used in the title, so make it meaningful,
        # we don't take it from the output (as that will ba a
        # directory, and may well be ".")

        basename: str = os.path.splitext(os.path.split(ipath)[1])[0]
        svgmulti: str

        if opt['ifmt'] == 'grd':

            # input is a single grd file, convert it to a single
            # svg file with muliple gradients; then call convert()
            # with that file as input and burst = True (so that
            # we execute the ifmt == 'svg' case below).
            # counting the gradient to reduce the redundant zeros

            svgmulti = f'{tempdir}/{basename}.svg'
            clist = (
                ['pssvg'] +
                opt['backtrace'] +
                ['--output', svgmulti] +
                [ipath]
            )
            if opt['verbose']:
                print(f'  {" ".join(clist)}')
            if subprocess.call(clist) != 0:
                print(f'failed call to {clist[0]}: aborting')
                return False
            delfiles.append(svgmulti)

            opt['ifmt'] = 'svg'
            opt['burst'] = True

            return convert(svgmulti, opath, opt)

        elif opt['ifmt'] == 'qgs':

            # pretty-much the same as for grd

            svgmulti = f'{tempdir}/{basename}.svg'
            clist = (
                ['qgssvg'] +
                opt['backtrace'] +
                ['--output', svgmulti] +
                [ipath]
            )
            if opt['verbose']:
                print(f'  {" ".join(clist)}')
            if subprocess.call(clist) != 0:
                print(f'failed call to {clist[0]}: aborting')
                return False
            delfiles.append(svgmulti)

            opt['ifmt'] = 'svg'
            opt['burst'] = True

            return convert(svgmulti, opath, opt)

        elif opt['ifmt'] == 'svg':

            # input is a single svg file (which may be from the
            # case above, or an original infile).

            if opt['ofmt'] == 'svg':

                # final output is svg, burst to output directory

                clist = (
                    ['svgsvg'] +
                    opt['backtrace'] +
                    ['--output', opath] +
                    ['--all'] +
                    [ipath]
                )
                clist.extend(opt['subopts']['svgsvg'])
                if not run_clist(clist, opath, opt['verbose']):
                    return False
            else:

                # final output is not svg, so burst into a temp
                # directory, then call convert() on each file in
                # that directory, this time with burst = False

                svgdir: str = f'{tempdir}/{basename}'
                os.mkdir(svgdir)
                deldirs.append(svgdir)
                clist = (
                    ['svgsvg'] +
                    opt['backtrace'] +
                    ['--output', svgdir] +
                    ['--all'] +
                    [ipath]
                )
                if not run_clist(clist, None, opt['verbose']):
                    return False
                opt['ifmt'] = 'svg'
                opt['burst'] = False
                for svg in os.listdir(svgdir):
                    svgbase: str = os.path.splitext(svg)[0]
                    oext: str = gradient_exts[opt["ofmt"]]
                    ipath2: str = f'{svgdir}/{svg}'
                    opath2: str = f'{opath}/{svgbase}.{oext}'
                    delfiles.append(ipath2)
                    if not convert(ipath2, opath2, opt):
                        return False
            return True

    # basename for non-burst is only used in temporary files

    basename = os.path.splitext(os.path.split(opath)[1])[0]

    # trap non-conversions, note that this does not affect 'bursting'
    # svg to svg (that is handled above)

    if opt['ifmt'] == opt['ofmt']:
        print(f'converting {opt["ifmt"]} to {opt["ofmt"]} seems, pointless?')
        return False

    def pairs(L: Any) -> Generator[Tuple[Any, Any], None, None]:
        '''
        Helper method which generates consecutive pairs from
        anything iterable
        '''
        i: Any = iter(L)
        prev: Any = next(i)
        item: Any = prev
        for item in i:
            yield prev, item
            prev = item

    # create the system-call sequence, first we create
    # a list of dictionaries of call data

    callpath: List[str] | None = shortest_path(
        gradient_graph,
        opt['ifmt'],
        opt['ofmt']
    )

    if callpath is None:
        print(f'cannot convert {opt["ifmt"]} to {opt["ofmt"]} yet')
        formats_supported(gradient_matrix, gradient_names)
        return False

    cdlist: List[Dict[str, str]] = [
        {
            'fromtype': t0,
            'totype': t1,
            'program': gradient_matrix[t0][t1]
        }
        for t0, t1 in pairs(callpath)
    ]

    # add input/output filenames

    cdlist[0]['frompath'] = ipath
    cdlist[-1]['topath'] = opath

    # add temporary filenames (also added to the global
    # delfiles list used by the cleanup function)

    for cd0, cd1 in pairs(cdlist):
        totype: str = cd0['totype']
        path: str = f'{tempdir}/{basename}.{gradient_exts[totype]}'
        cd0['topath'] = path
        cd1['frompath'] = path
        delfiles.append(path)

    # now run through the call data and make the calls

    for cd in cdlist:

        program: str = cd['program']
        topath: str = cd['topath']
        frompath: str = cd['frompath']

        clist = (
            [program] +
            opt['backtrace'] +
            ['--output', topath] +
            opt['subopts'][program] +
            [frompath]
        )

        if not run_clist(clist, topath, opt['verbose']):
            return False

    return True


def cleanup_dirs(verbose: bool) -> None:
    global deldirs
    for path in reversed(deldirs):
        if verbose:
            print(path)
        os.rmdir(path)


def cleanup_files(verbose: bool) -> None:
    global delfiles
    for path in delfiles:
        if verbose:
            print(path)
        if os.path.isfile(path):
            os.unlink(path)


class CapAction(Action):
    def __call__(
            self,
            parser: ArgumentParser,
            namespace: Namespace,
            values: Union[str, Sequence[Any], None],
            opts: Optional[str] = None
    ) -> None:
        capabilities()
        sys.exit(0)


class GraphVizAction(Action):
    def __call__(
            self,
            parser: ArgumentParser,
            namespace: Namespace,
            values: Union[str, Sequence[Any], None],
            opts: Optional[str] = None
    ) -> None:
        graphviz()
        sys.exit(0)


def main() -> None:
    description = 'Convert colour gradients to other formats'
    epilog = 'For details, see gradient-convert(1)'
    parser = ArgumentParser(
        description=description,
        epilog=epilog
    )
    parser.add_argument(
        '-b', '--background',
        metavar='RGB',
        help='set background (cpt)'
    )
    parser.add_argument(
        '--backtrace-file',
        metavar='PATH',
        dest='backtrace_path',
        help='write backtrace to PATH on error'
    )
    parser.add_argument(
        '--backtrace-format',
        metavar='FORMAT',
        help='format of backtrace'
    )
    parser.add_argument(
        '-B', '--burst',
        action='store_true',
        dest='burst',
        help='burst multiple gradients'
    )
    parser.add_argument(
        '-c', '--capabilities',
        action=CapAction,
        nargs=0,
        help='print program capabilites in YAML format'
    )
    parser.add_argument(
        '-f', '--foreground',
        metavar='RGB',
        help='set foreground (cpt)'
    )
    parser.add_argument(
        '-g', '--geometry',
        metavar='WxH',
        help='set geometry (png, svg)'
    )
    parser.add_argument(
        '-G', '--graphviz',
        action=GraphVizAction,
        nargs=0,
        help='print conversion graph in GraphViz dot format'
    )
    parser.add_argument(
        '-i', '--input-format',
        metavar='FORMAT',
        choices=gradient_types.values(),
        help='format of the input'
    )
    parser.add_argument(
        '-n', '--nan',
        metavar='RGB',
        help='set NaN colour (cpt)'
    )
    parser.add_argument(
        '-o', '--output-format',
        metavar='FORMAT',
        choices=gradient_types.values(),
        help='format of the output'
    )
    parser.add_argument(
        '-p', '--preview',
        action='store_true',
        dest='preview',
        help='add preview (svg)'
    )
    parser.add_argument(
        '-T', '--transparency',
        metavar='RGB',
        help='replace transparency (cpt, gpt, sao)'
    )
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        dest='verbose',
        help='print runtime information'
    )
    parser.add_argument(
        '-V', '--version',
        action='version',
        version=f'gradient-convert {version}',
        help='print version and exit'
    )
    parser.add_argument(
        '-z', '--zip',
        action='store_true',
        dest='zipped',
        help='create zip file of multiple outputs'
    )
    parser.add_argument(
        '-4', '--gmt4',
        action='store_true',
        dest='gmt4',
        help='use GMT4 format in output (cpt)'
    )
    parser.add_argument(
        '-5', '--gmt5',
        action='store_true',
        dest='gmt5',
        help='use GMT5 format in output (cpt)'
    )
    parser.add_argument(
        '-6', '--gmt6',
        action='store_true',
        dest='gmt6',
        help='use GMT6 format in output (cpt)'
    )
    parser.add_argument(
        'input_path',
        metavar='INPUT-PATH',
        help='path to input'
    )
    parser.add_argument(
        'output_path',
        metavar='OUTPUT-PATH',
        help='path to output'
    )
    opts = parser.parse_args()

    # the above extracts the options and handles --version, --help
    # and --capabilities; we now proceed as with the earlier version
    # of the code creating an 'opt' (note singular) and passing that
    # to the unmodified code above

    ipath: str = opts.input_path
    opath: str = opts.output_path
    verbose: bool = opts.verbose
    ifmt: str | None = opts.input_format
    ofmt: str | None = opts.output_format
    burst: bool = opts.burst
    zipped: bool = opts.zipped
    subopts: Dict[str, List[Any]] = dict((p, []) for p in programs)
    backtrace: List[str] = []

    # hello all

    if verbose:
        print(f'This is gradient-convert (version {version})')

    # progs_<x> accept the -<x> option

    progs_p: Tuple[str, ...] = ('cptsvg', 'gimpsvg', 'pspsvg', 'svgsvg')
    progs_g: Tuple[str, ...] = progs_p + ('svgpng',)
    progs_bfn: Tuple[str, ...] = ('svgcpt', 'gplcpt')
    progs_456: Tuple[str, ...] = ('svgcpt', 'gplcpt')
    progs_T: Tuple[str, ...] = ('svgcpt', 'svggpt', 'svgsao')

    # options to propagate to some programs in the call sequence

    if opts.preview:
        for prog in progs_p:
            subopts[prog].extend(['--preview'])

    if opts.geometry is not None:
        for prog in progs_g:
            subopts[prog].extend(['--geometry', opts.geometry])

    if opts.background is not None:
        for prog in progs_bfn:
            subopts[prog].extend(['--background', opts.background])

    if opts.foreground is not None:
        for prog in progs_bfn:
            subopts[prog].extend(['--foreground', opts.foreground])

    if opts.nan is not None:
        for prog in progs_bfn:
            subopts[prog].extend(['--nan', opts.nan])

    if opts.transparency:
        for prog in progs_T:
            subopts[prog].extend(['--transparency', opts.transparency])

    if opts.gmt4:
        for prog in progs_456:
            subopts[prog].extend(['--gmt4'])

    if opts.gmt5:
        for prog in progs_456:
            subopts[prog].extend(['--gmt5'])

    if opts.gmt6:
        for prog in progs_456:
            subopts[prog].extend(['--gmt6'])

    # options to propagate to all programs in the call sequence

    if opts.backtrace_path is not None:
        backtrace.extend(['--backtrace-file', opts.backtrace_path])

    if opts.backtrace_format is not None:
        backtrace.extend(['--backtrace-format', opts.backtrace_format])

    # formats

    if ifmt is None:
        ifmt = type_from_path(ipath)

    if ofmt is None:
        if burst or zipped:
            print('Output format must be specified (see -o option)')
            sys.exit(1)
        ofmt = type_from_path(opath)

    # exit handlers

    atexit.register(cleanup_dirs, False)
    atexit.register(cleanup_files, False)

    # ready to go

    if verbose:
        print(f'input: {gradient_names[ifmt]}')
        print(f'  {ipath}')
        print(f'output: {gradient_names[ofmt]}')
        print(f'  {opath}')
        print('call sequence:')

    opt: Dict[str, Any] = {
        'verbose': verbose,
        'subopts': subopts,
        'ifmt': ifmt,
        'ofmt': ofmt,
        'burst': burst,
        'zipped': zipped,
        'backtrace': backtrace
    }

    success: bool = convert(ipath, opath, opt)

    if success:
        if verbose:
            print('done.')
        sys.exit(0)
    else:
        if verbose:
            print('failed.')
        sys.exit(1)


if __name__ == '__main__':
    main()
