import re

with open("gateway.py", "r") as f:
    code = f.read()

# I see what's happening. When `llama_auth_headers()` gets called by `gateway.py`, it now returns `{"Authorization": f"Bearer {LLAMA_API_KEY}"} if LLAMA_API_KEY else {}`.
# Because `LLAMA_API_KEY` is not set, it returns `{}`.
# But `llama-swap` IS running in the space. And it seems to require an API key even if we didn't specify one?
# Wait! In `llama-swap.yaml`, there's no api key. But `llama-swap` uses the `API_KEY` environment variable as its API key by default!
# So if `API_KEY` is set in the environment (which it is, presumably configured by the user as a secret), `llama-swap` ENFORCES IT!
# And since we made `gateway.py` NOT send the `API_KEY` to `llama-swap` when `LLAMA_API_KEY` is empty, `llama-swap` rejects it with 401.

new_headers = """def llama_auth_headers() -> dict[str, str]:
    key = LLAMA_API_KEY or API_KEY
    return {"Authorization": f"Bearer {key}"} if key else {}"""

old_headers = """def llama_auth_headers() -> dict[str, str]:
    return {"Authorization": f"Bearer {LLAMA_API_KEY}"} if LLAMA_API_KEY else {}"""

code = code.replace(old_headers, new_headers)

with open("gateway.py", "w") as f:
    f.write(code)