import argparse
import requests
from flask import Flask, request, Response, send_from_directory

app = Flask(__name__, static_folder="static", template_folder="templates")

EDITOR_BASE = "http://localhost:5555"
API_BASE = "http://localhost:5556"

# ---------------------------------------------------------
# 1. Combined UI page
# ---------------------------------------------------------
@app.route("/")
def index():
    return send_from_directory("templates", "index.html")


# ---------------------------------------------------------
# 2. Reverse proxy to ConlluEditor (UI + static)
# ---------------------------------------------------------
@app.route("/editor/", defaults={"path": ""}, methods=["GET", "POST"])
@app.route("/editor/<path:path>", methods=["GET", "POST"])
def proxy_editor(path):
    url = f"{EDITOR_BASE}/{path}"
    if request.query_string:
        url += "?" + request.query_string.decode()

    if request.method == "GET":
        r = requests.get(url, headers=_filtered_headers(), stream=True)
    else:
        r = requests.post(
            url,
            data=request.get_data(),
            headers={"Content-Type": request.headers.get("Content-Type")},
            stream=True
        )

    return Response(
        r.iter_content(chunk_size=8192),
        status=r.status_code,
        headers=_proxied_response_headers(r)
    )


# Handle /editor (no slash)
@app.route("/editor", methods=["GET", "POST"])
def proxy_editor_noslash():
    return proxy_editor("")


# ---------------------------------------------------------
# 3. Reverse proxy for ConlluEditor internal API (/edit/*)
# ---------------------------------------------------------
@app.route("/edit/<path:path>", methods=["GET", "POST"])
def proxy_editor_api(path):
    url = f"{EDITOR_BASE}/edit/{path}"
    if request.query_string:
        url += "?" + request.query_string.decode()

    if request.method == "GET":
        r = requests.get(url, headers=_filtered_headers(), stream=True)
    else:
        r = requests.post(
            url,
            data=request.get_data(),
            headers={"Content-Type": request.headers.get("Content-Type")},
            stream=True
        )

    return Response(
        r.iter_content(chunk_size=8192),
        status=r.status_code,
        headers=_proxied_response_headers(r)
    )


# ---------------------------------------------------------
# 4. Reverse proxy to Upload API (/api/*)
# ---------------------------------------------------------
@app.route("/api/<path:path>", methods=["GET", "POST"])
def proxy_api(path):
    url = f"{API_BASE}/api/{path}"
    if request.query_string:
        url += "?" + request.query_string.decode()

    if request.method == "GET":
        r = requests.get(url, headers=_filtered_headers(), stream=True)
    else:
        # Forward raw body + content-type so Flask can parse multipart uploads
        r = requests.post(
            url,
            data=request.get_data(),
            headers={"Content-Type": request.headers.get("Content-Type")},
            stream=True
        )

    return Response(
        r.iter_content(chunk_size=8192),
        status=r.status_code,
        headers=_proxied_response_headers(r)
    )


# ---------------------------------------------------------
# 5. Catch‑all: static assets for ConlluEditor
# ---------------------------------------------------------
@app.route("/<path:path>", methods=["GET", "POST"])
def proxy_static(path):
    url = f"{EDITOR_BASE}/{path}"
    if request.query_string:
        url += "?" + request.query_string.decode()

    if request.method == "GET":
        r = requests.get(url, headers=_filtered_headers(), stream=True)
    else:
        r = requests.post(
            url,
            data=request.get_data(),
            headers={"Content-Type": request.headers.get("Content-Type")},
            stream=True
        )

    return Response(
        r.iter_content(chunk_size=8192),
        status=r.status_code,
        headers=_proxied_response_headers(r)
    )


# ---------------------------------------------------------
# Header filtering helpers
# ---------------------------------------------------------
def _filtered_headers():
    excluded = {"host", "content-length", "content-encoding", "connection"}
    return {k: v for k, v in request.headers.items() if k.lower() not in excluded}


def _proxied_response_headers(r):
    excluded = {"content-encoding", "transfer-encoding", "connection"}
    return [(k, v) for k, v in r.headers.items() if k.lower() not in excluded]


# ---------------------------------------------------------
# Main entry
# ---------------------------------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()
    app.run(host="0.0.0.0", port=args.port)