import streamlit as st
import os
import pandas as pd
from datetime import datetime
from PIL import Image
from PIL import ImageFilter
from io import BytesIO
import zipfile
import base64
import deepdanbooru as dd
import huggingface_hub
import tensorflow as tf
import numpy as np

# ページごとの表示数
PAGE_SIZE = 20

# タグ付け関数
def predict_tags(image: Image.Image, score_threshold: float) -> tuple[dict[str, float], dict[str, float], str]:
    _, height, width, _ = model.input_shape
    image = np.asarray(image)
    image = tf.image.resize(image, size=(height, width), method=tf.image.ResizeMethod.AREA, preserve_aspect_ratio=True)
    image = image.numpy()
    image = dd.image.transform_and_pad_image(image, width, height)
    image = image / 255.
    probs = model.predict(image[None, ...])[0]
    probs = probs.astype(float)

    indices = np.argsort(probs)[::-1]
    result_all = dict()
    result_threshold = dict()
    for index in indices:
        label = labels[index]
        prob = probs[index]
        result_all[label] = prob
        if prob < score_threshold:
            break
        result_threshold[label] = prob
    result_text = ', '.join(result_all.keys())
    return result_threshold, result_all, result_text

# ページングと並び替えのためのデータを取得する関数
def load_data():
    index_df = pd.read_csv(index_file_path)
    return index_df

# アップロードされた写真を表示する関数
def display_photos(photos):
    for photo_info in photos.iterrows():
        row = photo_info[1]
        photo_path = os.path.join("photos", row["File Name"])
        image = Image.open(photo_path)

        # タグを予測して表示
        result_threshold, result_all, result_text = predict_tags(image, 0.7)

        st.image(image.filter(ImageFilter.BLUR), caption=row["File Name"], use_column_width=True)
        
        # タグを表示
        st.write("タグ:", ", ".join(result_all.keys()))

# フォルダーの中身を zip ファイルとしてダウンロード
def download_photos_as_zip(file_paths):
    # Zip ファイル作成
    with BytesIO() as zip_buffer:
        with zipfile.ZipFile(zip_buffer, "a", zipfile.ZIP_DEFLATED) as zip_file:
            for file_path in file_paths:
                zip_file.write(file_path, os.path.basename(file_path))
        
        # ダウンロードリンク表示
        st.markdown(
            f"**[ダウンロード ZIPファイル](data:application/zip;base64,{base64.b64encode(zip_buffer.getvalue()).decode()})**",
            unsafe_allow_html=True
        )

# モデルとラベルをダウンロードする関数
def load_model():
    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru', 'model-resnet_custom_v3.h5')
    model = tf.keras.models.load_model(path)
    return model

def load_labels():
    path = huggingface_hub.hf_hub_download('public-data/DeepDanbooru', 'tags.txt')
    with open(path) as f:
        labels = [line.strip() for line in f.readlines()]
    return labels

# Streamlit アプリケーションのメイン部分
def main():
    st.sidebar.title("アップロードオプション")
    uploaded_photos = st.sidebar.file_uploader("写真をアップロードしてください", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

    # モデルとラベルをダウンロードする
    global model, labels
    model = load_model()
    labels = load_labels()

    if uploaded_photos:
        for uploaded_photo in uploaded_photos:
            file_name = f"{datetime.now().strftime('%Y%m%d%H%M%S%f')}.png"
            save_uploaded_photo(uploaded_photo, file_name)

            # インデックスに追加
            global photo_df  # globalキーワードを追加
            photo_df = pd.concat([photo_df, pd.DataFrame([[os.path.basename(file_name), datetime.now(), ""]], columns=["File Name", "Timestamp", "Tags"])], ignore_index=True)

        # インデックスを更新
        photo_df.sort_values(by="Timestamp", inplace=True, ascending=False)
        photo_df.reset_index(drop=True, inplace=True)
        photo_df.to_csv(index_file_path, index=False)

    st.subheader("アップロードされた写真")
    
    # ページングと並び替えのためのデータを読み込む
    index_df = load_data()

    # ページング
    page_num = st.sidebar.number_input("ページ番号", value=1, min_value=1, max_value=(len(index_df) // PAGE_SIZE) + 1)
    start_idx = (page_num - 1) * PAGE_SIZE
    end_idx = min(start_idx + PAGE_SIZE, len(index_df))
    current_page = index_df.iloc[start_idx:end_idx]

    # 並び替え
    sort_options = ["TimeStamp 昇順", "TimeStamp 降順", "名前 昇順", "名前 降順"]
    selected_sort = st.sidebar.selectbox("並び替え", sort_options)
    
    if selected_sort == "TimeStamp 昇順":
        current_page = current_page.sort_values(by="Timestamp", ascending=True)
    elif selected_sort == "TimeStamp 降順":
        current_page = current_page.sort_values(by="Timestamp", ascending=False)
    elif selected_sort == "名前 昇順":
        current_page = current_page.sort_values(by="File Name", ascending=True)
    elif selected_sort == "名前 降順":
        current_page = current_page.sort_values(by="File Name", ascending=False)

    display_photos(current_page)

    # Next ボタン
    if st.button("Next"):
        page_num += 1
        st.experimental_rerun()

    # Previous ボタン
    if st.button("Previous"):
        page_num -= 1
        st.experimental_rerun()

    # Zip ダウンロードボタン
    if st.button("写真をダウンロード (ZIP)"):
        file_paths = [os.path.join(photos_folder, file) for file in os.listdir(photos_folder)]
        download_photos_as_zip(file_paths)

if __name__ == "__main__":
    main()