import gradio as gr
from PIL import Image
import requests
from io import BytesIO
import json
import html

css = """
.html-container span {
    color: var(--color-accent);
}
.json-node {
    --string-color: #b5bd68;
    --separator-color: var(--block-border-color);
}
.json-node .null {
    color: var(--body-text-color-subdued);
}
.json-node .bool {
    color: #cc99cc;
}
"""

NL = "\n"

def parse_key_value_pairs(input_string):
    pairs = []
    key = None
    value = ""
    brace_count = 0
    in_string = False
    escaping = False

    for i, char in enumerate(input_string):
        if char == '"' and not escaping:
            in_string = not in_string
        elif char == "\\" and not escaping:
            escaping = True
        else:
            escaping = False

        if char == "{" and not in_string:
            brace_count += 1
        elif char == "}" and not in_string:
            brace_count -= 1

        if (char == "," and brace_count == 0 and not in_string) or i == len(input_string) - 1:
            if i == len(input_string) - 1:
                value += char
            if key is not None and value:
                pairs.append((key, value))
            elif key is None and value and pairs:
                prev_key, prev_value = pairs.pop()
                pairs.append((prev_key, f"{prev_value},{value}"))
            key, value = None, ""
        elif char == ":" and brace_count == 0 and not in_string:
            key, value = value, ""
        else:
            value += char

    return [(k.strip(), v.strip()) for k, v in pairs]

def process_webui(info_data):
    lines = info_data.split("\n")
    positive_prompt = []
    negative_prompt = []
    settings_params = []
    current_section = "positive"

    for line in lines:
        if current_section == "positive" and line.startswith("Negative prompt: "):
            current_section = "negative"
            line = line[len("Negative prompt: "):]
        elif current_section != "settings" and line.startswith("Steps: "):
            current_section = "settings"

        if current_section == "positive":
            positive_prompt.append(line)
        elif current_section == "negative":
            negative_prompt.append(line)
        elif current_section == "settings":
            settings_params.append(line)

    html_data = f"""
<p><span>Prompt:</span> {html.escape(NL.join(positive_prompt)).replace(NL, '<br>')}</p>
<p><span>Negative prompt:</span> {html.escape(NL.join(negative_prompt)).replace(NL, '<br>')}</p>
"""

    for key, value in parse_key_value_pairs("\n".join(settings_params)):
        html_data += f"<p><span>{html.escape(key).replace(NL, '<br>')}:</span> {html.escape(value).replace(NL, '<br>')}</p>"

    return html_data

def process_nai(info_data):
    comment = json.loads(info_data["Comment"])
    html_data = f"""
<p><span>Prompt:</span> {html.escape(comment.get('prompt', '')).replace(NL, '<br>')}</p>
<p><span>Undesired Content:</span> {html.escape(comment.get('uc', '')).replace(NL, '<br>')}</p>
<p><span>Resolution:</span> {comment.get('width', '')}x{comment.get('height', '')}</p>
<p><span>Seed:</span> {comment.get('seed', '')}</p>
<p><span>Steps:</span> {comment.get('steps', '')}</p>
<p><span>Sampler:</span> {comment.get('sampler', '')} ({comment.get('noise_schedule', '')})</p>
<p><span>Prompt Guidance:</span> {comment.get('scale', '')}</p>
<p><span>Prompt Guidance Rescale:</span> {comment.get('cfg_rescale', '')}</p>
<p><span>Undesired Content Strength:</span> {comment.get('uncond_scale', '')}</p>
<p><span>Request Type:</span> {comment.get('request_type', '')}</p>
<p style="margin-top: 16px;"><span>Software:</span> {info_data.get('Software', '')}</p>
<p><span>Source:</span> {info_data.get('Source', '')}</p>
<p><span>Title:</span> {info_data.get('Title', '')}</p>
<p><span>Generation Time:</span> {info_data.get('Generation time', info_data.get('Generation_time', ''))}</p>
"""

    return html_data, comment

def process_image(file=None, url=None):
    try:
        if url:
            response = requests.get(url)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content))
        elif file:
            img = file
        else:
            return (
                "No input provided.",
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=False)
            )

        info_data = img.info

        if info_data:
            if "parameters" in info_data:
                parameters = info_data["parameters"]
                html_data = process_webui(parameters)
                return (
                    html_data,
                    gr.update(visible=True, value=f"<p>{html.escape(parameters).replace(NL, '<br>')}</p>"),
                    gr.update(visible=False),
                    gr.update(visible=True, value=img) if url else gr.update(visible=False)
                )
            elif "Comment" in info_data and info_data.get("Software") == "NovelAI":
                simple, raw = process_nai(info_data)
                return (
                    simple,
                    gr.update(visible=False),
                    gr.update(visible=True, value=raw),
                    gr.update(visible=True, value=img) if url else gr.update(visible=False)
                )
            else:
                return (
                    "No metadata found.",
                    gr.update(visible=False),
                    gr.update(visible=False),
                    gr.update(visible=True, value=img) if url else gr.update(visible=False)
                )
        else:
            return (
                "No metadata found.",
                gr.update(visible=False),
                gr.update(visible=False),
                gr.update(visible=True, value=img) if url else gr.update(visible=False)
            )
    except Exception as e:
        return (
            f"Error: {str(e)}",
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=False)
        )

with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Upload an Image or Provide a URL")
            input_url = gr.Textbox(label="Source Image URL")
            btn = gr.Button("Submit", variant="primary")
            input_img = gr.Image(label="Source Image", type="pil", sources="upload", height=240)
        with gr.Column():
            gr.Markdown("### Image Info")
            output_img = gr.Image(label="URL Image", interactive=False, visible=False, height=240)
            with gr.Tabs():
                with gr.Tab("Metadata"):
                    output_html = gr.HTML()
                with gr.Tab("Raw Parameters"):
                    output_json = gr.JSON(container=False, visible=False)
                    output_params = gr.HTML(visible=False)

    input_img.change(
        fn=process_image,
        inputs=input_img,
        outputs=[output_html, output_params, output_json, output_img],
        api_name="interrogate"
    )

    btn.click(process_image, inputs=[input_img, input_url], outputs=[output_html, output_params, output_json, output_img])

demo.launch()