import os
import requests
from flask import Flask, render_template, request, jsonify

# 配置 Flask
app = Flask(__name__)
UPLOAD_FOLDER = "static/uploads"
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER

# 你的推理服务器地址
SERVER_URL = "http://127.0.0.1:5000/infer"  # 如果你的服务器在远程，改为公网 IP 或 ngrok 地址

# 确保上传目录存在
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        # 获取上传的图片和文本
        image = request.files["image"]
        text = request.form["text"]
        conv_mode = request.form.get("conv_mode", "vicuna_v1")

        if image:
            image_path = os.path.join(app.config["UPLOAD_FOLDER"], image.filename)
            image.save(image_path)  # 保存图片

            # 发送请求到推理服务器
            payload = {
                "image_path": image_path,
                "text": text,
                "conv_mode": conv_mode
            }
            response = requests.post(SERVER_URL, json=payload)

            if response.status_code == 200:
                result = response.json()["response"]
            else:
                result = "Error: Server did not respond correctly."

            return render_template("index.html", image_url=image_path, text=text, result=result)

    return render_template("index.html", image_url=None, text="", result="")

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080, debug=True)