import streamlit as st
import os
import zipfile
import csv
from PIL import Image, ImageFilter
import base64
import datetime
from gradio_client import Client  # インポートを追加

if 'blur_option' not in st.session_state:
    st.session_state.blur_option = True  

# Initialize Gradio client
client_nsfw = Client("https://ozoneasai-falconsai-nsfw-image-detection.hf.space/--replicas/hrcrr/")

# File uploader interface
uploaded_files = st.file_uploader("画像ファイルをアップロードしてください", type=["jpg", "jpeg", "png"], accept_multiple_files=True)

# Initialize pagination
if 'page' not in st.session_state:
    st.session_state.page = 1

# Initialize blur option as a toggle
if 'blur_option' not in st.session_state:
    st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", value=st.session_state.blur_option, key="blur_toggle")

# Initialize sort options
if 'sort_option' not in st.session_state:
    st.session_state.sort_option = {"column": "timestamp", "ascending": True}

# Function to generate a download link for a file
def get_binary_file_downloader_html(file_path, label="Download"):
    with open(file_path, 'rb') as f:
        data = f.read()
    b64 = base64.b64encode(data).decode()
    href = f'<a href="data:file/csv;base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
    return href

# Function to get Gradio NSFW prediction
def get_gradio_nsfw_prediction(image_path):
    result = client_nsfw.predict(image_path, api_name="/predict")
    return result[0]["label"] if result and result[0] and "label" in result[0] else "unknown"

# Function to save uploaded files and get Gradio predictions
def save_uploaded_files(uploaded_files):
    if not os.path.exists("temp"):
        os.makedirs("temp")
    
    for uploaded_file in uploaded_files:
        file_path = os.path.join("temp", uploaded_file.name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        
        # Get Gradio predictions
        gradio_nsfw_prediction = get_gradio_nsfw_prediction(file_path)
        
        # Get timestamp
        timestamp = datetime.datetime.now()

        # Append the information to the CSV file
        with open("temp/index.csv", mode='a', newline='', encoding='utf-8') as csv_file:
            csv_writer = csv.writer(csv_file)
            csv_writer.writerow([file_path, gradio_nsfw_prediction, timestamp])

# Function to apply blur to an image
def apply_blur(image_path):
    img = Image.open(image_path)
    img = img.filter(ImageFilter.GaussianBlur(radius=5))
    return img

# Function to paginate files
def paginate_files(files, page, files_per_page):
    start_index = (page - 1) * files_per_page
    end_index = start_index + files_per_page
    return files[start_index:end_index]

# Function to display images with predictions and options
def display_images(images, rows):
    for i, file_path in enumerate(images):
        file_ext = os.path.splitext(file_path)[1].lower()

        # Check if index is within the range
        if os.path.exists("temp/index.csv") and i < len(rows):
            gradio_nsfw_prediction = rows[i][1]
            timestamp = rows[i][2]

            st.write(f"**Prediction for {os.path.basename(file_path)} (NSFW):** {gradio_nsfw_prediction}")
            st.write(f"**Timestamp:** {timestamp}")

            if st.session_state.blur_option and gradio_nsfw_prediction.lower() == "nsfw":
                blurred_img = apply_blur(file_path)
                st.image(blurred_img, caption=os.path.basename(file_path), use_column_width=True)
            else:
                if file_ext in [".jpg", ".png"]:
                    st.image(Image.open(file_path), caption=os.path.basename(file_path), use_column_width=True)
                col1, col2 = st.columns([4, 1])
                with col1:
                    if file_ext in [".jpg", ".png"]:
                        st.image(Image.open(file_path), caption=os.path.basename(file_path), use_column_width=True)
                with col2:
                    if col2.button("削除", key=file_path):
                        os.remove(file_path)
                        # Update the displayed images after deletion
                        st.experimental_rerun()

# Main part of the app
if uploaded_files:
    save_uploaded_files(uploaded_files)

files_per_page = 20
num_pages = (len(os.listdir("temp")) - 1) // files_per_page + 1

col1, col2 = st.columns(2)
with col1:
    if st.button("前へ", key="prev_page"):
        st.session_state.page = max(1, st.session_state.page - 1)
with col2:
    if st.button("次へ", key="next_page"):
        st.session_state.page = min(num_pages, st.session_state.page + 1)

selected_page = st.number_input(
    "移動するページを指定してください",
    min_value=1,
    max_value=num_pages,
    value=max(1, min(st.session_state.page, num_pages)),  # Ensure the value is within the specified range
    key="selected_page"
)
if selected_page != st.session_state.page:
    st.session_state.page = selected_page

st.write(f"現在のページ: {st.session_state.page}")

# Retrieve the list of images from the CSV file
with open("temp/index.csv", mode='r', newline='', encoding='utf-8') as csv_file:
    csv_reader = csv.reader(csv_file)
    rows = list(csv_reader)

# Display images with predictions and options
current_page_files = paginate_files([row[0] for row in rows], st.session_state.page, files_per_page)
st.write("ファイル一覧:")
display_images(current_page_files, rows)

# Blurのトグルを更新


st.session_state.blur_option = st.checkbox("NSFW画像にBlurをかける", value=st.session_state.blur_option)

# CSVに保存ボタン
if st.button("CSVに保存"):
    csv_filename = "gradio_predictions_and_timestamps.csv"
    with open(csv_filename, mode='w', newline='', encoding='utf-8') as csv_file:
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(['ファイル名', 'NSFWの予測', '作成時間'])
        for row in rows:
            csv_writer.writerow([row[0], row[1], row[2]])

    st.markdown(get_binary_file_downloader_html(csv_filename, label="CSVをダウンロード"), unsafe_allow_html=True)

# ファイルを一括ダウンロード
if st.button("ファイルを一括ダウンロード"):
    zip_filename = "files.zip"
    with zipfile.ZipFile(zip_filename, "w") as zipf:
        for file_path in [row[0] for row in rows]:
            zipf.write(file_path, os.path.basename(file_path))

    st.markdown(get_binary_file_downloader_html(zip_filename, label="Zipファイルをダウンロード"), unsafe_allow_html=True)

st.success("ファイルがアップロードされました。このページのURLを他のクライアントと共有してください。")