import httpx
import asyncio
from websockets import ClientConnection
from websockets.asyncio.client import connect
from fastapi import FastAPI, Request, Response, WebSocket
from fastapi.responses import JSONResponse

app = FastAPI(title="Telegram Proxy")

TELEGRAM_HOST = 'square-term-f712.valerazbanovqs.workers.dev'

async def autocancel(request, coroutine):

    async def listen_for_disconnect():
        while True:
            message = await request.receive()
            if message["type"] == "http.disconnect":
                break
            await asyncio.sleep(1)
        return JSONResponse(
            status_code=499,
            content=dict(user_message="Запрос был отменен"),
        )

    async def process_request():
        return await coroutine

    process_request_task = asyncio.create_task(process_request())

    tasks = [
        process_request_task, 
        asyncio.create_task(listen_for_disconnect()),
    ]
    tasks = [task for task in tasks if task is not None]

    done, pending = await asyncio.wait(
        tasks,
        return_when=asyncio.FIRST_COMPLETED
    )

    for task in pending:
        task.cancel()

    for task in pending:
        try:
            await task
        except asyncio.CancelledError:
            ...
        except:
            pass

    return next(iter(done)).result()

def _prepare_headers(headers: dict):
    if "authorization" in headers: del headers["authorization"]
    headers["host"] = TELEGRAM_HOST
    return headers

async def _proxy_request(request: Request) -> Response:
    base_url = f'https://{TELEGRAM_HOST}/'
    url = '/' + request.path_params['url']

    headers = _prepare_headers(dict(request.headers))

    query_params = dict(request.query_params)

    content = await request.body()

    async with httpx.AsyncClient(base_url=base_url, trust_env=False) as session:
        response = await session.request(
            method=request.method,
            url=url,
            headers=headers,
            params=query_params,
            content=content,
            timeout=None,
        )
            
    return Response(
        content=response.content,
        status_code=response.status_code,
        headers=dict(response.headers)
    )

async def proxy_request(request):
    return await autocancel(request, _proxy_request(request))


@app.get("/base/telegram_api/{url:path}", summary='Проксирует все методы get в телеграм апи')
async def get_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.post("/base/telegram_api/{url:path}", summary='Проксирует все методы post в телеграм апи')
async def post_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.put("/base/telegram_api/{url:path}", summary='Проксирует все методы put в телеграм апи')
async def put_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.delete("/base/telegram_api/{url:path}", summary='Проксирует все методы delete в телеграм апи')
async def delete_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.patch("/base/telegram_api/{url:path}", summary='Проксирует все методы patch в телеграм апи')
async def patch_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.options("/base/telegram_api/{url:path}", summary='Проксирует все методы options в телеграм апи')
async def options_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.head("/base/telegram_api/{url:path}", summary='Проксирует все методы head в телеграм апи')
async def head_base_service_path(
    request: Request,
    url: str,
):
    return await proxy_request(request)


@app.websocket("/base/telegram_api/{url:path}")
async def proxy_websocket(websocket: WebSocket, url: str):
    await websocket.accept()
    
    target_url = f"wss://{TELEGRAM_IP}/{url}"
    headers = _prepare_headers(dict(websocket.headers))

    async def forward(source: ClientConnection, destination: ClientConnection):
        async for message in source:
            await destination.send(message)

    try:
        async with connect(target_url, additional_headers=headers, server_hostname=TELEGRAM_HOST) as target_ws:
            await asyncio.gather(
                forward(websocket, target_ws),
                forward(target_ws, websocket)
            )
    finally:
        await websocket.close()

@app.get("/test-connection-google")
async def test_connection_google():
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get("https://google.com", timeout=5.0)
            return {"status": "google_ok", "code": resp.status_code}
    except Exception as e:
        return {"status": "error", "message": str(e)}


@app.get("/test-connection-telegram")
async def test_connection_telegram():
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get("https://api.telegram.org", timeout=5.0)
            return {"status": "telegram_ok", "code": resp.status_code}
    except Exception as e:
        return {"status": "error", "message": str(e)}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="localhost", port=7860)