import os
import json
from flask import Flask, jsonify
from PIL import Image
from zipfile import ZipFile

app = Flask(__name__)

# FIX: 使用 /data/ 路径来避免权限错误
FIG_DIR = "/data/FIG"
ZIP_PATH = "/app/FIG.zip"  # 上传的图像包路径

# 解压 FIG.zip 到 FIG_DIR（只运行一次）
if not os.path.exists(FIG_DIR):
    os.makedirs(FIG_DIR, exist_ok=True)
    with ZipFile(ZIP_PATH, 'r') as zip_ref:
        zip_ref.extractall(FIG_DIR)

@app.route('/')
def home():
    return "✅ Figure extraction service is running. Visit /extract to get JSON."

@app.route('/extract')
def extract_figures():
    results = []
    for filename in os.listdir(FIG_DIR):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif')):
            image_url = f"https://{os.environ.get('SPACE_ID')}.hf.space/file/FIG/{filename}"
            caption = os.path.splitext(filename)[0]
            results.append({
                "image": image_url,
                "caption": caption
            })
    return jsonify(results)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)