import streamlit as st

from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
import requests
from dashscope import ImageSynthesis
import os


prompt = st.text_area("请输入提示词")
img_dir = os.path.join(os.getcwd(), 'imgs')
if not os.path.exists(img_dir):
    os.makedirs(img_dir)

if "file_ready" not in st.session_state:
    st.session_state.file_ready = False

width = st.number_input("图片宽度", value=512, min_value=512, max_value=1440)
height = st.number_input("图片高度", value=512, min_value=512, max_value=1440)

# 512, 1440
if st.button("生成"):
    with st.spinner('生成中...',show_time=True):
        rsp = ImageSynthesis.call(api_key=os.getenv("DASHSCOPE_API_KEY"),
                                model="wanx2.1-t2i-turbo",
                                prompt=prompt,
                                n=1,
                                size=f'{width}*{height}')
        print('response: %s' % rsp)
        if rsp.status_code == HTTPStatus.OK:
            # 在当前目录下保存图片
            for result in rsp.output.results:
                file_name = PurePosixPath(unquote(urlparse(result.url).path)).parts[-1]
                img_path = os.path.join(img_dir, file_name)
                print(f'save image to {img_path}')
                with open(img_path, 'wb+') as f:
                    f.write(requests.get(result.url).content)
                st.image(img_path)
                st.write(result.actual_prompt)
                # st.session_state.can_download = True
                
                # with open(img_path, 'rb') as file:
                #     st.session_state.file_data = file.read()
                #     st.session_state.file_name = file_name
                #     st.session_state.file_ready = True
                
        else:
            print('sync_call Failed, status_code: %s, code: %s, message: %s' %
                (rsp.status_code, rsp.code, rsp.message))
            
# 下载按钮逻辑
# if st.session_state.file_ready:
#     st.download_button(
#         label="下载图片",
#         data=st.session_state.file_data,
#         file_name=st.session_state.file_name,
#         mime='image/jpeg'
#     )