from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse
import subprocess
import io
import os
from typing import Optional
from urllib.parse import urlparse

app = FastAPI()

@app.post("/process")
async def process_image(
    file: Optional[UploadFile] = File(None),
    url: Optional[str] = Form(None),
    options: Optional[str] = Form("-q:v 2"),
    filename: Optional[str] = Form(None) # 新增：支持自定义文件名
):
    input_source = ""
    input_data = None

    # 1. 确定输入源
    if file:
        input_data = await file.read()
        input_source = "pipe:0"
        # 如果没传 filename，则使用上传文件的原始名，但后缀改为 .jpg
        if not filename:
            filename = os.path.splitext(file.filename)[0] + ".jpg"
    elif url:
        input_source = url
        # 如果没传 filename，从 URL 中解析文件名，后缀改为 .jpg
        if not filename:
            path = urlparse(url).path
            base_name = os.path.basename(path) or "downloaded_image"
            filename = os.path.splitext(base_name)[0] + ".jpg"
    else:
        raise HTTPException(status_code=400, detail="Missing source")

    # 2. 执行 FFmpeg 转换 (PNG -> JPG)
    command = ['ffmpeg', '-i', input_source]
    if options:
        command.extend(options.split())
    command.extend(['-f', 'image2', 'pipe:1'])

    try:
        process = subprocess.Popen(
            command,
            stdin=subprocess.PIPE if input_data else None,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        out, err = process.communicate(input=input_data)

        # 3. 构建响应，添加 Content-Disposition Header
        # 这会让 n8n 和浏览器识别出文件名
        headers = {
            'Content-Disposition': f'attachment; filename="{filename}"'
        }
        
        return StreamingResponse(
            io.BytesIO(out), 
            media_type="image/jpeg", 
            headers=headers
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))