from PIL import Image
from utils import pipe_image
import streamlit as st
from streamlit_drawable_canvas import st_canvas
import gettext
import numpy as np

_ = gettext.gettext

#language = st.selectbox(_('Choose your language'), ['en', 'it'])
# try:
#   localizator = gettext.translation('base', localedir='locales', languages=[language])
#   localizator.install()
#   _ = localizator.gettext 
# except:
#     pass

st.markdown("<h1 style='text-align:center;'> Magic Board - CPU version</h1>", unsafe_allow_html=True)
# Specify canvas parameters in application
drawing_mode = st.sidebar.selectbox(
    _("Drawing tool:"), ("freedraw", "line", "rect", "circle", "transform")
)

stroke_width = st.sidebar.slider(_("Stroke width: "), 1, 25, 3)
stroke_color = st.sidebar.color_picker(_("Stroke color hex: "))
bg_color = st.sidebar.color_picker(_("Background color hex: "), "#eee")
bg_image = st.sidebar.file_uploader(_("Background image:"), type=["png", "jpg"])
strength = st.sidebar.slider(_("Strength value: "), 0.0, 1.0, 0.1)

# Create a canvas component
canvas_result = st_canvas(
    fill_color="rgba(255, 165, 0, 0.3)",  # Fixed fill color with some opacity
    stroke_width=stroke_width,
    stroke_color=stroke_color,
    background_color=bg_color,
    background_image=Image.open(bg_image) if bg_image else None,
    update_streamlit=True,
    height=450,
    width=680,
    drawing_mode=drawing_mode,
    #point_display_radius=point_display_radius if drawing_mode == 'point' else 0,
    key="canvas",
)

prompt = st.text_input(_('Prompt to generate your cool image'), 
               placeholder=_('write you text here to improve the image with your favourite style'))

if st.button('Submit'):
    img_size = (512,512)
    if bg_image:
        img_back=Image.open(bg_image).convert('RGB').resize(img_size)
        comp_img = img_back
    elif np.any(canvas_result.image_data):
        img_draw=Image.fromarray(canvas_result.image_data).convert('RGB').resize(img_size)
        comp_img = img_draw
    elif bg_image and np.any(canvas_result.image_data):
        img_back=Image.open(bg_image).convert('RGB').resize(img_size)
        img_draw=Image.fromarray(canvas_result.image_data).convert('RGB').resize(img_size)
        comp_img = Image.new("RGB", img_size)
        comp_img = Image.blend(img_draw, img_back, alpha=0.5)
    else:
        st.write("Upload an image or draw something")
    
    if prompt is None:
        st.write("No prompt, no magic! :(")

    with st.spinner('Wait for it... ~ 5 minutes '):
        images = pipe_image(prompt=prompt,
                       init_image=comp_img,
                       strength=strength)
    if images["nsfw_content_detected"]:
        st.write("NSFW content detected, retry with another prompt or image")
    else:
        st.image(images.images)

st.warning("If you don't see any image could be the NSFW content checker or a Memory error. "
            "Please, re-run the submission and modify the prompt. "
            "Taking care that, each word in the prompt result in much more token for the model"
            "which means much more memory usage.")