import gradio as gr
from PIL import Image, ImageChops
import img2pdf
import io

import os

# --- 核心逻辑 (复用之前的代码) ---
def trim_whitespace(im, fuzz_level=20):
    bg = Image.new(im.mode, im.size, (255, 255, 255))
    diff = ImageChops.difference(im, bg)
    diff = ImageChops.add(diff, diff, 1, -fuzz_level)
    bbox = diff.getbbox()
    if bbox:
        return im.crop(bbox)
    return im

def process_pipeline(files, quality, fuzz_level, progress=gr.Progress()):
    if not files: return None
    output_files = []
    
    for filepath in progress.tqdm(files, desc="处理中"):
        try:
            with Image.open(filepath) as img:
                # 处理透明背景
                if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                    bg = Image.new('RGB', img.size, (255, 255, 255))
                    if img.mode!= 'RGBA': img = img.convert('RGBA')
                    bg.paste(img, mask=img.split()[3])
                    img = bg
                else:
                    img = img.convert('RGB')
                
                # 裁边
                trimmed = trim_whitespace(img, fuzz_level)
                
                # 获取原始文件名并构建输出路径
                original_name = os.path.basename(filepath)
                name, _ = os.path.splitext(original_name)
                # 保持文件名，改为 .pdf 后缀
                output_filename = f"{name}_cropped.pdf"
                
                # 保存为 PDF
                # img2pdf 需要 bytes 或文件路径，这里我们先存为临时图片再转，或者直接用 PIL save pdf
                # PIL save PDF 支持单张
                trimmed.save(output_filename, "PDF", resolution=100.0, save_all=True)
                output_files.append(output_filename)
                
        except Exception as e:
            print(f"Error processing {filepath}: {e}")
            continue

    return output_files

# --- 模块化 UI 接口 ---
def create_ui():
    """
    每个工具模块都需要暴露这个函数。
    注意：不要在这里创建 gr.Blocks()，直接写 Row/Column 即可，
    因为它们会被嵌入到主程序的 Tab 中。
    """
    with gr.Row():
        with gr.Column():
            file_input = gr.File(file_count="multiple", file_types=["image"], label="上传图片 (支持多选)")
            quality = gr.Slider(10, 100, 90, label="质量 (仅用于压缩，当前直接转PDF可忽略)")
            fuzz = gr.Slider(0, 100, 30, label="容差")
            btn = gr.Button("开始处理", variant="primary")
        with gr.Column():
            # 输出文件列表
            output = gr.File(label="下载结果 (点击文件名下载)", file_count="multiple")
            # 增加一个 Zip 下载选项，方便用户
            # 注意：这里我们暂时不实现 Zip 打包逻辑，因为用户明确说“不要打包”
            # 但为了方便“一次性下载”，通常 Zip 是唯一解。
            # 如果用户坚持不要 Zip，那只能列表展示。
            
    btn.click(process_pipeline, [file_input, quality, fuzz], output)