import asyncio
import random
from fake_useragent import UserAgent
import aiohttp
import logging

class ProxyForbiddenException(Exception):
    pass


class GrassConnect:
    def __init__(self, proxy: str = None):
        self.user_agent = UserAgent().random
        self.proxy = proxy
        self.session = None
        self.websocket = None
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
        self.logger = logging.getLogger(__name__)

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def connect(self):
        uri = f"wss://proxy.wynd.network"
        headers = {
            'Pragma': 'no-cache',
            'Origin': 'chrome-extension://ilehaonighjijnmpnagapkhpcdbhclfg',
            'Accept-Language': 'uk-UA,uk;q=0.9,en-US;q=0.8,en;q=0.7',
            'User-Agent': self.user_agent,
            'Upgrade': 'websocket',
            'Cache-Control': 'no-cache',
            'Connection': 'Upgrade',
            'Sec-WebSocket-Version': '13',
            'Sec-WebSocket-Extensions': 'permessage-deflate; client_max_window_bits'
        }
        while True:
            self.logger.info(f"Attempting to connect to {uri} using proxy {self.proxy}")
            try:
                self.websocket = await self.session.ws_connect(
                    uri,
                    headers=headers,
                    proxy=self.proxy,
                    proxy_headers=headers
                )
                self.logger.info(f"Successfully connected to {uri}")
                return self.websocket
            except aiohttp.ClientResponseError as e:
                self.logger.error(f"HTTP error occurred: {e}")
                if e.status == 403:
                    raise ProxyForbiddenException(f"Low proxy score. Can't connect. Error: {e}")
            except Exception as e:
                self.logger.error(f"Unexpected error occurred: {e}")
            self.logger.info("Connection failed. Retrying in 3 seconds...")
            await asyncio.sleep(3)

    async def close(self):
        if self.websocket:
            await self.websocket.close()
            self.logger.info("Websocket connection closed.")
        if self.session:
            await self.session.close()
            self.logger.info("Client session closed.")


# Example usage
async def main():
    proxy = "socks5://192.168.31.201:53012"
    async with GrassConnect(proxy=proxy) as connector:
        websocket = await connector.connect()
        # 继续处理websocket连接

if __name__ == "__main__":
    asyncio.run(main())