import io
import re
import sys
import json
import yaml
import argparse
import piexif
import piexif.helper
from PIL import Image
from rich.console import Console

console = Console()

IGNORED_INFO_KEYS = {
    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
    'icc_profile', 'chromaticity', 'photoshop',
}

def read_geninfo(image: Image.Image) -> tuple[str | None, dict]:
    items = (image.info or {}).copy()

    geninfo = items.pop('parameters', None)

    if "exif" in items:
        exif_data = items["exif"]
        try:
            exif = piexif.load(exif_data)
        except OSError:
            # memory / exif was not valid so piexif tried to read from a file
            exif = None
        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b'')
        try:
            exif_comment = piexif.helper.UserComment.load(exif_comment)
        except ValueError:
            exif_comment = exif_comment.decode('utf8', errors="ignore")

        if exif_comment:
            items['exif comment'] = exif_comment
            geninfo = exif_comment
    elif "comment" in items: # for gif
        geninfo = items["comment"].decode('utf8', errors="ignore")

    for field in IGNORED_INFO_KEYS:
        items.pop(field, None)

    if items.get("Software", None) == "NovelAI":
        try:
            json_info = json.loads(items["Comment"])
            sampler = json_info["sampler"]

            geninfo = f"""{items["Description"]}
Negative prompt: {json_info["uc"]}
Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""
        except Exception:
            console.print_exception("Error parsing NovelAI image generation parameters", exc_info=True)

    return geninfo, items

def parse_geninfo(geninfo):
    geninfo_lines = geninfo.splitlines()

    if geninfo_lines[0].startswith('Steps: '):
        remaining_geninfo_line = geninfo_lines[0]
    elif geninfo_lines[0].startswith('Negative prompt: '):
        remaining_geninfo_line = geninfo_lines[1:]
        geninfo_yaml_lines = [geninfo_lines[0]]
    else:
        geninfo_yaml_lines = [f'Positive prompt: {geninfo_lines[0]}']
        if geninfo_lines[1].startswith('Steps: '):
            remaining_geninfo_line = geninfo_lines[1]
        else:
            remaining_geninfo_line = geninfo_lines[2]
            geninfo_yaml_lines.append(geninfo_lines[1])

    if remaining_geninfo_line.startswith('Steps: '):
        unquoted_comma_regex = r'''(,)(?=(?:[^"]|"[^"]*")*$)'''
        remaining_geninfo_fields = re.split(unquoted_comma_regex, remaining_geninfo_line)

        for field in remaining_geninfo_fields:
            if field[-1] == ',':
                field = field[:-1]
            field = field.strip()
            if field:
                geninfo_yaml_lines.append(field)

    if len(geninfo_yaml_lines) == 0:
        console.print_exception(f'Error parsing geninfo:', geninfo)
        sys.exit(1)

    return '\n'.join(geninfo_yaml_lines)

def yaml_to_json(yaml_str):
    with io.StringIO(yaml_str) as f:
        yaml_obj = yaml.safe_load(f)

    return json.dumps(yaml_obj)

def main():
    parser = argparse.ArgumentParser(description='Python CLI for reading metadata embedded in images generated by Stable Diffusion.')
    parser.add_argument('-i', '--input', type=str, required=False, help='Path to input image from which to read metadata, alternatively you can pass the input path on stdin')
    parser.add_argument('-c', '--color', type=str, required=False, default='auto', help='Whether or not the JSON output should be printed in color, to pretty print the json result in color, defaults to "--color=auto" which will pretty print when stdout is detected as a terminal, pass "--color=never" to force printing in plain text, pass "--color=always" to force pretty printing')
    parser.add_argument('-m', '--minify', action='store_true', required=False, help='Whether or not to minify the json result, defaults to False')
    parser.add_argument('-f', '--format', type=str, required=False, default='json', choices=['json', 'yaml'], help='The format to output the metadata in, json or yaml, defaults to json')
    args = parser.parse_args()

    input_path = args.input
    if not input_path:
        if sys.stdin.readable():
            input_path = sys.stdin.readline().strip()
        else:
            console.print_exception('No input image path provided')
            sys.exit(1)

    input_image = Image.open(input_path)

    geninfo_str, additional_metadata = read_geninfo(input_image)

    geninfo_yaml = parse_geninfo(geninfo_str)

    highlight = None
    if args.color == 'always':
        highlight = True
    elif args.color == 'never':
        highlight = False

    if args.format == 'yaml':
        console.print(geninfo_yaml, highlight=highlight)
        sys.exit(0)

    geninfo_json = yaml_to_json(geninfo_yaml)

    output_obj = {
        'generation_metadata': json.loads(geninfo_json),
        'additional_metadata': additional_metadata
    }

    indent = 4
    if args.minify:
        indent = None

    output_str = json.dumps(output_obj, indent=indent)
    console.print(output_str, highlight=highlight)

if __name__ == '__main__':
    main()
