import gradio as gr
import numpy as np
from PIL import Image

from utils.preprocess import preprocess_image
from utils.fill_regions import fill_regions
from utils.sharpen import sharpen_image
from utils.validation import quality_report
from utils.image_io import save_png
from utils.image_io import save_bmp


##################################################
# MAIN AGENT PIPELINE
##################################################

def generate_design(image):

    if image is None:
        return None, None, None, None

    # PIL -> numpy
    img = np.array(image)

    # Step 1 : Preprocess
    binary = preprocess_image(img)

    # Step 2 : Fill regions
    filled = fill_regions(binary)

    # Step 3 : Restore black boundaries
    # filled = overlay_boundaries(
    #   filled,
    #   binary
    #)

    # Step 4 : Sharpen
    final_image = sharpen_image(
        filled
    )

    # Step 5 : Validation
    report = quality_report(
        binary,
        final_image
    )

    report_text = (
        f"Regions Detected : {report['regions']}\n"
        f"Unfilled Pixels : {report['unfilled_pixels']}\n"
        f"Quality Score : {report['quality_score']} %"
    )

    # Step 6 : Save files
    png_file = save_png(
        final_image
    )

    bmp_file = save_bmp(
        final_image
    )

    return (
        Image.fromarray(final_image),
        report_text,
        png_file,
        bmp_file
    )


##################################################
# EXAMPLES
##################################################

examples = []

##################################################
# UI
##################################################

with gr.Blocks(
    title="Saree Design AI Agent"
) as demo:

    gr.Markdown(
        """
        # 🎨 Saree Design AI Agent

        Convert Step2 loom design into Step3
        designer-ready colour filled image.

        Standard Colours:
        - Black → Boundaries
        - Red → Filled Regions
        - White → Background
        """
    )

    with gr.Row():

        input_image = gr.Image(
            type="pil",
            label="Step2 Image"
        )

        output_image = gr.Image(
            label="Generated Step3"
        )

    with gr.Row():

        report = gr.Textbox(
            label="Quality Report",
            lines=5
        )

    with gr.Row():

        png_file = gr.File(
            label="Download PNG"
        )

        bmp_file = gr.File(
            label="Download BMP"
        )

    generate_btn = gr.Button(
        "Generate Step3 Image"
    )

    generate_btn.click(
        fn=generate_design,
        inputs=input_image,
        outputs=[
            output_image,
            report,
            png_file,
            bmp_file
        ]
    )

    if len(examples) > 0:
        gr.Examples(
            examples=examples,
            inputs=input_image
        )

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