from flask import Flask, request, Response, send_from_directory
import itertools
import threading
import queue
import requests
import time
import os

app = Flask(__name__, static_folder=None)

BACKENDS = [
    "http://117.50.89.68:28889",
    "http://117.50.226.22:28889",
]

# One queue per backend, auto-sized from BACKENDS
queues = [queue.Queue() for _ in BACKENDS]
_rr = itertools.count()

EXCLUDED_REQ_HEADERS = {"host", "content-length", "transfer-encoding", "connection"}
EXCLUDED_RESP_HEADERS = {"content-encoding", "content-length", "transfer-encoding", "connection"}

STATIC_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")


def _now():
    return time.strftime("%Y-%m-%d %H:%M:%S")


def worker(q: queue.Queue, backend_url: str):
    """Single worker thread per queue — processes requests one at a time."""
    while True:
        task = q.get()
        if task is None:
            break

        method, url, headers, data, result_event, result_holder = task
        t_start = time.time()
        start_ts = _now()

        try:
            resp = requests.request(
                method=method,
                url=url,
                headers=headers,
                data=data,
                allow_redirects=False,
                timeout=60,
                stream=True,
            )
            result_holder["response"] = resp
            result_holder["error"] = None
        except Exception as e:
            result_holder["response"] = None
            result_holder["error"] = str(e)
        finally:
            elapsed = time.time() - t_start
            end_ts = _now()
            status = resp.status_code if result_holder.get("response") else "ERR"
            print(
                f"[worker] {method} {url} | backend={backend_url} | "
                f"start={start_ts} end={end_ts} | "
                f"{elapsed:.2f}s | status={status}"
            )
            result_event.set()
            q.task_done()


# Start one worker thread per queue
for i, backend in enumerate(BACKENDS):
    t = threading.Thread(target=worker, args=(queues[i], backend), daemon=True, name=f"worker-{i}")
    t.start()
    print(f"[startup] worker-{i} started → {backend}")


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


@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"])
def proxy(path):
    idx = next(_rr) % len(queues)

    # Build the full target URL (preserves query string)
    target_url = BACKENDS[idx] + request.full_path.rstrip("?")

    # Filter hop-by-hop and host headers before forwarding
    headers = {k: v for k, v in request.headers.items() if k.lower() not in EXCLUDED_REQ_HEADERS}

    result_event = threading.Event()
    result_holder: dict = {}

    task = (request.method, target_url, headers, request.get_data(), result_event, result_holder)
    queues[idx].put(task)
    client_ip = request.remote_addr
    print(f"[proxy] {_now()} | {client_ip} | queued → queue-{idx} ({BACKENDS[idx]}) | {request.method} {request.full_path}")

    # Block until the worker finishes (or timeout)
    if not result_event.wait(timeout=65):
        print(f"[proxy] {_now()} | {client_ip} | TIMEOUT → queue-{idx} ({BACKENDS[idx]}) | {request.method} {request.full_path}")
        return Response("Gateway Timeout", status=504)

    if result_holder.get("error"):
        print(f"[proxy] {_now()} | {client_ip} | ERROR → {result_holder['error']}")
        return Response(f"Bad Gateway: {result_holder['error']}", status=502)

    resp = result_holder["response"]
    resp_headers = {k: v for k, v in resp.headers.items() if k.lower() not in EXCLUDED_RESP_HEADERS}

    return Response(resp.content, status=resp.status_code, headers=resp_headers)


if __name__ == "__main__":
    print(f"[startup] proxy listening on 0.0.0.0:28888")
    for i, b in enumerate(BACKENDS):
        print(f"[startup] queue-{i} → {b}")
    app.run(host="0.0.0.0", port=28888, threaded=True)