from PIL import Image
import json
import html
import re

re_param_code = r'\s*([\w ]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)'
re_param = re.compile(re_param_code)
re_imagesize = re.compile(r"^(\d+)x(\d+)$")
IGNORED_INFO_KEYS = {
    'jfif', 'jfif_version', 'jfif_unit', 'jfif_density', 'dpi', 'exif',
    'loop', 'background', 'timestamp', 'duration', 'progressive', 'progression',
    'icc_profile', 'chromaticity', 'photoshop',
}

def plaintext_to_html(text, classname=None):
    content = "<br>\n".join(html.escape(x) for x in text.split('\n'))

    return f"<p class='{classname}'>{content}</p>" if classname else f"<p>{content}</p>"

def to_digit(v):
    try:
        return float(v)
    except:
        return v
    
def read_info_from_image(image) -> tuple[str | None, dict]:
    if image is None:
        return '', {}
    elif type(image) == str:
        image = Image.open(image)
    items = (image.info or {}).copy()
    info =''

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

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

        geninfo = f"""{items["Description"]}
Negative prompt: {json_info["uc"]}
Steps: {json_info["steps"]}, Sampler: {json_info['sampler']}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""

        items = {**{'parameters': geninfo}, **items}

        for key, text in items.items():
            info += f"""
    <div>
    <p><b>{plaintext_to_html(str(key))}</b></p>
    <p>{plaintext_to_html(str(text))}</p>
    </div>
    """.strip()+"\n"

        if len(info) == 0:
            message = "Nothing found in the image."
            info = f"<div><p>{message}<p></div>"
    elif 'parameters' in items:
        res = {}

        prompt = ""
        negative_prompt = ""

        done_with_prompt = False

        *lines, lastline = items['parameters'].strip().split("\n")
        if len(re_param.findall(lastline)) < 3:
            lines.append(lastline)
            lastline = ''

        for line in lines:
            line = line.strip()
            if line.startswith("Negative prompt:"):
                done_with_prompt = True
                line = line[16:].strip()
            if done_with_prompt:
                negative_prompt += ("" if negative_prompt == "" else "\n") + line
            else:
                prompt += ("" if prompt == "" else "\n") + line

        res["prompt"] = prompt
        res["negative prompt"] = negative_prompt

        for k, v in re_param.findall(lastline):
            try:
                if v[0] == '"' and v[-1] == '"':
                    v = to_digit(json.loads(v))

                m = re_imagesize.match(v)
                if m is not None:
                    res[f"{k.lower()}-1"] = to_digit(m.group(1))
                    res[f"{k.lower()}-2"] = to_digit(m.group(2))
                else:
                    res[k.lower()] = v
            except Exception:
                print(f"Error parsing \"{k}: {v}\"")

        return items['parameters'], res
    else:
        return info, {}
    return info, json.loads(items['Comment'])


def send_paras(*args):
    if len(args[0]) == 0:
        return args[1:]
    items, prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale = args
    paras = [prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale]
    for i, keys in enumerate([('prompt',), ('',), ('uc', 'negative prompt'), ('seed',), ('scale', 'cfg scale'), ('width', 'size-1'), ('height', 'size-2'), ('steps',), ('sampler',), ('noise_schedule',), ('sm',), ('sm_dyn',), ('dynamic_thresholding',), ('cfg_rescale',)]):
        for key in keys:
            if key in items:
                paras[i] = items[key]
    prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale = paras
    if prompt.endswith(', ' + quality_tags):
        prompt = prompt[:-(len(quality_tags) + 2)]
    return prompt, quality_tags, neg_prompt, seed, scale, width, height, steps, sampler, scheduler, smea, dyn, dyn_threshold, cfg_rescale