"""
This is the main file for the streamlit app.
Author: Sungjin Kim
Date: 2023-05-19
"""
import os
from typing import Union
from datetime import datetime

import streamlit as st
import requests
from PIL import Image
from langchain.callbacks import get_openai_callback
import openai

def download_image(image_url:str, 
                   fname:str=None, 
                   default_fold:str="images/") -> Union[str, None]:
    """
    download image from url
    """
    try:
        image = requests.get(image_url, timeout=5)
        if fname:
            name_ext = fname
        else:
            name_ext = image_url.split("/")[-1]
        name_ext = default_fold + name_ext
        with open(name_ext, "wb") as f:
        # with open(name, "wb") as f:
            f.write(image.content)
        return name_ext
    except Exception as e_name:
        print("An error occured", e_name)
        return None

def save_rgba_image(uploaded_file_name:str, uploaded_file_buffer:object,
                    default_fold:str="images/"):
    """
    # Save the uploaded file to a specific location
    """
    file_path = os.path.join(default_fold[:-1], uploaded_file_name)
    with open(file_path, "wb") as file:
        file.write(uploaded_file_buffer)
    image_file = Image.open(file_path)
    image_file_rgba = image_file.convert("RGBA")
    image_file_rgba.save(file_path)
    return file_path

def edit_image(img_name:str,
    prompt:str="change backgroud to a beautiful sunset over the ocean")-> str:
    """
    code is generated by being
    im_size : one of ['256x256', '512x512', '1024x1024'] - 'size'
    """
    response = openai.Image.create_edit(
        image=open(img_name, "rb"),
        prompt=prompt,
        size="256x256",
        n=1,
        response_format="url"
    )
    img_url = response["data"][0]["url"]
    return img_url

def main():
    """
    Main function for the streamlit app.
    """
    default_fold = "images/"
    # load_dotenv()
    st.set_page_config(page_title="Avatar Generation")
    st.header("Artistic Avatar Generation 💖")

    api_key = st.text_input("Enter your OpenAI API key:", type="password")
    openai.api_key = api_key

    mode = st.radio("Choose your option", ("Upload photo by file", "Upload photo by url"))
    photo_filename = None

    if mode == "Upload photo by file":
        uploaded_file = st.file_uploader("Upload your photo", type=["png"])
        if uploaded_file is not None:
            st.image(uploaded_file, width=256, caption="Your photo")
            # Save the uploaded file to a specific location
            photo_filename = save_rgba_image(uploaded_file.name, uploaded_file.getbuffer(),
                                             default_fold=default_fold)
    else:
        photo_url = st.text_input("Enter your photo url:",
                              value="https://pngimg.com/uploads/monkey/monkey_PNG18720.png") 
        if photo_url:
            st.image(photo_url, width=256, caption="Your photo")
            photo_filename = download_image(photo_url, default_fold=default_fold)
            # e.g., photo_url = "images/monkey_PNG18720.png"

    if photo_filename:
        image = Image.open(photo_filename)
        im_size = f"{image.size[0]}x{image.size[1]}" # width x height
        st.write(f"Your photo filename: {photo_filename} with size of {im_size}")
        prompt = st.text_input("Prompt to transform to avatar: ",
                                value="Change background to a beautiful sunset over the ocean")
        if prompt:
            with get_openai_callback() as oai_call_back:
                # img_url = gen_image(prompt)
                img_url = edit_image(photo_filename, prompt)
                print(oai_call_back)
            print(f"Avatar image url: {img_url}")
            st.image(img_url, width=256, caption=f"prompt: {prompt}")
            fname = f"{prompt}_{datetime.now().strftime('%Y%m%d%H%M%S')}.png"
            output_filename = download_image(img_url, fname=fname)
            if output_filename:
                st.write(f"Your avatar filename: {output_filename}")
                st.write("Done!")

if __name__ == '__main__':
    main()