import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

import requests
from flask import Flask, request, Response

TARGET = 'agentrouter.org'
PORT = 5000

app = Flask(__name__)

@app.route('/', defaults={'path': ''})
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def proxy(path):
    print(f"[{request.method}] /{path}")

    # Диагностика ключа
    auth_header = request.headers.get('Authorization')
    if auth_header:
        print(f"🔑 Токен передан: {auth_header[:15]}...")
    else:
        print("⚠️ Внимание: Запрос без токена Authorization")

    # 1. Фильтруем входящие заголовки от плагина
    headers = {}
    forbidden_incoming_headers = [
        'host', 'content-length', 'user-agent', 'x-title', 
        'http-referer', 'origin', 'sec-fetch-dest', 'sec-fetch-mode', 
        'sec-fetch-site', 'accept-language', 'accept-encoding'
    ]
    
    for key, value in request.headers:
        if key.lower() not in forbidden_incoming_headers:
            headers[key] = value

    # 2. Атака на WAF: маскируемся под разрешенного клиента
    headers['Host'] = TARGET
    headers['User-Agent'] = 'codex_cli_rs/0.101.0'
    headers['Originator'] = 'codex_cli_rs'
    
    # 3. ИСПРАВЛЕНИЕ СТРИМИНГА: Запрещаем серверу сжимать ответ
    headers['Accept-Encoding'] = 'identity'
    
    headers['Accept'] = '*/*'
    headers['Connection'] = 'keep-alive'

    url = f'https://{TARGET}/{path}'

    try:
        # 4. Выполняем запрос
        resp = requests.request(
            method=request.method,
            url=url,
            headers=headers,
            data=request.get_data(),
            params=request.args,
            stream=True,
            verify=False
        )

        # 5. Обрабатываем ответ сервера
        excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
        response_headers = []
        for name, value in resp.headers.items():
            if name.lower() not in excluded_headers:
                response_headers.append((name, value))

        # 6. Возвращаем CORS-заголовки
        response_headers.append(('Access-Control-Allow-Origin', '*'))
        response_headers.append(('Access-Control-Allow-Methods', '*'))
        response_headers.append(('Access-Control-Allow-Headers', '*'))

        # 7. ИСПРАВЛЕНИЕ СТРИМИНГА: chunk_size=None для моментальной отдачи
        def generate():
            # None позволяет отдавать данные сразу, как только они получены из сокета, 
            # не дожидаясь накопления определенного количества байт.
            for chunk in resp.iter_content(chunk_size=None):
                if chunk:
                    yield chunk

        print(f"✅ Статус ответа сервера: {resp.status_code}")
        return Response(generate(), status=resp.status_code, headers=response_headers)

    except Exception as e:
        print(f'❌ Proxy error: {e}')
        return f'Proxy error: {str(e)}', 500

if __name__ == '__main__':
    print(f"🚀 Roo Code Proxy (WAF Bypass + SSE Fix) is running on http://127.0.0.1:{PORT}")
    print(f"📡 Forwarding to https://{TARGET} as 'codex_cli_rs'")
    app.run(host='127.0.0.1', port=PORT, debug=False, threaded=True)