halldyll_core/

orchestrator.rs

//! Orchestrator - Main scraper orchestration with full component integration

use std::sync::Arc;
use url::Url;

use crate::crawl::{CanonicalResolver, ContentDedup, CrawlEntry, Frontier, UrlDedup, UrlNormalizer};
use crate::fetch::{FetchResponse, HttpClient};
use crate::observe::{MetricsCollector, StructuredLogger};
use crate::parse::{
    AudioExtractor, HtmlParser, ImageExtractor, JsonLdExtractor, LinkExtractor,
    MetadataExtractor, OpenGraphExtractor, TextExtractor, VideoExtractor,
};
use crate::politeness::{DomainThrottler, RobotsChecker};
use crate::render::{BrowserPool, RenderChecker, RenderDecision};
use crate::security::{DomainAllowlist, IpBlocker, ResourceLimits};
use crate::storage::{NormalizedStore, RawSnapshot, SnapshotStore};
use crate::types::document::RobotsDirectives;
use crate::types::error::Result;
use crate::types::{Config, Document, Error};

/// Scrape result with additional metadata
#[derive(Debug)]
pub struct ScrapeResult {
    /// Extracted document
    pub document: Document,
    /// Whether content was a duplicate
    pub is_duplicate: bool,
    /// Whether JS rendering was used
    pub js_rendered: bool,
    /// Canonical URL (if different from source)
    pub canonical_url: Option<Url>,
    /// Links discovered (for crawling)
    pub discovered_links: Vec<Url>,
}

/// Main orchestrator
pub struct Orchestrator {
    /// Configuration
    config: Arc<Config>,
    /// HTTP client
    client: HttpClient,
    /// URL normalizer
    normalizer: UrlNormalizer,
    /// Crawl frontier
    frontier: Frontier,
    /// URL dedup
    url_dedup: UrlDedup,
    /// Content dedup
    content_dedup: ContentDedup,
    /// Canonical resolver
    canonical_resolver: CanonicalResolver,
    /// robots.txt checker
    robots_checker: RobotsChecker,
    /// Domain throttler
    throttler: DomainThrottler,
    /// Render checker
    render_checker: RenderChecker,
    /// Browser pool for JS rendering
    browser_pool: BrowserPool,
    /// Domain allowlist
    allowlist: DomainAllowlist,
    /// IP blocker
    ip_blocker: IpBlocker,
    /// Resource limits
    limits: ResourceLimits,
    /// Snapshot store
    snapshot_store: SnapshotStore,
    /// Document store
    document_store: NormalizedStore,
    /// Metrics collector
    metrics: MetricsCollector,
    /// Logger
    logger: StructuredLogger,
}

impl Orchestrator {
    /// Create a new orchestrator
    pub fn new(config: Config) -> Result<Self> {
        let client = HttpClient::new(config.clone())?;

        let frontier = Frontier::new(
            config.crawl.strategy.clone(),
            config.crawl.max_depth,
            config.crawl.max_urls_per_domain,
            config.crawl.max_urls_total,
        );

        let throttler = DomainThrottler::new(
            config.politeness.default_delay_ms,
            config.politeness.max_concurrent_per_domain as usize,
            config.politeness.max_concurrent_total as usize,
            config.politeness.adaptive_delay,
            config.politeness.rate_limit_pause_ms,
        );

        let robots_checker = RobotsChecker::new(
            &config.fetch.user_agent,
            config.politeness.robots_cache_ttl_secs,
        );

        let logger = StructuredLogger::new();

        Ok(Self {
            config: Arc::new(config),
            client,
            normalizer: UrlNormalizer::default(),
            frontier,
            url_dedup: UrlDedup::new(),
            content_dedup: ContentDedup::default(),
            canonical_resolver: CanonicalResolver::default(),
            robots_checker,
            throttler,
            render_checker: RenderChecker::default(),
            browser_pool: BrowserPool::default(),
            allowlist: DomainAllowlist::new(),
            ip_blocker: IpBlocker::default(),
            limits: ResourceLimits::default(),
            snapshot_store: SnapshotStore::default(),
            document_store: NormalizedStore::default(),
            metrics: MetricsCollector::new(),
            logger,
        })
    }

    /// Configure the browser pool for JS rendering
    pub fn with_browser_pool(mut self, pool: BrowserPool) -> Self {
        self.browser_pool = pool;
        self
    }

    /// Scrape a single URL with full pipeline
    pub async fn scrape(&self, url: &Url) -> Result<ScrapeResult> {
        let start = std::time::Instant::now();

        // 1. Normalize URL
        let normalized_url = self.normalizer.normalize(url);
        self.logger.log_request_url(&normalized_url, "GET");

        // 2. Check URL deduplication
        if !self.url_dedup.check_and_mark(normalized_url.as_str()) {
            self.logger.log_info(&format!("URL already seen: {}", normalized_url));
            return Err(Error::DuplicateUrl(normalized_url.to_string()));
        }

        // 3. Security validation
        self.validate_url(&normalized_url)?;

        // 4. Check robots.txt
        if self.config.politeness.respect_robots_txt {
            self.check_robots(&normalized_url).await?;
        }

        // 5. Throttling
        let crawl_delay = self.robots_checker.get_crawl_delay(&normalized_url);
        self.throttler.acquire(&normalized_url, crawl_delay).await;

        // 6. Fetch
        self.metrics
            .record_request(normalized_url.host_str().unwrap_or("unknown"));
        let response = self.fetch(&normalized_url).await?;
        let duration_ms = start.elapsed().as_millis() as u64;

        // 7. Release throttle and record metrics
        self.throttler
            .release(&normalized_url, duration_ms, response.is_rate_limited());

        if response.is_success() {
            self.metrics.record_success(
                normalized_url.host_str().unwrap_or("unknown"),
                response.body_size(),
                duration_ms,
            );
        } else {
            self.metrics
                .record_failure(normalized_url.host_str().unwrap_or("unknown"), duration_ms);
            return Err(Error::Http(response.status.as_u16()));
        }

        // 8. Check resource limits
        self.limits.check_response_size(response.body_size())?;

        // 9. Store raw snapshot
        let snapshot = RawSnapshot::from_response(
            normalized_url.clone(),
            response.status.as_u16(),
            response.headers_map(),
            response.body.to_vec(),
        );
        self.snapshot_store.store(snapshot);

        // 10. Get HTML content
        let html = response
            .text()
            .map_err(|_| Error::Parse("Invalid UTF-8".to_string()))?;

        // 11. Check content deduplication
        let is_duplicate = !self.content_dedup.check_and_mark(&html);
        if is_duplicate && self.config.crawl.dedup_content {
            self.logger
                .log_info(&format!("Content duplicate: {}", normalized_url));
        }

        // 12. Check if JS rendering needed
        let render_decision = self.render_checker.check(&html);
        let (final_html, js_rendered) = if render_decision != RenderDecision::Static {
            // Try browser rendering if available
            match self
                .browser_pool
                .render(&normalized_url, None)
                .await
            {
                Ok(browser_response) => (browser_response.html, true),
                Err(_) => {
                    // Fall back to static HTML
                    self.logger.log_warn(&format!(
                        "JS rendering needed but browser unavailable: {}",
                        normalized_url
                    ));
                    (html, false)
                }
            }
        } else {
            (html, false)
        };

        // 13. Parse and extract
        let mut document = self
            .process_response(&normalized_url, &response, &final_html)
            .await?;

        // 14. Resolve canonical URL
        let canonical_url = if self.config.crawl.respect_canonicals {
            self.canonical_resolver
                .resolve_from_html(&final_html, &response.final_url)
        } else {
            None
        };

        if let Some(ref canonical) = canonical_url {
            document.canonical_url = Some(canonical.clone());
            // Mark canonical as seen to avoid re-crawling
            self.url_dedup.mark_seen(canonical.as_str());
        }

        // 15. Extract discovered links for crawling
        let discovered_links: Vec<Url> = document
            .out_links
            .iter()
            .filter_map(|link| {
                let url = self.normalizer.normalize(&link.url);
                // Only include if not seen and allowed
                if self.url_dedup.is_duplicate(url.as_str()) {
                    None
                } else if !self.allowlist.is_allowed(&url) {
                    None
                } else {
                    Some(url)
                }
            })
            .collect();

        // 16. Store normalized document
        self.document_store.store(document.clone());
        self.metrics.record_document();

        self.logger.log_response_parts(
            &normalized_url,
            response.status.as_u16(),
            response.body_size(),
            duration_ms,
        );

        Ok(ScrapeResult {
            document,
            is_duplicate,
            js_rendered,
            canonical_url,
            discovered_links,
        })
    }

    /// Crawl starting from seed URLs
    pub async fn crawl(&self, seeds: Vec<Url>, max_pages: Option<usize>) -> Result<Vec<Document>> {
        // Add seeds to frontier
        for seed in seeds {
            let normalized = self.normalizer.normalize(&seed);
            self.frontier.push(CrawlEntry::new(normalized, 0, 100));
        }

        let mut documents = Vec::new();
        let max = max_pages.unwrap_or(usize::MAX);

        while let Some(entry) = self.frontier.pop() {
            if documents.len() >= max {
                break;
            }

            match self.scrape(&entry.url).await {
                Ok(result) => {
                    // Add discovered links to frontier
                    for link in result.discovered_links {
                        self.frontier.push(CrawlEntry::new(
                            link,
                            entry.depth + 1,
                            50, // Lower priority for deeper pages
                        ));
                    }
                    documents.push(result.document);
                }
                Err(e) => {
                    self.logger.log_error_parts(&entry.url, &e.to_string(), true);
                }
            }
        }

        Ok(documents)
    }

    /// Validate a URL before scraping
    fn validate_url(&self, url: &Url) -> Result<()> {
        // Check scheme
        if !url.scheme().starts_with("http") {
            return Err(Error::Config(format!(
                "Unsupported scheme: {}",
                url.scheme()
            )));
        }

        // Check allowlist
        if !self.allowlist.is_allowed(url) {
            return Err(Error::DomainNotAllowed(
                url.host_str().unwrap_or("").to_string(),
            ));
        }

        // Check blocked IPs
        if self.ip_blocker.is_url_hostname_blocked(url) {
            return Err(Error::IpBlocked(
                url.host_str().unwrap_or("").to_string(),
            ));
        }

        Ok(())
    }

    /// Check robots.txt
    async fn check_robots(&self, url: &Url) -> Result<()> {
        // Check cache first
        if self.robots_checker.cache().get(url).is_none() {
            // Fetch robots.txt
            if let Some(robots_url) = RobotsChecker::robots_url(url) {
                if let Ok(response) = self.fetch(&robots_url).await {
                    if response.is_success() {
                        if let Ok(text) = response.text() {
                            self.robots_checker.cache_robots(url, &text);
                        }
                    }
                }
            }
        }

        if !self.robots_checker.is_allowed(url, None) {
            return Err(Error::RobotsBlocked(url.to_string()));
        }

        Ok(())
    }

    /// Fetch a URL
    async fn fetch(&self, url: &Url) -> Result<FetchResponse> {
        let request = crate::fetch::RequestBuilder::new(url.clone()).with_config(&self.config);

        let (url, headers) = request.build();

        let response = self
            .client
            .client()
            .get(url.as_str())
            .headers(headers)
            .send()
            .await?;

        let status = response.status();
        let headers = response.headers().clone();
        let final_url = response.url().clone();
        let body = response.bytes().await?;

        Ok(FetchResponse::new(final_url, status, headers, body))
    }

    /// Process an HTTP response
    async fn process_response(
        &self,
        source_url: &Url,
        response: &FetchResponse,
        html: &str,
    ) -> Result<Document> {
        let parse_start = std::time::Instant::now();

        // Create base document
        let mut document = Document::new(source_url.clone(), response.final_url.clone());
        document.status_code = response.status.as_u16();
        document.content_type = response.content_type();
        document.content_length = Some(response.body_size());

        // Provenance
        document.provenance.response_headers = response.headers_map();
        document.provenance.etag = response.etag();
        document.provenance.last_modified = response.last_modified();
        document.provenance.cache_control = response.cache_control();
        document.provenance.content_hash =
            Some(crate::crawl::dedup::ContentDedup::hash_content(html));
        document.provenance.timings = response.timings.clone();

        // Check parsing time limit
        let check_time = || {
            let elapsed = parse_start.elapsed().as_millis() as u64;
            self.limits.check_parse_time(elapsed)
        };

        // Parse HTML
        let _parser = HtmlParser::parse(html);

        // Metadata
        let metadata = MetadataExtractor::new().extract(html);
        document.title = metadata.title;
        document.language = metadata.language;

        check_time()?;

        // Main text
        let text_extractor = TextExtractor::default()
            .with_chunking(self.config.parse.segment_text, self.config.parse.chunk_size);
        let extracted = text_extractor.extract(html);
        document.main_text = extracted.full_text;
        document.provenance.text_hash = Some(crate::crawl::dedup::ContentDedup::hash_content(
            &document.main_text,
        ));

        check_time()?;

        // Links
        if self.config.parse.extract_links {
            let link_extractor = LinkExtractor::default();
            document.out_links = link_extractor.extract(html, &response.final_url);
        }

        check_time()?;

        // Structured data
        if self.config.parse.extract_json_ld {
            let jsonld = JsonLdExtractor::new().extract(html);
            document.structured_data.json_ld = jsonld;
        }

        if self.config.parse.extract_open_graph {
            let og = OpenGraphExtractor::new().extract(html);
            document.structured_data.open_graph = og;
        }

        check_time()?;

        // Robots directives
        document.structured_data.robots_directives = self.parse_robots_directives(html, response);

        // Assets
        if self.config.parse.extract_images {
            let img_extractor = ImageExtractor::default()
                .with_options(self.config.parse.resolve_lazy_loading, false);
            document.assets.images = img_extractor.extract(html, &response.final_url);
        }

        if self.config.parse.extract_videos {
            let video_extractor = VideoExtractor::new();
            document.assets.videos = video_extractor.extract(html, &response.final_url);
        }

        if self.config.parse.extract_audios {
            let audio_extractor = AudioExtractor::new();
            document.assets.audios = audio_extractor.extract(html, &response.final_url);
        }

        Ok(document)
    }

    /// Parse robots directives
    fn parse_robots_directives(&self, html: &str, response: &FetchResponse) -> RobotsDirectives {
        let parser = HtmlParser::parse(html);

        let meta_robots = parser.attr(r#"meta[name="robots"]"#, "content");
        let x_robots_tag = response.x_robots_tag();

        let combined = format!(
            "{} {}",
            meta_robots.as_deref().unwrap_or(""),
            x_robots_tag.as_deref().unwrap_or("")
        )
        .to_lowercase();

        RobotsDirectives {
            meta_robots,
            x_robots_tag,
            index: !combined.contains("noindex"),
            follow: !combined.contains("nofollow"),
            archive: !combined.contains("noarchive"),
            snippet: !combined.contains("nosnippet"),
            image_index: !combined.contains("noimageindex"),
        }
    }

    /// Access to metrics
    pub fn metrics(&self) -> &MetricsCollector {
        &self.metrics
    }

    /// Access to document store
    pub fn documents(&self) -> &NormalizedStore {
        &self.document_store
    }

    /// Access to snapshot store
    pub fn snapshots(&self) -> &SnapshotStore {
        &self.snapshot_store
    }

    /// Access to frontier
    pub fn frontier(&self) -> &Frontier {
        &self.frontier
    }

    /// Access to URL dedup
    pub fn url_dedup(&self) -> &UrlDedup {
        &self.url_dedup
    }

    /// Access to content dedup
    pub fn content_dedup(&self) -> &ContentDedup {
        &self.content_dedup
    }

    /// Configure the allowlist
    pub fn allowlist_mut(&mut self) -> &mut DomainAllowlist {
        &mut self.allowlist
    }

    /// Configure IP blocker
    pub fn ip_blocker_mut(&mut self) -> &mut IpBlocker {
        &mut self.ip_blocker
    }

    /// Configure resource limits
    pub fn limits_mut(&mut self) -> &mut ResourceLimits {
        &mut self.limits
    }

    /// Get configuration
    pub fn config(&self) -> &Config {
        &self.config
    }

    /// Get logger
    pub fn logger(&self) -> &StructuredLogger {
        &self.logger
    }
}