import pandas as pd
import streamlit as st
from PIL import Image
# from streamlit_drawable_canvas import st_canvas
import replicate
from pathlib import Path
# Import the 'replicate' module, which likely contains functions for running AI models or tasks.
import replicate
import urllib.request
from pathlib import Path
import os
from PIL import Image


st.title("StyleSnap")

with st.sidebar:
    with st.form('Replicate'):
        # User selects the model (OpenAI/Cohere) and enters API keys
        rep_key = st.text_input('Enter Replicate Key', type="password")
        submitted = st.form_submit_button("Submit")

if rep_key:
    os.environ["REPLICATE_API_TOKEN"] = rep_key


def resize_image(image, size):
    img = Image.open(image)
    img_resized = img.resize(size, Image.LANCZOS)
    return img_resized

def gen(path,prompt,clothing):
    output = replicate.run(
        "naklecha/fashion-ai:4e7916cc6ca0fe2e0e414c32033a378ff5d8879f209b1df30e824d6779403826",
        input={"image": open(path,'rb'),
            "clothing": clothing,
            "prompt": prompt}
    )
    urllib.request.urlretrieve(output[-1], "image.png")
    image = Image.open("image.png")
    st.image(image)
    image.save("inpainting.png")

clothing = [
    "topwear",
    "bottomwear"
]
clothing = st.selectbox("Select clothing", clothing)
prompt = st.text_input('Enter prompt')

uploaded_file = st.file_uploader(f"Upload image:", type=['png','jpg'])

if (st.button("Submit")):
    if uploaded_file is not None: 
        st.write("File uploaded successfully!")
        file_contents = uploaded_file.read()
        save_path = uploaded_file.name
        # image = Image.open(uploaded_file)
        resized_size = (512, 512)
        resized_image = resize_image(uploaded_file, resized_size)
        st.image(resized_image, caption="Uploaded Image", use_column_width=True)
        resized_image.save(save_path)
        gen(save_path,prompt,clothing)