import gradio as gr
import shutil
import zipfile
from datetime import datetime
from pathlib import Path
from PIL import Image
import pillow_heif
import tempfile
import os

# Setup base folders
BASE_DIR = Path(tempfile.gettempdir()) / "17eMIT"
UPLOAD_FOLDER = BASE_DIR / "images"
FILES_FOLDER = BASE_DIR / "files"
UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
FILES_FOLDER.mkdir(parents=True, exist_ok=True)

# Allowed file types
IMAGE_EXTENSIONS = ['.png', '.jpg', '.jpeg', '.heic', '.webp', '.bmp', '.tiff']
FILE_EXTENSIONS = ['.mp3', '.m4a', '.mp4', '.pdf', '.docx', '.txt', '.xlsx', '.pptx']

def get_image_paths(limit=30):
    files = sorted(UPLOAD_FOLDER.glob("*"), key=lambda x: x.stat().st_mtime, reverse=True)
    return [str(f) for f in files if f.suffix.lower() in IMAGE_EXTENSIONS][:limit]

def get_file_list():
    files = sorted(FILES_FOLDER.glob("*"), key=lambda x: x.stat().st_mtime, reverse=True)
    return [f.name for f in files]

def save_files(file_paths):
    if not file_paths:
        return "❗請選擇至少一個檔案", get_image_paths(), gr.update(choices=get_file_list())
    
    messages = []
    for file_path in file_paths:
        try:
            ext = Path(file_path).suffix.lower()
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S%f")
            filename = f"{timestamp}{ext}"

            if ext in IMAGE_EXTENSIONS:
                if ext == ".heic":
                    heif_file = pillow_heif.read_heif(file_path)
                    image = Image.frombytes(heif_file.mode, heif_file.size, heif_file.data, "raw")
                else:
                    image = Image.open(file_path).convert("RGB")

                filepath = UPLOAD_FOLDER / f"{timestamp}.png"
                image.save(filepath, format="PNG")
            elif ext in FILE_EXTENSIONS:
                filepath = FILES_FOLDER / filename
                shutil.copy(file_path, filepath)
            else:
                messages.append(f"❌ 不支援的檔案格式: {ext}")
                continue

            messages.append(f"✅ 檔案已上傳: {filename}")
        except Exception as e:
            messages.append(f"❌ 上傳失敗 ({file_path}): {str(e)}")
    
    return "\n".join(messages), get_image_paths(), gr.update(choices=get_file_list())

def clear_images():
    try:
        shutil.rmtree(UPLOAD_FOLDER)
        shutil.rmtree(FILES_FOLDER)
        UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
        FILES_FOLDER.mkdir(parents=True, exist_ok=True)
        return "🧹 所有圖片與檔案已清除", [], gr.update(choices=[])
    except Exception as e:
        return f"❌ 清除失敗: {str(e)}", get_image_paths(), gr.update(choices=get_file_list())

def download_all_files():
    all_files = list(UPLOAD_FOLDER.glob("*")) + list(FILES_FOLDER.glob("*"))
    if not all_files:
        return None, "⚠️ 沒有可下載的檔案"

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    temp_zip = Path(tempfile.gettempdir()) / f"all_uploads_{timestamp}.zip"

    try:
        with zipfile.ZipFile(temp_zip, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for filepath in all_files:
                zipf.write(filepath, arcname=filepath.name)
        return str(temp_zip), "✅ ZIP 文件已生成"
    except Exception as e:
        return None, f"❌ ZIP 文件生成失敗: {str(e)}"

def cleanup_old_zips():
    temp_dir = Path(tempfile.gettempdir())
    for zip_file in temp_dir.glob("all_uploads_*.zip"):
        try:
            zip_file.unlink()
        except Exception:
            pass

# --- UI Starts ---
with gr.Blocks(title="📁 圖片與檔案分享平台") as demo:

    with gr.Row():
        file_input = gr.File(label="上傳檔案", type="filepath", file_count="multiple")
        upload_button = gr.Button("📤 上傳", variant="primary")
        output_msg = gr.Textbox(label="狀態", lines=4, interactive=False)

    with gr.Row():
        clear_button = gr.Button("🚫 清除所有", variant="stop")
        download_button = gr.Button("📦 下載全部", variant="primary")
        download_file = gr.File(label="點此下載 ZIP")
        zip_path_display = gr.Textbox(label="ZIP 文件訊息", lines=2, interactive=False)

    gr.Markdown("### 📄 上傳檔案列表")
    file_list = gr.Dropdown(label="所有檔案", choices=get_file_list(), interactive=False)

    gallery = gr.Gallery(label="📷 已上傳圖片", show_label=True, columns=3, object_fit="contain", height="500px")

    upload_button.click(
        fn=save_files,
        inputs=file_input,
        outputs=[output_msg, gallery, file_list]
    )

    clear_button.click(
        fn=clear_images,
        outputs=[output_msg, gallery, file_list]
    )

    download_button.click(
        fn=download_all_files,
        outputs=[download_file, zip_path_display]
    )

    demo.load(fn=get_image_paths, outputs=gallery)
    demo.load(fn=cleanup_old_zips)

demo.launch()