from crawl4ai import AsyncWebCrawler
from urllib.parse import urlparse
import aiohttp
import asyncio
from asyncio.exceptions import TimeoutError as async_timeout
from fast_async import make_async
from bs4 import BeautifulSoup, NavigableString
import secrets
from datetime import datetime
import random
import os
import re
import uuid
from typing import List, Dict, Tuple, Optional
from io import BytesIO
import PyPDF2
from fake_useragent import FakeUserAgent
from htmlrag import clean_html, build_block_tree, EmbedHTMLPruner, BM25HTMLPruner
from transformers import AutoTokenizer, AutoConfig
import torch
import time

class Crawler:
    def __init__(self, user_dir=None, rate_limit=1, headless=True, verbose=False):
        self.session_pool = {}  # Track active sessions
        self.verbose = verbose
        self.rate_limit = rate_limit
        self.user_dir = user_dir
        self.headless = headless
        self.crawler = AsyncWebCrawler(
            context_options={"userDataDir": self.user_dir}, 
            headless=self.headless, 
            verbose=self.verbose
        )

        # Browser context management
        self._browser_contexts = {}
        self._context_locks = {}

    async def get_browser_context(self, session_id):
        """Get or create a browser context with proper locking"""
        if session_id not in self._context_locks:
            self._context_locks[session_id] = asyncio.Lock()
        
        async with self._context_locks[session_id]:
            if session_id not in self._browser_contexts:
                context = await self.crawler.new_context()
                self._browser_contexts[session_id] = context
            return self._browser_contexts[session_id]
        
    async def cleanup_browser_context(self, session_id):
        """Safely cleanup browser context"""
        if session_id in self._context_locks:
            async with self._context_locks[session_id]:
                if session_id in self._browser_contexts:
                    try:
                        await asyncio.shield(
                            self._browser_contexts[session_id].close()
                        )
                    except Exception as e:
                        print(f"Error cleaning up browser context: {e}")
                    finally:
                        del self._browser_contexts[session_id]

    def create_session(self):
        """Create a new session with secure ID"""
        session_id = secrets.token_urlsafe(32)  # Secure session ID
        self.session_pool[session_id] = {
            'created_at': datetime.now(),
            'last_used': datetime.now(),
            'requests_count': 0
        }
        return session_id

    def rotate_session(self, session_id):
        """Implement session rotation logic"""
        if self.session_pool[session_id]['requests_count'] > 100:
            self.cleanup_session(session_id)
            return self.create_session()
        return session_id

    def is_dynamic_page(self, html_content: str) -> Tuple[bool, Optional[str]]:
        """Analyzes HTML content to determine if a webpage is dynamically loaded"""
        def _check_structural_indicators(soup: BeautifulSoup) -> Dict[str, int]:
            """Check structural indicators of dynamic content loading."""
            scores = {
                'empty_containers': 0,
                'repeated_structures': 0,
                'api_endpoints': 0,
                'state_management': 0
            }
            
            # 1. Check for empty content containers
            main_containers = soup.find_all(['main', 'div', 'section'], 
                                        class_=lambda x: x and any(term in str(x).lower() 
                                        for term in ['content', 'main', 'feed', 'list', 'container']))
            
            for container in main_containers:
                # Check if container is empty or has minimal content
                if len(container.find_all()) < 3:
                    scores['empty_containers'] += 1
                
                # Check for repeated similar structures (common in dynamic lists)
                children = container.find_all(recursive=False)
                if children:
                    first_child_class = children[0].get('class', [])
                    similar_siblings = [c for c in children[1:] 
                                    if c.get('class', []) == first_child_class]
                    if len(similar_siblings) > 0:
                        scores['repeated_structures'] += 1

            # 2. Check for API endpoints in scripts
            scripts = soup.find_all('script', {'src': True})
            api_patterns = ['/api/', '/graphql', '/rest/', '/v1/', '/v2/']
            for script in scripts:
                if any(pattern in script['src'] for pattern in api_patterns):
                    scores['api_endpoints'] += 1

            # 3. Look for state management setup
            state_patterns = [
                r'window\.__INITIAL_STATE__',
                r'window\.__PRELOADED_STATE__',
                r'__REDUX_STATE__',
                r'__NUXT__',
                r'__NEXT_DATA__',
                r'window\.__data'
            ]
            
            inline_scripts = soup.find_all('script')
            for script in inline_scripts:
                if script.string:
                    for pattern in state_patterns:
                        if re.search(pattern, script.string):
                            scores['state_management'] += 1

            return scores

        def _check_modern_framework_indicators(soup: BeautifulSoup) -> Dict[str, int]:
            """Check for indicators of modern web frameworks and dynamic loading patterns."""
            scores = {
                'framework_roots': 0,
                'hydration': 0,
                'routing': 0
            }
            
            # 1. Framework-specific root elements
            framework_roots = {
                'react': ['react-root', 'react-app', 'root', '__next'],
                'angular': ['ng-version', 'ng-app'],
                'vue': ['v-app', '#app', 'nuxt-app'],
                'modern': ['app-root', 'application', 'spa-root']
            }
            
            for framework, identifiers in framework_roots.items():
                for id_value in identifiers:
                    if (soup.find(attrs={'id': re.compile(id_value, re.I)}) or 
                        soup.find(attrs={'class': re.compile(id_value, re.I)}) or
                        soup.find(attrs={'data-': re.compile(id_value, re.I)})):
                        scores['framework_roots'] += 1

            # 2. Check for hydration indicators
            hydration_patterns = [
                r'hydrate',
                r'createRoot',
                r'reactive',
                r'observable'
            ]
            
            scripts = soup.find_all('script')
            for script in scripts:
                if script.string:
                    for pattern in hydration_patterns:
                        if re.search(pattern, script.string):
                            scores['hydration'] += 1

            # 3. Check for dynamic routing setup
            router_patterns = [
                'router-view',
                'router-link',
                'route-link',
                'history.push',
                'navigation'
            ]
            
            for pattern in router_patterns:
                if soup.find(class_=re.compile(pattern, re.I)) or \
                soup.find(id=re.compile(pattern, re.I)):
                    scores['routing'] += 1

            return scores

        def _check_dynamic_loading_patterns(soup: BeautifulSoup) -> Dict[str, int]:
            """Check for various dynamic content loading patterns."""
            scores = {
                'infinite_scroll': 0,
                'load_more_buttons': 0,
                'pagination': 0,
                'lazy_loading': 0,
                'loading_indicators': 0
            }
            
            # 1. Check for infinite scroll indicators
            scroll_indicators = [
                'infinite-scroll',
                'data-infinite',
                'data-virtualized',
                'virtual-scroll',
                'scroll-container',
                'scroll-viewport'
            ]
            
            for indicator in scroll_indicators:
                elements = soup.find_all(
                    lambda tag: any(indicator.lower() in str(v).lower() 
                                for v in tag.attrs.values())
                )
                if elements:
                    scores['infinite_scroll'] += len(elements)

            # 2. Check for load more buttons
            button_patterns = [
                r'load[_-]?more',
                r'show[_-]?more',
                r'view[_-]?more',
                r'see[_-]?more',
                r'more[_-]?posts',
                r'more[_-]?results'
            ]
            
            for pattern in button_patterns:
                elements = soup.find_all(
                    ['button', 'a', 'div', 'span'],
                    text=re.compile(pattern, re.I)
                )
                if elements:
                    scores['load_more_buttons'] += len(elements)

            # 3. Check for pagination
            pagination_patterns = [
                'pagination',
                'page-numbers',
                'page-nav',
                'page-links'
            ]
            
            for pattern in pagination_patterns:
                elements = soup.find_all(class_=re.compile(pattern, re.I))
                if elements:
                    scores['pagination'] += len(elements)

            # 4. Check for lazy loading
            lazy_patterns = ['lazy', 'data-src', 'data-lazy']
            for pattern in lazy_patterns:
                elements = soup.find_all(
                    lambda tag: any(pattern.lower() in str(v).lower() 
                                for v in tag.attrs.values())
                )
                if elements:
                    scores['lazy_loading'] += len(elements)

            # 5. Check for loading indicators
            loading_patterns = [
                'loading',
                'spinner',
                'skeleton',
                'placeholder',
                'shimmer'
            ]
            
            for pattern in loading_patterns:
                elements = soup.find_all(class_=re.compile(pattern, re.I))
                if elements:
                    scores['loading_indicators'] += len(elements)

            return scores

        def _evaluate_dynamic_indicators(
            structural: Dict[str, int],
            framework: Dict[str, int],
            loading: Dict[str, int]
        ) -> Tuple[bool, Optional[str]]:
            """Evaluate dynamic indicators and return JavaScript instructions."""
            methods = []
            js_snippets = []

            # Infinite Scroll
            if loading['infinite_scroll'] > 0:
                methods.append("scroll")
                js_snippets.append(
                    """
window.scrollTo(0, document.body.scrollHeight);
await new Promise(resolve => setTimeout(resolve, 1000));
                    """.strip().replace('\n', '')
                )

            # Load More Buttons
            if loading['load_more_buttons'] > 0:
                methods.append("button")
                js_snippets.append(
                    """
const button = Array.from(document.querySelectorAll('button, a, div, span')).find(
    el => /load[_-]?more|show[_-]?more/i.test(el.textContent)
);
if (button) {
    button.click();
    await new Promise(resolve => setTimeout(resolve, 1000));
} else {
    console.warn("No 'Load More' button found.");
}
                    """.strip().replace('\n', '')
                )

            # Paginated Interfaces
            if loading.get('pagination', 0) > 0:
                methods.append("pagination")
                js_snippets.append(
                    """
  const nextPage = document.querySelector('a[rel="next"], .pagination-next, .page-next');
if (nextPage) {
    nextPage.click();
    await new Promise(resolve => setTimeout(resolve, 1000));
} else {
    console.warn("No pagination link found.");
}
                    """.strip().replace('\n', '')
                )

            # Lazy Loading
            if loading.get('lazy_loading', 0) > 0:
                methods.append("lazy")
                js_snippets.append(
                    """
if (window.__INITIAL_STATE__ || window.__REDUX_STATE__ || window.__NUXT__ || window.__NEXT_DATA__) {
    console.log('Framework state detected. Consider monitoring network requests for further actions.');
}
                    """.strip().replace('\n', '')
                )

            # Framework and State Management Indicators
            if framework['framework_roots'] > 0 or structural['state_management'] > 0:
                methods.append("stateful")
                js_snippets.append(
                    """
if (window.__INITIAL_STATE__ || window.__REDUX_STATE__ || window.__NUXT__ || window.__NEXT_DATA__) {
    console.log('Detected stateful framework data loading.');
                    }
                    """.strip().replace('\n', '')
                )

            # API-Driven Content
            if structural['api_endpoints'] > 0:
                methods.append("api")
                js_snippets.append(
                    """
console.log('API requests detected. Use browser devtools to inspect network activity for specific endpoints.');
                    """.strip().replace('\n', '')
                )

            # Aggregate and finalize
            if methods:
                js_code = "\n".join(js_snippets)
                return True, js_code
            
            return False, None

        # Main execution
        soup = BeautifulSoup(html_content, 'html.parser')
        
        # Run all checks
        structural_scores = _check_structural_indicators(soup)
        framework_scores = _check_modern_framework_indicators(soup)
        loading_scores = _check_dynamic_loading_patterns(soup)
        
        # Evaluate results
        return _evaluate_dynamic_indicators(structural_scores, framework_scores, loading_scores)

    async def crawl(
            self, 
            url, 
            depth=2, 
            max_pages=5, 
            session_id=None, 
            human_simulation=True, 
            rotate_user_agent=True, 
            rotate_proxy=True, 
            return_html=False
        ):
        if not session_id:
            session_id = self.create_session()

        session_id = self.rotate_session(session_id)

        # List of rotating user agents
        user_agents = [
            'Chrome/115.0.0.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36',
            'Chrome/115.0.0.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36', 
            'Chrome/115.0.0.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36',
            'Chrome/115.0.0.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36',
            'Chrome/115.0.0.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36'
        ]

        # List of rotating proxies
        proxies = [
            "http://50.62.183.123:80",
            "http://104.129.60.84:6516", 
            "http://156.228.118.163:3128",
            "http://142.111.104.97:6107",
            "http://156.228.99.99:3128"
        ]

        try:
            async with self.crawler as crawler:
                # Rotate user agent and optimize headers for each attempt
                headers = {
                    "User-Agent": random.choice(user_agents) if rotate_user_agent else user_agents[0],
                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
                    "Accept-Language": "en-US,en;q=0.5",
                    "Accept-Encoding": "gzip, deflate",
                    "Connection": "keep-alive",
                    "Upgrade-Insecure-Requests": "1",
                    "Sec-Fetch-Dest": "document",
                    "Sec-Fetch-Mode": "navigate",
                    "Sec-Fetch-Site": "none",
                    "Sec-Fetch-User": "?1",
                    "Cache-Control": "max-age=0"
                }
            
                # Update crawler headers for rotation
                crawler.crawler_strategy.headers = headers

                if rotate_proxy:
                    # Update crawler proxy for rotation
                    crawler.crawler_strategy.proxy = random.choice(proxies)

                result_1 = await crawler.arun(
                    session_id=session_id,
                    url=url,
                    magic=True if human_simulation else False,
                    simulate_user=True if human_simulation else False,
                    override_navigator=True if human_simulation else False,
                    depth=depth,
                    max_pages=max_pages,
                    bypass_cache=True,
                    remove_overlay_elements=True,
                    delay_before_retrieve_html=1.0,
                    verbose=self.verbose
                )

            # Update session metrics
            self.session_pool[session_id]['requests_count'] += 1
            self.session_pool[session_id]['last_used'] = datetime.now()

            if result_1.success:
                if hasattr(result_1, 'html'):
                    success, js_code = self.is_dynamic_page(result_1.html)

                    if success:
                        async with crawler as crawler:
                            # Update crawler headers for rotation
                            crawler.crawler_strategy.headers = headers

                            if rotate_proxy:
                                # Update crawler proxy for rotation
                                crawler.crawler_strategy.proxy = random.choice(proxies)

                            print(f"Executing JS code: {js_code}")
                            result_2 = await crawler.arun(
                                session_id=session_id,
                                url=url,
                                magic=True if human_simulation else False,
                                simulate_user=True if human_simulation else False,
                                override_navigator=True if human_simulation else False,
                                depth=depth,
                                max_pages=max_pages,
                                js_code=js_code,
                                bypass_cache=True,
                                remove_overlay_elements=True,
                                delay_before_retrieve_html=1.0,
                                verbose=self.verbose
                            )

                        if result_2.success:
                            result = result_2
                        else:
                            result = result_1

                        # Update session metrics
                        self.session_pool[session_id]['requests_count'] += 1
                        self.session_pool[session_id]['last_used'] = datetime.now()

                    else:
                        result = result_1
            
                    if return_html and hasattr(result, 'html'):
                        return result.html
                    elif hasattr(result, 'fit_markdown'):
                        return result.fit_markdown
                    elif hasattr(result, 'markdown'):
                        return self.extract_content(result.markdown)
                
        except Exception as e:
            print(f"Error crawling {url}: {str(e)}")

        return None
        
    async def crawl_with_retry(
            self, 
            url, 
            depth=2, 
            max_pages=5, 
            max_retries=3, 
            backoff_factor=1, 
            session_id=None,
            human_simulation=True,
            rotate_user_agent=True, 
            rotate_proxy=True,
            return_html=False,
            timeout=10.0
        ):
        """Crawl with retry logic and anti-blocking measures"""

        async def attempt_crawl(attempt):
            try:
                async with async_timeout.timeout(timeout):
                    context = await self.get_browser_context(session_id)
                    return await self.crawl(
                        context, 
                        url, 
                        depth, 
                        max_pages, 
                        session_id, 
                        human_simulation, 
                        rotate_user_agent,
                        rotate_proxy,
                        return_html
                    )
            except asyncio.TimeoutError:
                print(f"Timeout on attempt {attempt} for {url}")
                raise
            except Exception as e:
                print(f"Error on attempt {attempt} for {url}: {e}")
                raise

        if not self.is_valid_url(url) and not self.is_html_url(url):
            print(f"Invalid URL: {url}")
            return f"No web results found for query: {url}"

        for attempt in range(max_retries):
            try:
                if attempt > 0:
                    # Add delay between retries with exponential backoff
                    delay = backoff_factor * (2 ** (attempt - 1))
                    await asyncio.sleep(delay)
                
                return await attempt_crawl(attempt + 1)
            except Exception as e:
                if attempt == max_retries - 1:
                    print(f"Max retries ({max_retries}) reached for {url}")
                    return f"Failed to crawl after {max_retries} attempts: {url}"
                continue

        return f"No content found after {max_retries} attempts for: {url}"

    def extract_content(self, html_content):
        soup = BeautifulSoup(html_content, 'html.parser')
        for script in soup(["script", "style"]):
            script.decompose()
        text = soup.get_text()
        lines = (line.strip() for line in text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        text = '\n'.join(chunk for chunk in chunks if chunk)
        return text
    
    def cleanup_session(self, session_id):
        """Clean up a session"""
        print(f"Cleaning up session {session_id}")
        if session_id in self.session_pool:
            self.crawler.crawler_strategy.kill_session(session_id)
            del self.session_pool[session_id]

    def cleanup_expired_sessions(self):
        """Regular cleanup of expired sessions using proper time calculation"""
        try:
            current_time = datetime.now()
            expired_sessions = []
            
            for sid, data in self.session_pool.items():
                # Calculate time difference in seconds
                time_diff = (current_time - data['last_used']).total_seconds()
                
                # Check if more than 1 hour (3600 seconds)
                if time_diff > 3600:
                    expired_sessions.append(sid)
            
            # Cleanup expired sessions
            for session_id in expired_sessions:
                self.cleanup_session(session_id)
                
        except Exception as e:
            if self.verbose:
                print(f"Error during session cleanup: {str(e)}")
    
    @staticmethod
    def is_valid_url(url):
        try:
            result = urlparse(url)
            return all([result.scheme, result.netloc])
        except ValueError:
            return False
        
    @staticmethod
    def is_html_url(url):
        return url.endswith(".html") or url.endswith(".htm")
    
class CustomCrawler:
    def __init__(
            self, 
            embed_model: str = "HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1", 
            max_concurrent_requests: int = 10, 
            verbose: bool = True
        ):
        print(f"🦀 Initializing the crawler") if verbose else None
        time.sleep(1)
        self.embed_model = embed_model
        self.max_concurrent_requests = max_concurrent_requests
        self.verbose = verbose
        self.ua = FakeUserAgent()
        self.semaphore = asyncio.Semaphore(self.max_concurrent_requests)
        self.sessions = {}

        # Intilizing HTML Pruners and Tokenizer
        print(f"🔃 Loading HTML Pruners and Tokenizer with {self.embed_model}") if self.verbose else None
        self.bm25_html_pruner = BM25HTMLPruner()
        self.embed_html_pruner = EmbedHTMLPruner(
            embed_model=self.embed_model,
            local_inference=True
        )
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.embed_model, 
            use_fast=True, 
            trust_remote_code=True,
            device="cuda" if torch.cuda.is_available() else "cpu"
        )

        # Get the model config and set the max context length for the model
        print(f"🛠️  Getting model configuration for {self.embed_model}") if self.verbose else None
        self.config = AutoConfig.from_pretrained(self.embed_model)
        self.tokenizer.max_seq_length = self.config.max_position_embeddings
        print(f"📏 Setting max context length to {self.tokenizer.max_seq_length}") if self.verbose else None

    async def create_session(self):
        session_id = str(uuid.uuid4())
        timeout = aiohttp.ClientTimeout(total=600) # Set a 10-minute timeout
        connector = aiohttp.TCPConnector(limit=self.max_concurrent_requests) # Connection pool
        self.sessions[session_id] = aiohttp.ClientSession(timeout=timeout, connector=connector)
        print(f"🔗 Created session: {session_id}") if self.verbose else None
        return session_id

    async def close_session(self, session_id):
        session = self.sessions.pop(session_id, None)
        if session:
            await session.close()
            print(f"🌂 Closed session: {session_id}") if self.verbose else None

    async def cleanup_expired_sessions(self, expiration_time: int = 600): # Default 10 minutes
        current_time = time.time()
        expired_sessions = []
        print("🔍 Checking for expired sessions") if self.verbose else None
        for session_id, (session, creation_time) in self.sessions.items():
            if current_time - creation_time > expiration_time:
                expired_sessions.append(session_id)

        for session_id in expired_sessions:
            await self.close_session(session_id)

        print(f"🗑️  Successfully cleaned up all expired sessions") if self.verbose else None

    @make_async
    def html_rag(
            self, 
            query: str, 
            html: str, 
            max_context_length: int = 32000, 
            buffer: int = 2000
        ) -> str:
        if not html:
            raise Exception("No HTML contents provided.")
        
        # Validate HTML structure
        try:
            BeautifulSoup(html, 'html.parser')
        except Exception as e:
            raise Exception(f"Invalid HTML content: {e}")

        prompt_for_retrieval = \
"""Given a query, your task is to retrieve the most relevant passages that answers and/or is relevant to the query.

Query:"""

        self.embed_html_pruner.query_instruction_for_retrieval = prompt_for_retrieval

        print(f"🧹 Pruning HTML for query: {query}") if self.verbose else None
        cleaned_html = clean_html(html)
        block_tree, cleaned_html = build_block_tree(cleaned_html, max_node_words=10)

        block_rankings = self.bm25_html_pruner.calculate_block_rankings(query, cleaned_html, block_tree)
        
        max_context_window = max_context_length - buffer
        pruned_html = self.embed_html_pruner.prune_HTML(
            cleaned_html,
            block_tree,
            block_rankings,
            self.tokenizer,
            max_context_window
        )
        print(f"👍 Successfully pruned HTML for query: {query}") if self.verbose else None
        return pruned_html

    async def fetch_page_contents(
        self,
        urls: List[str],
        query: Optional[str] = None,
        session_id: Optional[str] = None,
        max_attempts: int = 3,
        delay: float = 1.0,
        timeout: float = 10.0,
        return_type: str = "markdown",
        rotate_headers: bool = True,
    ) -> List[Optional[str]]:
        async def fetch_single_page(url, proxies, session=None, query=query):
            for attempt in range(max_attempts):
                print(f"🔍 Attempt {attempt + 1}/{max_attempts}: Fetching content from {url}") if self.verbose else None
                content = await self._fetch_page_contents(
                    url=url,
                    query=query,
                    timeout=timeout,
                    return_type=return_type,
                    rotate_headers=rotate_headers,
                    proxies=proxies,
                    session=session
                )
                
                if content:
                    print(f"✅ Successfully fetched content from {url}") if self.verbose else None
                    return content
                else:
                    print(f"🚫 Failed to fetch content from {url}. Retrying in {delay} seconds...") if self.verbose else None
                    await asyncio.sleep(delay)

            print(f"🚫 Failed to fetch content from {url} after {max_attempts} attempts.") if self.verbose else None
            return None
        
        proxy_list = self.load_proxies()  # Load proxies from environment variables
        if proxy_list:
            proxies = proxy_list
        else:
            proxies = None

        if not urls:
            raise Exception("No URLs provided!")
            
        if return_type == "fit_markdown" and query is None:
            raise Exception("Query must be provided when return_type is 'fit_markdown'!")

        if session_id:  # Use existing session if provided
            if session_id not in self.sessions:
                raise ValueError(f"Invalid session ID: {session_id}")
            session = self.sessions[session_id]
            tasks = [fetch_single_page(url, proxies, session) for url in urls]  # Pass session to tasks
        else:  # No session handling if session_id is None
            tasks = [fetch_single_page(url, proxies) for url in urls]  # No session passed

        results = await asyncio.gather(*tasks)
        return [result for result in results if result is not None]
    
    async def _fetch_page_contents(
        self, 
        url: str, 
        query: Optional[str] = None,
        timeout: float = 10.0,
        return_type: str = "markdown",
        rotate_headers: bool = True,
        proxies: Optional[List[str]] = None,
        session: Optional[aiohttp.ClientSession] = None
    ) -> Optional[str]:
        async def get_content(response, return_type=return_type):
            print(f"📄 Getting content from {url}") if self.verbose else None
            if return_type == "html":
                return await response.text()

            response.raise_for_status()
            content_type = response.headers.get('Content-Type', '').lower()

            if 'application/pdf' in content_type:
                content = await response.read()
                text = self.extract_text_from_pdf(content)
                return text
            elif 'text/html' in content_type:
                html_content = await response.text()
                if return_type == "fit_markdown":
                    html_content = self.html_rag(query, html_content).wait()

                soup = BeautifulSoup(html_content, "html.parser")
                for script_or_style in soup(["script", "style"]):
                    script_or_style.decompose()
                text = self.html_to_markdown(soup)
                return text.strip()
            else:
                print(f"🚫 Unsupported content type {content_type} for URL {url}") if self.verbose else None
                return None

        headers = self.get_headers() if rotate_headers else {}
        proxy = self.get_proxy(proxies) if proxies else None

        # Split timeout into connection and processing phases
        timeout_config = aiohttp.ClientTimeout(
            total=None,  # No total timeout
            connect=timeout,  # Connection timeout
            sock_read=timeout  # Socket read timeout
        )

        try:
            # Use provided session if available
            if session:
                async with session.get(url, proxy=proxy, timeout=timeout_config, headers=headers) as response:
                    return await get_content(response)
            # Otherwise, create a new session for each request
            else:
                async with aiohttp.ClientSession() as new_session:
                    async with new_session.get(url, proxy=proxy, timeout=timeout_config, headers=headers) as response:
                        return await get_content(response)

        except aiohttp.ClientError as e:
            print(f"🚫 Request Exception for {url}: {e}") if self.verbose else None
            return None
        except asyncio.TimeoutError as e:
            print(f"🚫 Timeout error for {url}") if self.verbose else None
            return None
        except Exception as e:
            print(f"🚫 Unexpected error fetching {url}: {e}") if self.verbose else None
            return None
        
    def load_proxies(self) -> Optional[List[str]]:
        # Get all environment variables
        env_vars = dict(os.environ)

        # Load proxies from environment variables
        proxy_pattern = re.compile(r"PROXY_\d+")
        proxies = [env_vars[key] for key in env_vars if proxy_pattern.match(key)]

        if proxies:
            print(f"🔌 Loaded {len(proxies)} proxies from environment variables") if self.verbose else None
            return proxies
        else:
            return None
    
    def get_proxy(self, proxies: List[str]) -> str:
        if proxies:  # Check if the proxies list is not empty
            return next(iter(proxies))
        return None  # Or raise an exception, handle differently, etc.

    def get_headers(self) -> Dict[str, str]:
        return {'User-Agent': self.ua.random}

    def extract_text_from_pdf(self, pdf_content: bytes) -> str:
        try:
            print(f"📕 Extracting text from PDF") if self.verbose else None
            pdf_reader = PyPDF2.PdfReader(BytesIO(pdf_content))
            text = ''
            for page in pdf_reader.pages:
                text += page.extract_text()
            print(f"💪 Successfully extracted text from PDF") if self.verbose else None
            return text
        except Exception as e:
            print(f"🚫 Error extracting text from PDF: {e}") if self.verbose else None
            return ""
        
    def html_to_markdown(self, soup):
        markdown_text = ""
        print(f"📟 Converting HTML to Markdown") if self.verbose else None
        def process_element(element, indent=0):
            nonlocal markdown_text

            if isinstance(element, NavigableString):
                text = str(element).strip()
                if text:
                    markdown_text += text + " "
                return

            tag = element.name

            if tag == "h1":
                markdown_text += "# " + element.text.strip() + "\n\n"
            elif tag == "h2":
                markdown_text += "## " + element.text.strip() + "\n\n"
            elif tag == "h3":
                markdown_text += "### " + element.text.strip() + "\n\n"
            elif tag == "h4":
                markdown_text += "#### " + element.text.strip() + "\n\n"
            elif tag == "h5":
                markdown_text += "##### " + element.text.strip() + "\n\n"
            elif tag == "h6":
                markdown_text += "###### " + element.text.strip() + "\n\n"
            elif tag == "p":
                markdown_text += element.text.strip() + "\n\n"
            elif tag == "br":
                markdown_text += "\n"
            elif tag == "ul":
                for li in element.find_all("li", recursive=False):
                    markdown_text += "  " * indent + "- "
                    process_element(li, indent + 1)
                    markdown_text += "\n"
                markdown_text += "\n"
            elif tag == "ol":
                for i, li in enumerate(element.find_all("li", recursive=False), 1):
                    markdown_text += "  " * indent + f"{i}. "
                    process_element(li, indent + 1)
                    markdown_text += "\n"
                markdown_text += "\n"
            elif tag == "table":
                rows = element.find_all("tr")
                for row in rows:
                    cells = row.find_all(["td", "th"])
                    row_text = [cell.text.strip() for cell in cells]
                    markdown_text += "| " + " | ".join(row_text) + " |\n"
                    if row == rows[0]:  # Header row separator
                        markdown_text += "| " + " | ".join(["---"] * len(cells)) + " |\n"
                markdown_text += "\n"
            elif tag == "blockquote":
                markdown_text += "> " + element.text.strip().replace("\n", "\n> ") + "\n\n"
            elif tag == "strong" or tag == "b":
                markdown_text += "**" + element.text.strip() + "**"
            elif tag == "em" or tag == "i":
                markdown_text += "*" + element.text.strip() + "*"
            elif tag == "code":
                markdown_text += "`" + element.text.strip() + "`"
            elif tag == "pre":
                markdown_text += "```\n" + element.text + "\n```\n\n"
            elif tag == "hr":
                markdown_text += "---\n\n"
            else:
                for child in element.children:
                    process_element(child, indent)

        process_element(soup)
        print(f"👌 Successfully converted HTML to Markdown") if self.verbose else None
        return markdown_text
    
if __name__ == "__main__":
    import time
    import winloop

    URLS = [
        "https://en.wikipedia.org/wiki/Treaty_Principles_Bill#:~:text=The%20Treaty%20Principles%20Bill%2C%20or,of%20the%20Treaty%20of%20Waitangi.",
        "https://www.parliament.nz/en/pb/sc/make-a-submission/document/54SCJUST_SCF_227E6D0B-E632-42EB-CFFE-08DCFEB826C6/principles-of-the-treaty-of-waitangi-bill",
        "https://en.wikipedia.org/wiki/Waitangi_Tribunal",
        "https://aljazeera.com/news/2024/11/19/why-are-new-zealands-maori-protesting-over-colonial-era-treaty-bill",
        "https://downiewenjack.ca/treaty-of-waitangi-treaty-principles-bill/"
    ]#  * 10   # Make 50 requests

    query = "What is the Treaty of Waitangi Bill?"
    loop = asyncio.get_event_loop()
    custom_crawler = CustomCrawler(max_concurrent_requests=1000)
    session_id = loop.run_until_complete(custom_crawler.create_session())
    start = time.perf_counter()
    winloop.install()
    result = loop.run_until_complete(custom_crawler.fetch_page_contents(
            URLS, 
            query,
            session_id=session_id, 
            timeout=20,
            max_attempts=1,
            return_type="fit_markdown",
        )
    )
    end = time.perf_counter()
    loop.run_until_complete(custom_crawler.close_session(session_id))
    loop.run_until_complete(custom_crawler.cleanup_expired_sessions())
    print("\n\n".join([f"Document {i+1}:\n\n{result[i]}" for i in range(len(result))]))  
    print(f"\n\nTime taken: {end - start} seconds")