import re

with open('app/routes/printer_routes.py', 'r', encoding='utf-8') as f:
    content = f.read()

# Replace the WS Proxy header logic
target_str = """            ws_headers['Host'] = request.host
            ws_headers['X-Real-IP'] = request.remote_addr
            forwarded_for = request.headers.get('X-Forwarded-For', '')
            ws_headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
            ws_headers['X-Forwarded-Proto'] = request.scheme
            ws_headers['X-Forwarded-Host'] = request.host
            ws_headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))"""

replacement_str = """            ws_headers['Host'] = request.host
            ws_headers['X-Real-IP'] = request.remote_addr
            forwarded_for = request.headers.get('X-Forwarded-For', '')
            ws_headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
            ws_headers['X-Forwarded-Proto'] = request.scheme
            ws_headers['X-Forwarded-Host'] = request.host
            ws_headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))
            ws_headers['X-Script-Name'] = '/printer/proxy'"""

new_content = content.replace(target_str, replacement_str)

origin_str = """            # Mask Origin/Referer to bypass Octoprint CSRF if needed
            if 'Origin' in ws_headers:
                ws_headers['Origin'] = base_url
            if 'Referer' in ws_headers:
                ws_headers['Referer'] = f"{base_url}/{path}\""""

replacement_origin = """            # Match Tornado's expectations for Origin to avoid 400 Bad Request
            parsed_base = urlparse(base_url)
            ws_headers['Host'] = parsed_base.netloc
            if 'Origin' in ws_headers:
                ws_headers['Origin'] = base_url
            if 'Referer' in ws_headers:
                ws_headers['Referer'] = f"{base_url}/{path}\""""

new_content = new_content.replace(origin_str, replacement_origin)

with open('app/routes/printer_routes.py', 'w', encoding='utf-8') as f:
    f.write(new_content)

print('Patched!')