"""
URL Harvesting Module - Collect URLs using multiple tools
"""

from pathlib import Path
from typing import List, Set
from rich.console import Console

from ..utils.shell import run_command, run_with_input_file
from ..utils.logger import log_progress
from ..utils.hashing import DeduplicatorHash

console = Console()

class URLHarvester:
    """Harvest URLs from multiple sources"""
    
    def __init__(self, output_dir: Path):
        self.output_dir = output_dir
        self.deduplicator = DeduplicatorHash()
        
    def run_gau(self, domains: List[str]) -> List[str]:
        """Run gau tool"""
        log_progress("Running gau...")
        
        all_urls = []
        for domain in domains:
            success, stdout, stderr = run_command(["gau", domain])
            if success and stdout:
                urls = [url.strip() for url in stdout.split('\n') if url.strip()]
                all_urls.extend(urls)
        
        return all_urls
    
    def run_waybackurls(self, domains: List[str]) -> List[str]:
        """Run waybackurls tool"""
        log_progress("Running waybackurls...")
        
        all_urls = []
        for domain in domains:
            success, stdout, stderr = run_command(["waybackurls", domain])
            if success and stdout:
                urls = [url.strip() for url in stdout.split('\n') if url.strip()]
                all_urls.extend(urls)
        
        return all_urls
    
    def run_hakrawler(self, domains: List[str]) -> List[str]:
        """Run hakrawler tool"""
        log_progress("Running hakrawler...")
        
        all_urls = []
        for domain in domains:
            success, stdout, stderr = run_command([
                "hakrawler", "-url", domain, "-depth", "2", "-plain"
            ])
            if success and stdout:
                urls = [url.strip() for url in stdout.split('\n') if url.strip()]
                all_urls.extend(urls)
        
        return all_urls
    
    def run_katana(self, domains: List[str]) -> List[str]:
        """Run katana tool"""
        log_progress("Running katana...")
        
        domain_input = '\n'.join(domains)
        success, stdout, stderr = run_command(
            ["katana", "-list", "-", "-depth", "2", "-silent"],
            input_data=domain_input
        )
        
        if success and stdout:
            return [url.strip() for url in stdout.split('\n') if url.strip()]
        
        return []
    
    def harvest_urls(self, domains: List[str]) -> List[str]:
        """
        Harvest URLs from all sources
        
        Args:
            domains: List of domains to harvest
            
        Returns:
            Deduplicated list of URLs
        """
        log_progress("Harvesting URLs (gau, waybackurls, hakrawler, katana)")
        
        all_urls = []
        
        # Run all harvesting tools
        try:
            all_urls.extend(self.run_gau(domains))
        except Exception as e:
            console.print(f"[yellow]Warning: gau failed - {e}[/yellow]")
        
        try:
            all_urls.extend(self.run_waybackurls(domains))
        except Exception as e:
            console.print(f"[yellow]Warning: waybackurls failed - {e}[/yellow]")
        
        try:
            all_urls.extend(self.run_hakrawler(domains))
        except Exception as e:
            console.print(f"[yellow]Warning: hakrawler failed - {e}[/yellow]")
        
        try:
            all_urls.extend(self.run_katana(domains))
        except Exception as e:
            console.print(f"[yellow]Warning: katana failed - {e}[/yellow]")
        
        # Deduplicate URLs
        unique_urls = self.deduplicator.deduplicate_list(all_urls)
        
        log_progress(f"Harvested {len(unique_urls)} unique URLs from {len(all_urls)} total")
        
        # Save raw URLs
        urls_file = self.output_dir / "harvested_urls.txt"
        with open(urls_file, 'w') as f:
            for url in unique_urls:
                f.write(f"{url}\n")
        
        return unique_urls