from flask import Flask, send_from_directory, request, abort
import os
import requests

app = Flask(__name__)

LOCAL_FOLDER = "serve"
BASE_URL = "https://llava.hliu.cc"


@app.route("/", methods=["GET"])
def serve_home():
    return send_from_directory(LOCAL_FOLDER, "index.html")


@app.route("/<path:filepath>", methods=["GET"])
def serve_file(filepath):
    # Check if the file exists locally
    local_path = os.path.join(LOCAL_FOLDER, filepath)
    if os.path.exists(local_path):
        # Serve the file from the local folder
        return send_from_directory(LOCAL_FOLDER, filepath)
    else:
        # Fetch the file from the external URL
        url = os.path.join(BASE_URL, filepath)
        response = requests.get(url)

        if response.status_code == 200:
            # Optionally, save the file locally
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            with open(local_path, "wb") as f:
                f.write(response.content)

            # Serve the file content
            return (
                response.content,
                200,
                {"Content-Type": response.headers["content-type"]},
            )
        else:
            abort(404)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)))