#!/usr/bin/env python3

import hashlib
import os
import pprint
import subprocess
import sys
import time
import traceback


def update_image(imagename, srcfilename):

    tags = {
        'Author': 'SkyWater PDK Authors',
        'Copyright': 'SPDX-License-Identifier: Apache-2.0',
    }
    if os.path.exists(srcfilename):
        sha1sum = hashlib.sha1(open(srcfilename, 'rb').read()).hexdigest()
        source_cmt = "{1}  {0}".format(os.path.basename(srcfilename), sha1sum)
        tags['Source'] = source_cmt

    tmpfile = imagename+'.tmp.png'
    if os.path.exists(tmpfile):
        os.unlink(tmpfile)

    cmd = ['exiftool']
    cmd.append('-o')
    cmd.append(tmpfile)
    for tag, v in tags.items():
        cmd.append('-{}="{}"'.format(tag, v))
    cmd.append(imagename)
    success = False
    for i in range(0, 10):
        p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output, _ = p.communicate()
        output = output.decode('utf-8')
        if p.returncode != 0:
            if 'Temporary file already exists' in output:
                print(time.now(), "Temporary file collision (trying again):", imagename, flush=True)
                time.sleep(1)
                continue
            msg = 'Failed to update metadata for {}'.format(imagename)
            sys.stderr.write(msg)
            sys.stderr.write('\n')
            sys.stderr.write('-'*75)
            sys.stderr.write('\n')
            sys.stderr.write(output)
            sys.stderr.write('-'*75)
            sys.stderr.write('\n')
            sys.stderr.flush()
            raise IOError(msg)
        success = True
        break
    if not success:
        msg = 'Unknown error, failed to update metadata for {}'.format(imagename)
        sys.stderr.write(msg)
        sys.stderr.write('\n')
        sys.stderr.flush()
        raise IOError(msg)
    os.rename(tmpfile, imagename)


def update_schematic(pathname):
    assert pathname.endswith("sch.png"), pathname
    # sky130_fd_sc_hs__clkbuf.sch.png
    # sky130_fd_sc_hs__clkbuf.simple.v
    vsrc = pathname.replace('.sch.png', '.full.v')
    update_image(pathname, vsrc)


def update_module(pathname):
    # sky130_fd_sc_hs__clkbuf.module.png
    # sky130_fd_sc_hs__clkbuf.blackbox.v
    vsrc = pathname.replace('.module.png', '.full.v')
    update_image(pathname, vsrc)


def main(argv):
    errors = []
    for f in argv:
        pathname = os.path.abspath(f)
        assert os.path.exists(pathname), pathname
        try:
            if pathname.endswith("sch.png"):
                update_schematic(pathname)
            elif pathname.endswith("module.png"):
                update_module(pathname)
            else:
                raise ValueError("Can't process {}".format(pathname))
            print("Added image metadata to", pathname, flush=True)
        except Exception as e:
            msg = traceback.format_exc()
            print('Failure while processing:', pathname, file=sys.stderr)
            sys.stderr.write(msg)
            sys.stderr.flush()
            errors.append(msg)
    if errors:
        return len(errors)
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except Exception as e:
        import traceback
        traceback.print_exc()
        sys.exit(1)