crawler/engine.py

import asyncio
import logging
from typing import Set
from urllib.parse import urlparse
from .config import CrawlerConfig
from .anti_blocking import AntiBlockingManager
from .state import CrawlerState
from .client import RobustHTTPClient
from .parser import PageParser

logger = logging.getLogger("crawler.engine")

class CrawlEngine:
    def __init__(self, config: CrawlerConfig, db_path: str = "crawler_state.db"):
        self.config = config
        self.state = CrawlerState(db_path)
        self.anti_blocking = AntiBlockingManager()
        self.client = RobustHTTPClient(config, self.anti_blocking)
        self.parser = PageParser(config.domain, config.scheme)
        self.semaphore = asyncio.Semaphore(config.concurrency_limit)
        
        # Track pages scraped in active run
        self.pages_scraped_count = 0
        
        # Prepare queue - reset interrupted running tasks, seed base URL
        self.state.reset_running()
        self.state.add_to_queue(config.base_url, depth=0)

    async def _process_url(self, url: str, depth: int):
        """Worker task processing a single URL: get, parse, store, enqueue new links."""
        async with self.semaphore:
            # Check maximum depth limits
            if depth > self.config.max_depth:
                self.state.mark_visited(url, depth, error="Exceeded max crawl depth")
                return

            # Check if max pages target reached
            if self.config.max_pages and self.pages_scraped_count >= self.config.max_pages:
                logger.info(f"Crawl limit reached ({self.config.max_pages} pages). Skipping {url}")
                # Re-queue it to handle if resumed later
                self.state.mark_failed(url, depth, "Skipped due to page limit")
                return

            try:
                # 1. Fetch Page
                response = await self.client.get(url)
                
                # Check response header type (only parse html text)
                content_type = response.headers.get("content-type", "").lower()
                if "text/html" not in content_type:
                    self.state.mark_visited(url, depth, status_code=response.status_code, error="Not HTML content")
                    return
                
                # 2. Parse Data
                html_content = response.text
                parsed_data = self.parser.parse(html_content, url)
                
                # 3. Store Data
                self.state.save_scraped_data(url, parsed_data)
                self.pages_scraped_count += 1
                logger.info(f"Successfully scraped [{response.status_code}]: {url} ({parsed_data['title']})")
                
                # 4. Discover Links (only if we haven't reached depth limit)
                if depth < self.config.max_depth:
                    new_links = self.parser.extract_links(html_content, url)
                    links_to_add = [(link, depth + 1) for link in new_links]
                    added = self.state.add_multiple_to_queue(links_to_add)
                    if added > 0:
                        logger.info(f"Queued {added} new internal links found on {url}")

                # 5. Mark as complete
                self.state.mark_visited(url, depth, status_code=response.status_code)

            except Exception as e:
                logger.error(f"Error crawling {url}: {str(e)}")
                self.state.mark_failed(url, depth, error=str(e))

    async def run(self):
        """Starts the main async crawling loop, coordinating workers."""
        logger.info(f"Starting crawler for domain: {self.config.domain}")
        logger.info(f"Concurrency limit: {self.config.concurrency_limit}, delay range: {self.config.delay_range}")
        
        active_tasks = set()
        
        while True:
            # Check limits
            if self.config.max_pages and self.pages_scraped_count >= self.config.max_pages:
                logger.info("Configured max pages scrape limit reached. Terminating crawl.")
                break

            # Fill up worker slots from queue
            free_slots = self.config.concurrency_limit - len(active_tasks)
            for _ in range(free_slots):
                next_item = self.state.get_next_url()
                if next_item:
                    url, depth = next_item
                    task = asyncio.create_task(self._process_url(url, depth))
                    active_tasks.add(task)
                    # Add callback to self-clean finished tasks
                    task.add_done_callback(active_tasks.discard)
                else:
                    break

            # If no tasks are currently executing and queue is empty, we are done
            if not active_tasks:
                # Double-check queue in case of race condition or delay
                next_item = self.state.get_next_url()
                if not next_item:
                    logger.info("Crawl complete. No more pending links in the queue.")
                    break
                else:
                    # Queue populated, continue loop
                    continue

            # Wait briefly for any task to complete before next loop iteration
            await asyncio.sleep(0.1)

        # Cleanup client connection pools
        await self.client.close()
        logger.info("Crawler run finished.")
        
        # Print crawl statistics summary
        stats = self.state.get_stats()
        print("\n--- Crawl Statistics ---")
        print(f"Scraped & Saved:   {stats['scraped']} pages")
        print(f"Visited (Success): {stats['successful']} pages")
        print(f"Visited (Failed):  {stats['failed']} pages")
        print(f"Pending in Queue:  {stats['pending']} pages")
        print("------------------------")