import re

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

# Find everything between @sock.route('/proxy', bp=printer_bp) and def octo_proxy(path):
start_str = "@sock.route('/proxy', bp=printer_bp)"
end_str = "@printer_bp.route('/proxy', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])"

pre_content = content[:content.find(start_str)]
post_content = content[content.find(end_str):]

new_proxy = """@printer_bp.route('/proxy', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
@printer_bp.route('/proxy/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH'])
@login_required
def octo_proxy(path):
    if not current_user.is_admin:
        return "Unauthorized", 403

    url_config = SystemConfig.query.filter_by(key='octoprint_url').first()
    if not url_config or not url_config.value:
        return "OctoPrint URL not configured", 404

    base_url = url_config.value.rstrip('/')
    
    # --- WebSocket Proxy Logic ---
    if request.headers.get('Upgrade', '').lower() == 'websocket':
        from flask_sock import Server, ConnectionClosed
        
        # Check if environment supports WebSockets
        try:
            ws = Server(request.environ)
        except Exception as e:
            return "WebSocket Upgrade Failed", 400
            
        def handle_ws():
            if base_url.startswith('https://'):
                ws_base = base_url.replace('https://', 'wss://', 1)
            else:
                ws_base = base_url.replace('http://', 'ws://', 1)
                
            target_url = f"{ws_base}/{path}"
            if request.query_string:
                target_url = f"{target_url}?{request.query_string.decode('utf-8')}"

            # Forward essential headers like NGINX Proxy
            headers = {
                'Host': request.host,
                'X-Real-IP': request.remote_addr,
            }
            forwarded_for = request.headers.get('X-Forwarded-For', '')
            headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
            headers['X-Forwarded-Proto'] = request.scheme
            headers['X-Forwarded-Host'] = request.host
            headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))
            
            if request.headers.get('Cookie'):
                headers['Cookie'] = request.headers.get('Cookie')

            try:
                remote_ws = ws_connect(target_url, additional_headers=headers)
            except Exception as e:
                ws.close(1011, str(e))
                return

            def recv_loop():
                try:
                    for message in remote_ws:
                        ws.send(message)
                except Exception:
                    pass
                finally:
                    try: remote_ws.close()
                    except: pass
                    try: ws.close()
                    except: pass

            t = threading.Thread(target=recv_loop)
            t.daemon = True
            t.start()

            try:
                while True:
                    data = ws.receive()
                    if data is None:
                        break
                    remote_ws.send(data)
            except Exception:
                pass
            finally:
                try: remote_ws.close()
                except: pass
                try: ws.close()
                except: pass
        
        try:
            handle_ws()
        except ConnectionClosed:
            pass
        except Exception:
            pass
        finally:
            try: ws.close() 
            except: pass

        class WebSocketResponse(Response):
            def __call__(self, *args, **kwargs):
                if getattr(ws, 'mode', 'werkzeug') == 'werkzeug':
                    return super().__call__(*args, **kwargs)
                return []
                
        return WebSocketResponse()

    # --- Standard HTTP Proxy Logic ---
    from urllib.parse import urlparse
    target_url = f"{base_url}/{path}"
    
    if request.query_string:
        target_url = f"{target_url}?{request.query_string.decode('utf-8')}"

    # Build headers for reverse proxy based on nginx config reference
    parsed_base = urlparse(base_url)
    headers = {k: v for k, v in request.headers if k.lower() not in ['host', 'content-length']}
    
    # NGINX equivalent proxy headers
    headers['Host'] = request.host
    headers['X-Real-IP'] = request.remote_addr
    headers['X-Real-Port'] = str(request.environ.get('REMOTE_PORT', ''))
    
    forwarded_for = request.headers.get('X-Forwarded-For', '')
    headers['X-Forwarded-For'] = f"{forwarded_for}, {request.remote_addr}" if forwarded_for else request.remote_addr
    
    headers['X-Forwarded-Protocol'] = request.scheme
    headers['X-Script-Name'] = "/printer/proxy"
    headers['X-Forwarded-Host'] = request.host
    headers['X-Forwarded-Port'] = str(request.environ.get('SERVER_PORT', '80'))
    headers['REMOTE-HOST'] = request.remote_addr
    
    if request.headers.get('Upgrade'):
        headers['Upgrade'] = request.headers.get('Upgrade')
    if request.headers.get('Connection'):
        headers['Connection'] = request.headers.get('Connection')
        
    try:
        # proxy_connect_timeout 60s, proxy_read_timeout 600s
        resp = requests.request(
            method=request.method,
            url=target_url,
            headers=headers,
            data=request.get_data(),
            cookies=request.cookies,
            allow_redirects=False,
            stream=True,
            timeout=(60, 600)
        )
    except requests.exceptions.RequestException as e:
        return f"Proxy connection error: {str(e)}", 502

    # Strip headers that might break the iframe or framing
    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection', 'x-frame-options', 'content-security-policy']
    response_headers = [(name, value) for (name, value) in resp.headers.items() if name.lower() not in excluded_headers]

    def generate():
        for chunk in resp.iter_content(chunk_size=8192):
            if chunk:
                yield chunk

    return Response(generate(), resp.status_code, response_headers)
"""

post_end_str = "    return Response(generate(), resp.status_code, response_headers)"
post_end_idx = post_content.find(post_end_str) + len(post_end_str)
final_content = pre_content + new_proxy + post_content[post_end_idx:]

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

print("Patched!")