import gradio as gr
import subprocess, shutil, os, zipfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent
OUTPUT_DIR = ROOT / "output"
ZIP_PATH = ROOT / "output.zip"

def run_pipeline(model_name_t, model_name_v, result_dir, paper_latex_root, arxiv_url, openai_key, gemini_key):
    # 清理旧输出
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    if ZIP_PATH.exists():
        ZIP_PATH.unlink()

    # 构造命令（传递两个 key）
    cmd = [
        "python", "pipeline.py",
        "--model_name_t", model_name_t,
        "--model_name_v", model_name_v,
        "--result_dir", result_dir,
        "--paper_latex_root", paper_latex_root,
        "--arxiv_url", arxiv_url,
        "--openai_key", openai_key,
        "--gemini_key", gemini_key
    ]
    
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
        logs = result.stdout + "\n" + result.stderr
    except subprocess.TimeoutExpired:
        return "❌ Pipeline timed out (30 min limit).", None
    except Exception as e:
        return f"❌ Error: {e}", None

    if not OUTPUT_DIR.exists():
        return "❌ No output generated.", None

    # 压缩 output 文件夹
    with zipfile.ZipFile(ZIP_PATH, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(OUTPUT_DIR):
            for file in files:
                file_path = Path(root) / file
                zipf.write(file_path, arcname=file_path.relative_to(OUTPUT_DIR))
    
    return logs, ZIP_PATH

# ===================== Gradio UI =====================
iface = gr.Interface(
    fn=run_pipeline,
    inputs=[
        gr.Textbox(label="Model Name (Text)", value="gpt-4.1"),
        gr.Textbox(label="Model Name (Vision)", value="gpt-4.1"),
        gr.Textbox(label="Result Dir", value="output"),
        gr.Textbox(label="Paper LaTeX Root", value="input/latex_proj"),
        gr.Textbox(label="ArXiv URL", value="https://arxiv.org/abs/2505.21497"),
        gr.Textbox(label="OpenAI API Key", placeholder="sk-...", type="password"),
        gr.Textbox(label="Gemini API Key", placeholder="AIza...", type="password"),
    ],
    outputs=[
        gr.Textbox(label="Logs"),
        gr.File(label="Download Output (.zip)")
    ],
    title="PaperShow Pipeline",
    description="输入 arXiv 链接和参数，自动生成 slides + poster，结果打包下载。"
)

if __name__ == "__main__":
    iface.launch()