app.py

import streamlit as st
from PIL import Image, UnidentifiedImageError
import io
import os
import zipfile

# --- 配置 ---
ICC_PROFILE_OPTIONS = {
    "原版青蛙": "qingwa.icc",
    "超亮版本": "6172.icc",
    "多兼容版本": "2020.icc",
}
DEFAULT_ICC_KEY = "原版青蛙"
SUPPORTED_FORMATS = ["jpg", "jpeg", "png", "tif", "tiff"]


# --- 辅助函数 ---

def get_icc_profile_bytes(selected_icc_filename):
    script_dir = os.path.dirname(__file__)
    icc_profile_path = os.path.join(script_dir, selected_icc_filename)
    if not os.path.exists(icc_profile_path):
        st.error(f"错误: ICC 配置文件 '{selected_icc_filename}' 未在脚本目录 '{script_dir}' 中找到。")
        return None
    try:
        with open(icc_profile_path, "rb") as f:
            return f.read()
    except Exception as e:
        st.error(f"读取 ICC 文件 '{selected_icc_filename}' 时出错: {e}")
        return None


def embed_icc_profile(image_bytes, icc_profile_bytes, original_filename):
    if icc_profile_bytes is None:
        st.error(f"由于无法加载ICC配置文件，未能处理图片 '{original_filename}'。")
        return None, None
    try:
        img = Image.open(io.BytesIO(image_bytes))
        img_format = img.format if img.format else "PNG"

        original_mode = img.mode
        exif_data = img.info.get("exif")

        if img.mode == 'P':
            has_transparency = "transparency" in img.info
            img = img.convert('RGBA' if has_transparency else 'RGB')

        output_image_bytes = io.BytesIO()
        params = {"icc_profile": icc_profile_bytes}
        if exif_data:
            params["exif"] = exif_data

        if img_format.upper() == "JPEG":
            params["quality"] = 95
            params["optimize"] = True

        try:
            img.save(output_image_bytes, format=img_format, **params)
        except Exception as e_save:
            st.warning(
                f"无法以格式 '{img_format}' (模式: {original_mode}) 保存 '{original_filename}' (原因: {e_save})。尝试转换为兼容模式并以PNG格式保存...")

            img_reloaded = Image.open(io.BytesIO(image_bytes))
            target_mode = 'RGB'
            if img_reloaded.mode in ['LA', 'RGBA', 'PA'] or \
                    (img_reloaded.mode == 'P' and "transparency" in img_reloaded.info):
                target_mode = 'RGBA'

            img_converted = img_reloaded
            if img_reloaded.mode != target_mode and img_reloaded.mode not in ['RGB', 'RGBA', 'L', 'CMYK', 'LAB']:
                st.caption(
                    f"提示: 将图像 '{original_filename}' 从 {img_reloaded.mode} 转换为 {target_mode} 以便用PNG保存。")
                img_converted = img_reloaded.convert(target_mode)
            elif img_reloaded.mode not in ['RGB', 'RGBA']:
                img_converted = img_reloaded.convert(target_mode)

            output_image_bytes = io.BytesIO()
            png_params = {"icc_profile": icc_profile_bytes}
            img_converted.save(output_image_bytes, format="PNG", **png_params)
            img_format = "PNG"

        output_image_bytes.seek(0)
        return output_image_bytes, img_format

    except UnidentifiedImageError:
        st.error(f"无法识别文件 '{original_filename}' 为有效图片。")
        return None, None
    except Exception as e:
        st.error(f"处理图片 '{original_filename}' 时发生错误: {e}")
        return None, None


# --- Streamlit 应用界面 ---
st.set_page_config(page_title="发光青蛙表情包 HDR 制作")
st.title("🐸 发光青蛙表情包 HDR 制作")

# ICC 文件选择 (使用单选按钮)
st.markdown("#### 1. 选择要嵌入的 ICC 配置文件：")
selected_icc_display_name = st.radio(
    label="选择 ICC Profile:",  # label 可以设为空字符串如果你觉得上面的 markdown 标题足够
    options=list(ICC_PROFILE_OPTIONS.keys()),
    index=list(ICC_PROFILE_OPTIONS.keys()).index(DEFAULT_ICC_KEY),  # 设置默认选项
    key="icc_radio_selector",
    horizontal=True,  # 让单选按钮水平排列
    label_visibility="collapsed"  # 如果上面的 markdown 标题已足够，可以隐藏 radio 的 label
)
# 根据用户选择的显示名称获取实际的文件名
SELECTED_ICC_FILENAME = ICC_PROFILE_OPTIONS[selected_icc_display_name]
icc_profile_base_name = SELECTED_ICC_FILENAME.split('.')[0]

st.info(f"当前选定的 ICC 配置文件：**{selected_icc_display_name}**")

# 文件上传组件
st.markdown("#### 2. 上传图片文件：")
uploaded_files = st.file_uploader(
    "请上传一个或多个图片文件:",
    type=SUPPORTED_FORMATS,
    accept_multiple_files=True,
    key="file_uploader_main",
    label_visibility="collapsed"  # 隐藏 label，因为上面的 markdown 标题已足够
)

if uploaded_files:
    icc_profile_data = get_icc_profile_bytes(SELECTED_ICC_FILENAME)

    if icc_profile_data is None:
        st.error(f"无法加载选定的 ICC 配置文件 '{SELECTED_ICC_FILENAME}'。请检查文件是否存在。处理已中止。")
        st.stop()

    processed_files_info = []

    with st.status(f"正在使用 '{selected_icc_display_name}' 处理图片...", expanded=True) as status_container:
        for i, uploaded_file in enumerate(uploaded_files):
            status_container.write(f"⏳ 处理: {uploaded_file.name} ({i + 1}/{len(uploaded_files)})")
            image_bytes = uploaded_file.getvalue()

            processed_image_bytes, processed_format = embed_icc_profile(image_bytes, icc_profile_data,
                                                                        uploaded_file.name)

            if processed_image_bytes:
                base, ext = os.path.splitext(uploaded_file.name)
                actual_ext = f".{processed_format.lower()}" if processed_format else ext
                new_filename = f"{base}_{icc_profile_base_name}_embedded{actual_ext}"
                mime_type = f"image/{processed_format.lower()}" if processed_format else uploaded_file.type

                processed_files_info.append({
                    "name": new_filename,
                    "data": processed_image_bytes,
                    "mime": mime_type,
                    "original_name": uploaded_file.name,
                    "id": i
                })
                status_container.write(f"✅ 完成: {new_filename}")
            else:
                status_container.write(f"❌ 失败: {uploaded_file.name}")
        status_container.update(label="所有文件处理完毕!", state="complete", expanded=False)

    if processed_files_info:
        st.markdown("#### 3. 处理结果：")

        tab_labels = [item["name"] for item in processed_files_info]

        if not tab_labels:
            st.info("没有成功处理的图片可供展示。")
        else:
            image_tabs = st.tabs(tab_labels)

            for i, item in enumerate(processed_files_info):
                with image_tabs[i]:
                    st.image(item["data"], use_container_width=True)

                    download_data_buffer = io.BytesIO(item["data"].getvalue())
                    download_data_buffer.seek(0)

                    st.download_button(
                        label=f"下载 {item['name']}",
                        data=download_data_buffer,
                        file_name=item["name"],
                        mime=item["mime"],
                        key=f"download_tab_{item['original_name']}_{item['id']}_{icc_profile_base_name}"
                    )
    else:
        if uploaded_files:
            st.warning("所有上传的文件都未能成功处理。")

    if len(processed_files_info) > 1:
        st.markdown("#### 下载所有处理过的文件 (.zip)")

        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
            for item in processed_files_info:
                item_data_for_zip = io.BytesIO(item["data"].getvalue())
                item_data_for_zip.seek(0)
                zf.writestr(item["name"], item_data_for_zip.read())

        zip_buffer.seek(0)
        st.download_button(
            label=f"下载全部 ({len(processed_files_info)} 张图片) .zip",
            data=zip_buffer,
            file_name=f"processed_images_{icc_profile_base_name}_batch.zip",
            mime="application/zip",
            key=f"download_all_zip_{icc_profile_base_name}"
        )

elif not uploaded_files:
    st.info("请选择 ICC 配置文件并上传图片文件开始处理。")