from PIL import Image, ExifTags, PngImagePlugin
from sd_parsers import ParserManager
from io import BytesIO
from pathlib import Path
from os import path
import sys
import time
from typing import Union

#def format_metadata(params: dict) -> str:
def format_metadata(params):
    lines = []

    lines.append(params['prompt'])

    loras = params.get('lora', {})
    lorainfo=[]
    for name, props in sorted(loras.items()):
        weight = props.get('weight')
        weight = float(weight) if weight else 0.0
        if name and weight:
            lorainfo.append('<lora:{}:{}>'.format(name,weight))
    if lorainfo:
        lines.append(''.join(lorainfo))

    negative_prompt = params.get('negative_prompt')
    if negative_prompt is not None:
        lines.append('Negative prompt: ' + negative_prompt)

    info = []

    info.append(('Steps', params['num_inference_steps']))

    sampler = params.get('sampler')
    if sampler:
        info.append(('Sampler', sampler))

    info.append(('CFG Scale', params['guidance_scale']))

    seed = params.get('seed', -1)
    if seed != -1:
        info.append(('Seed', seed))

    width, height = params.get('width'), params.get('height')
    if width and height:
        info.append(('Size', '{}x{}'.format(width, height)))

    model = params.get('model')
    if model is not None:
        if model.endswith('.safetensors'):
            # filename without extension
            model = path.basename(model).rsplit('.', 1)[0]
        # else assume it's a model id
        info.append(('Model', model))

    mode = params.get('mode') # Img2Img, Txt2Img
    if mode is not None:
        info.append(('Mode', mode))

        if mode == 'Img2Img':
            strength = params.get('strength')
            if strength is not None:
                info.append(('Denoising strength', strength))

    clip_skip = params.get('clip_skip', 0)
    if clip_skip >= 1:
        info.append(('Clip skip', clip_skip+1))

    controlnet_model = params.get('cnet_model', '')
    controlnet_conditioning_scale = params.get('controlnet_conditioning_scale', 0.0)
    if controlnet_model and controlnet_conditioning_scale >= 0.001:
        info.append(('ControlNet', controlnet_conditioning_scale))
        info.append(('ControlNet Conditioning Scale', controlnet_conditioning_scale))
        if params.get('guess_mode'):
            info.append(('ControlNet Guess Mode', 1))

    vae = params.get('vae')
    if vae is not None:
        info.append(('VAE', vae))

    for (opt, meta) in [
        ('compel', 'Compel'),
        ('freeu', 'FreeU'),
        ('resadapter_norm', 'ResAdapter Normalization'),
        ('vae_tiling', 'VAE Tiling'),
        ]:
        if params.get(opt):
            info.append((meta, 1))

    for (opt, meta) in [
        ('eta', 'Eta'),
        ('pag_adaptive_scale', 'PAG Adaptive Scale'),
        ('pag_scale', 'PAG Scale'),
        ('sag_scale', 'SAG Scale'),
        ('token_merging_ratio', 'Token Merging Ratio'),
        ]:
        val = params.get(opt, 0.0)
        if val and val >= 0.01:
            info.append((meta, val))

    lines.append(', '.join(['{}: {}'.format(k, v) for (k, v) in info]))

    return '\n'.join(lines)

#def add_metadata_to_pil_image(pil_image: Image, params: Union[dict, str])-> None:
def add_metadata_to_pil_image(pil_image, params):
    '''add generation parameters to the image info fields, in a Gradio-compatible way'''

    if isinstance(params, str):
        metadata = params
    else:
        metadata = format_metadata(params)

    pil_image.info['parameters'] = metadata

    # borrowed from piexif
    usercomment = b'UNICODE\0' + metadata.encode('utf_16_be', errors='replace')
    # The PIL Exif encoder detects both bytes and bytearrays as sequences of
    # integers, so they get encoded with the wrong type, and most tools won't
    # interpret that as text. A list wrapping a bytearray dodges those
    # heuristics, correctly storing the data as a byte sequence.
    usercomment = [bytearray(usercomment)]
    exif = pil_image.getexif()
    exif.setdefault(ExifTags.IFD.Exif, {})[ExifTags.Base.UserComment] = usercomment
    pil_image.info['exif'] = exif.tobytes()

def get_image_meta_str(image):
    res = []
    if image.filename:
        res.append('Filename: {}'.format(image.filename))
    if image.format:
        res.append('Format: {}'.format(image.format))
    res.append('Size: {}x{}'.format(*image.size))
    info = ParserManager().parse(image)
    if info:
        for prompt in info.prompts:
            if prompt:
                res.append('Prompt: {}'.format(prompt))
        for negative_prompt in info.negative_prompts:
            if negative_prompt:
                res.append('Negative Prompt: {}'.format(negative_prompt))
        for sampler in info.samplers:
            if sampler.name:
                res.append('Sampler: {}'.format(sampler.name))
            seed = sampler.parameters.get('seed')
            if seed:
                res.append('Seed: {}'.format(seed))
            steps = sampler.parameters.get('steps')
            if steps:
                res.append('Steps: {}'.format(steps))
            if sampler.model and sampler.model.name:
                res.append('Model: {}'.format(sampler.model.name))
        for (meta, value) in sorted(info.metadata.items()):
            res.append('{}: {}'.format(meta, value))
    return '\n'.join(res)

#def extract_meta_for_saving(pil_image: Image):
def extract_meta_for_saving(pil_image):

    meta_args = {}

    exif = pil_image.info.get('exif')
    if exif:
        meta_args["exif"] = exif

    pngtexts = [ (key, value) for key, value in pil_image.info.items()
        if isinstance(key, str) and isinstance(value, str) ]
    if pngtexts:
       pnginfo = PngImagePlugin.PngInfo()
       for key, value in pngtexts:
           pnginfo.add_text(key, value)
       meta_args["pnginfo"] = pnginfo

    return meta_args

def filebytes(pil_image, format):
    with BytesIO() as output_bytes:
        meta_params = extract_meta_for_saving(pil_image)
        if format in ['jpeg', 'jpg', 'webp']:
            meta_params['quality'] = 90
        pil_image.save(output_bytes, format=format, **meta_params)
        return output_bytes.getvalue()

def open_file_timestamp(basename, directory='.', extension='.png'):
    output = Path(directory)
    count = 0

    opt = True
    spbasename = basename.rsplit('-', 1)
    if len(spbasename) == 2 and len(spbasename[1]) == 4 and spbasename[1].isdigit():
        count = int(spbasename[1])
        basename = spbasename[0]
        opt = False

    def getfilepath(count):
        extra = '-{:04}'.format(count)
        return output / '{}{}.{}'.format(basename, extra, extension)
    filepath = getfilepath(count)

    # otimizacao
    if opt:
        existing = sorted(output.glob(basename + '*.' + extension))
        if existing:
            count = 1
            last = existing[-1].name[:len(extension)]
            sepnum = (last.rsplit('-', 1)+[''])[1]
            if sepnum.isdigit():
                count = int(sepnum)
            filepath = getfilepath(count)

    while True:
        try:
            return open(filepath, 'xb')
        except FileExistsError:
            count += 1
            filepath = getfilepath(count)

def save_image_with_meta(pil_image, basename, directory, format="png"):
    filedata = filebytes(pil_image, format)
    with open_file_timestamp(basename, directory, format) as outfile:
        outfile.write(filedata)
        return outfile.name


def save_image_timestamp(pil_image, directory, format="png"):
    outdir = Path(directory)
    basename = time.strftime('%Y%m%d-0001')
    return save_image_with_meta(pil_image, basename, outdir, format)