import pyperclip
import os
import io
from PIL import Image
from sd_parsers import ParserManager
import gradio as gr
import re
import json


parser_manager = ParserManager()

TITLE = "<h1><center>Workflow Info Reader - By Andy N Le 0908 23 11 81 </center></h1>\n"

def extract_lora_info(prompt):
    lora_pattern = r"<lora:(.*?):(.*?)>"
    loras = re.findall(lora_pattern, prompt)
    if not loras:
        return "Không xác định"
    lora_info = [f" {name} - Weight: {weight}" for name, weight in loras]
    return "\n".join(lora_info)

def generate_workflow_from_metadata(metadata):
    nodes = []
    connections = []

    if 'nodes' in metadata:
        for node in metadata['nodes']:
            nodes.append({
                'id': node['id'],
                'title': node.get('title', 'Unnamed Node'),
                'pos': node.get('pos', [0, 0]),
                'type': node.get('type', 'basic/node')
            })
    
    if 'connections' in metadata:
        for conn in metadata['connections']:
            connections.append({
                'from': conn['from'],
                'to': conn['to']
            })
    
    return {
        'nodes': nodes,
        'connections': connections
    }

def format_sampler_params(sampler_params):
    try:
        params = json.loads(sampler_params)
    except (json.JSONDecodeError, TypeError):
        return "Không xác định"
    
    scheduler = params.get("scheduler", "Không xác định")
    cfg_scale = params.get("cfg_scale", "Không xác định")
    steps = params.get("steps", "Không xác định")
    return f"scheduler: {scheduler}\nCFG: {cfg_scale}\nSteps: {steps}"

download_file = gr.File(label="Tải file đã tạo", visible=False)

def save_metadata_to_file(image_input, prompt, neg_prompt, model, loras, seed, sampler, sampler_params, metadata):
    if not any([prompt, neg_prompt, model, loras, seed, sampler, sampler_params, metadata]):
        return None, gr.Info("Không có thông tin để lưu!", duration=2)
    if image_input is None:
        return None, gr.Info("Không thể lưu tệp: Không có ảnh nào được tải lên.", duration=2)
    
    # Lấy tên file ảnh và đặt tên cho file txt
    image_filename = os.path.basename(image_input)
    txt_filename = os.path.splitext(image_filename)[0] + ".txt"
    txt_filepath = os.path.join("outputs", txt_filename)

    # Ghi thông tin metadata vào file txt
    with open(txt_filepath, "w", encoding="utf-8") as f:
        f.write(f"Prompt: {prompt}\n")
        f.write(f"Negative Prompt: {neg_prompt}\n")
        f.write(f"Model: {model}\n")
        f.write(f"Loras: {loras}\n")
        f.write(f"Seed: {seed}\n")
        f.write(f"Sampler: {sampler}\n")
        f.write(f"Sampler Parameters: {sampler_params}\n")
        f.write(f"Other Metadata:\n{metadata}\n")
    
    return txt_filepath, gr.Info(f"Đã lưu thông tin vào file: {txt_filename}", duration=2)


    
def read_image_metadata(image_path):
    try:
        if not image_path:
            raise ValueError("Không có ảnh được tải lên.")
        
        with Image.open(image_path) as img:
            prompt_info = parser_manager.parse(img)
            
            if not prompt_info:
                raise ValueError("Không thể đọc thông tin từ ảnh.")

            # Kiểm tra loại generator
            if prompt_info.generator == "AUTOMATIC1111":
                return handle_automatic1111(prompt_info)
            elif prompt_info.generator == "ComfyUI":
                return handle_comfyui(prompt_info)
            else:
                raise ValueError("Loại generator không được hỗ trợ.")
                
    except Exception as e:
        print(f"Lỗi khi xử lý ảnh: {str(e)}")  
        return "Không tìm thấy thông tin!", "", "", "", "", "", "", ""

def handle_automatic1111(prompt_info):
    prompt = prompt_info.full_prompt
    negative_prompt = prompt_info.full_negative_prompt if prompt_info.full_negative_prompt else "N/A"
    
    models_list = list(prompt_info.models) if isinstance(prompt_info.models, set) else prompt_info.models
    model = models_list[0].name if models_list else "Không xác định"
    
    loras = extract_lora_info(prompt)
    
    samplers_list = list(prompt_info.samplers) if isinstance(prompt_info.samplers, set) else prompt_info.samplers
    if samplers_list and len(samplers_list) > 0:
        sampler = samplers_list[0].name if samplers_list[0].name else "N/A"
        sampler_params = json.dumps(samplers_list[0].parameters) if samplers_list[0].parameters else "N/A"
        seed = samplers_list[0].parameters.get('seed', 'Không xác định')
    else:
        sampler = "N/A"
        sampler_params = "N/A"
        seed = "Không xác định"

    formatted_sampler_params = format_sampler_params(sampler_params)

    other_metadata = "\n".join([
        f"{key}: {value}" for key, value in prompt_info.metadata.items() 
        if isinstance(key, str) and not key.startswith("Module")
    ])
    
    return prompt, negative_prompt, model, loras, seed, sampler, formatted_sampler_params, other_metadata

def format_metadata(metadata):
    formatted_output = ""
    for key, value in metadata.items():
        node_title, node_id = key
        formatted_output += f"Node '{node_title}'\n"
        for prop_key, prop_value in value.items():
            formatted_output += f"    {prop_key}: {prop_value}\n"
    return formatted_output

def handle_comfyui(prompt_info):
    prompt = prompt_info.full_prompt or "Không xác định"
    negative_prompt = prompt_info.full_negative_prompt if prompt_info.full_negative_prompt else "N/A"
    
    models_list = list(prompt_info.models) if isinstance(prompt_info.models, set) else prompt_info.models
    model = models_list[0].name if models_list else "Không xác định"
    
    samplers_list = list(prompt_info.samplers) if isinstance(prompt_info.samplers, set) else prompt_info.samplers
    if samplers_list and len(samplers_list) > 0:
        sampler = samplers_list[0].name if samplers_list[0].name else "N/A"
        sampler_params = json.dumps(samplers_list[0].parameters) if samplers_list[0].parameters else "N/A"
        seed = samplers_list[0].parameters.get('seed', 'Không xác định')
    else:
        sampler = "N/A"
        sampler_params = "N/A"
        seed = "Không xác định"

    formatted_sampler_params = format_sampler_params(sampler_params)

    
    converted_metadata = {}
    for key, value in prompt_info.metadata.items():
        if isinstance(key, tuple):
            key = str(key)
        converted_metadata[key] = value

    other_metadata = format_metadata(prompt_info.metadata)
    
    if not other_metadata:
        other_metadata = "Không có metadata bổ sung."

    return prompt, negative_prompt, model, None, seed, sampler, formatted_sampler_params, other_metadata

output_dir = "outputs"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)

def copy_to_clipboard(prompt, neg_prompt, seed, copy_prompt, copy_neg_prompt, copy_seed):
    copied_text = ""
    if copy_prompt:
        copied_text += f"Prompt: {prompt}\n"
    if copy_neg_prompt:
        copied_text += f"Negative Prompt: {neg_prompt}\n"
    if copy_seed:
        copied_text += f"Seed: {seed}\n"

    if copied_text:
        try:
            pyperclip.copy(copied_text)
            return gr.Info("Sao chép thành công!", duration=2)
        except pyperclip.PyperclipException:
            return gr.Info("Không hỗ trợ tính năng này trong bản demo. Vui lòng liên hệ tác giả để biết chi tiết !", duration=2)
    else:
        return gr.Info("Không có gì để sao chép!", duration=2)


# Hàm hủy bỏ sao chép
def cancel_copy():
    try:
        pyperclip.copy("")  # Xóa clipboard
        return (
            gr.update(value=False),  # Uncheck all checkboxes
            gr.update(value=False),
            gr.update(value=False),
            gr.Info("Đã hủy bỏ sao chép!", duration=2)
        )
    except pyperclip.PyperclipException:
        # Nếu không thể sử dụng clipboard trên hệ thống, hiển thị thông báo
        return (
            gr.update(value=False),  # Uncheck all checkboxes
            gr.update(value=False),
            gr.update(value=False),
            gr.Info("Hệ thống không hỗ trợ chức năng sao chép tự động.", duration=2)
        )


output_dir = "outputs"
if not os.path.exists(output_dir):
    os.mkdir(output_dir)





def check_image_size(image_input):
    try:
        if not image_input:
            raise ValueError("Không có ảnh được tải lên.")
        
        with Image.open(image_input) as img:
            width, height = img.size
            if width > 5000 or height > 5000:
                raise ValueError("Kích thước ảnh vượt quá 5000 px ở chiều ngang hoặc chiều dọc.")
            return image_input, gr.Info("Ảnh hợp lệ.", duration=1)
    except Exception as e:
        return None, gr.Info(f"Lỗi khi xử lý ảnh: {str(e)}", duration=2)


def gradio_interface(image_input):
    try:
        # Đọc dữ liệu từ ảnh
        prompt, negative_prompt, model, loras, seed, sampler, formatted_sampler_params, other_metadata = read_image_metadata(image_input)
        
        # Kiểm tra xem dữ liệu có bị thiếu không cho từng trường hợp
        if prompt == "Không thể đọc thông tin" or not any([prompt, negative_prompt, model, seed]):
            raise ValueError("Thiếu dữ liệu khi xử lý ảnh.")

        return prompt, negative_prompt, model, loras, seed, sampler, formatted_sampler_params, other_metadata
    except Exception as e:
        # In ra lỗi và trả về nội dung lỗi
        print(f"Lỗi trong gradio_interface: {str(e)}")
        return "Không tìm thấy thông tin ", "", "", "", "", "", "", ""


js_func = """
            function refresh() {
                const url = new URL(window.location);

                if (url.searchParams.get('__theme') !== 'dark') {
                    url.searchParams.set('__theme', 'dark');
                    window.location.href = url.href;
                }
            }

            """
with gr.Blocks(js = js_func) as demo:
    gr.HTML(TITLE)
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Tải lên hình ảnh")
            read_button = gr.Button("Đọc thông tin")
            copy_prompt = gr.Checkbox(label="Sao chép lời mô tả", value=False)
            copy_neg_prompt = gr.Checkbox(label="Sao chép mô tả loại trừ", value=False)
            copy_seed = gr.Checkbox(label="Sao Chép Seed", value=False)
            with gr.Row():
                copy_button = gr.Button("Sao chép")
                cancel_button = gr.Button("Hủy sao chép")
                download_button = gr.Button("Lưu file .txt")
            message_output = gr.HTML()
    
        
        with gr.Column(): 
            prompt_output = gr.Textbox(label="Lời mô tả (prompt)")
            negative_prompt_output = gr.Textbox(label="Mô tả loại trừ (Negative Prompt)")
            model_output = gr.Textbox(label="Mô hình (Model)")
            lora_output = gr.Textbox(label="Lora (Tên & Trọng số)")
            seed_output = gr.Textbox(label="Seed")
            sampler_output = gr.Textbox(label="Phương pháp lấy mẫu")
            sampler_params_output = gr.Textbox(label="Thông số lấy mẫu")
            other_metadata_output = gr.Textbox(label="Thông tin khác", lines=10)
            download_file = gr.File(label="Tải file đã tạo", visible=False)
            
    
    read_button.click(
        fn=check_image_size,
        inputs=image_input,
        outputs=[image_input, message_output],
        show_progress=False
    )

    # Sau khi kiểm tra kích thước, nếu hợp lệ, đọc metadata từ ảnh
    read_button.click(
        fn=gradio_interface,
        inputs=image_input,
        outputs=[prompt_output, negative_prompt_output, model_output, lora_output, seed_output, sampler_output, sampler_params_output, other_metadata_output]
    )
    copy_button.click(
        copy_to_clipboard, 
        inputs=None, 
        outputs=message_output,
        show_progress=False
    )

    # Gắn sự kiện click cho nút hủy sao chép
    cancel_button.click(
        cancel_copy, 
        inputs=None, 
        outputs=message_output,
        show_progress=False
    )
    def notify_download():
        return gr.Info("Không hỗ trợ tính năng này trong bản demo. Vui lòng liên hệ tác giả để biết chi tiết!", duration=3)

    # Gắn sự kiện click cho nút download
    download_button.click(
        fn=notify_download,
        inputs=None,  # Không cần inputs
        outputs=message_output,  # Chỉ hiển thị thông báo
        show_progress=False
    )






if __name__ == "__main__":
    demo.launch(inbrowser=True)