"""
URL Harvesting Module - Parallel execution for maximum speed
"""

import asyncio
import subprocess
from pathlib import Path
from typing import List, Set, Dict
from rich.console import Console
from concurrent.futures import ThreadPoolExecutor

from ..utils.logger import log_progress
from ..utils.hashing import DeduplicatorHash

console = Console()

class URLHarvester:
    """Harvest URLs from multiple sources in parallel"""
    
    def __init__(self, output_dir: Path):
        self.output_dir = output_dir
        self.deduplicator = DeduplicatorHash()
        self.timeout = 120  # 2 minutes per tool
        
    async def run_tool_async(self, tool_name: str, domain: str) -> List[str]:
        """Run a single tool asynchronously"""
        try:
            if tool_name == "gau":
                cmd = ["gau", domain]
            elif tool_name == "waybackurls":
                cmd = ["waybackurls", domain]
            elif tool_name == "hakrawler":
                cmd = ["hakrawler", "-url", domain, "-depth", "2", "-plain"]
            elif tool_name == "katana":
                cmd = ["katana", "-u", domain, "-depth", "2", "-silent"]
            elif tool_name == "subjs":
                cmd = ["subjs", "-d", domain]
            else:
                return []
            
            # Run tool with timeout
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(), 
                    timeout=self.timeout
                )
                
                if process.returncode == 0 and stdout:
                    urls = [url.strip() for url in stdout.decode('utf-8', errors='ignore').split('\n') if url.strip()]
                    log_progress(f"{tool_name}: found {len(urls)} URLs for {domain}")
                    return urls
                else:
                    log_progress(f"{tool_name}: no results for {domain}")
                    return []
                    
            except asyncio.TimeoutError:
                process.kill()
                log_progress(f"{tool_name}: timeout for {domain}")
                return []
                
        except Exception as e:
            log_progress(f"{tool_name}: error for {domain} - {e}")
            return []
    
    async def harvest_domain_parallel(self, domain: str) -> List[str]:
        """Harvest URLs from a single domain using all tools in parallel"""
        tools = ["gau", "waybackurls", "hakrawler", "katana", "subjs"]
        
        # Run all tools in parallel
        tasks = [self.run_tool_async(tool, domain) for tool in tools]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Collect all URLs
        all_urls = []
        for i, result in enumerate(results):
            if isinstance(result, list):
                all_urls.extend(result)
            else:
                log_progress(f"{tools[i]}: exception for {domain} - {result}")
        
        return all_urls
    
    async def harvest_all_domains_parallel(self, domains: List[str]) -> List[str]:
        """Harvest URLs from all domains in parallel"""
        log_progress(f"Harvesting URLs from {len(domains)} domains in parallel")
        
        # Limit concurrent domains to avoid overwhelming the system
        semaphore = asyncio.Semaphore(3)  # Max 3 domains at once
        
        async def harvest_with_semaphore(domain):
            async with semaphore:
                return await self.harvest_domain_parallel(domain)
        
        # Run domain harvesting in parallel
        tasks = [harvest_with_semaphore(domain) for domain in domains]
        domain_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Collect all URLs
        all_urls = []
        for i, result in enumerate(domain_results):
            if isinstance(result, list):
                all_urls.extend(result)
                log_progress(f"Domain {domains[i]}: {len(result)} URLs")
            else:
                log_progress(f"Domain {domains[i]}: failed - {result}")
        
        return all_urls
    
    def harvest_urls(self, domains: List[str]) -> List[str]:
        """
        Harvest URLs from all sources in parallel (main entry point)
        
        Args:
            domains: List of domains to harvest
            
        Returns:
            Deduplicated list of URLs
        """
        log_progress(">> Starting PARALLEL URL harvesting (gau, waybackurls, hakrawler, katana, subjs)")
        
        # Run async harvesting
        try:
            all_urls = asyncio.run(self.harvest_all_domains_parallel(domains))
        except Exception as e:
            log_progress(f"Parallel harvesting failed, falling back to sequential: {e}")
            return self.harvest_urls_sequential(domains)
        
        # Deduplicate URLs
        unique_urls = self.deduplicator.deduplicate_list(all_urls)
        
        log_progress(f">> PARALLEL harvest complete: {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
    
    def harvest_urls_sequential(self, domains: List[str]) -> List[str]:
        """Fallback sequential harvesting if parallel fails"""
        log_progress("Running sequential URL harvesting (fallback)")
        
        all_urls = []
        tools = [
            ("gau", lambda d: ["gau", d]),
            ("waybackurls", lambda d: ["waybackurls", d]),
            ("hakrawler", lambda d: ["hakrawler", "-url", d, "-depth", "2", "-plain"]),
            ("katana", lambda d: ["katana", "-u", d, "-depth", "2", "-silent"]),
            ("subjs", lambda d: ["subjs", "-d", d])
        ]
        
        for domain in domains:
            for tool_name, cmd_func in tools:
                try:
                    result = subprocess.run(
                        cmd_func(domain),
                        capture_output=True,
                        text=True,
                        timeout=self.timeout
                    )
                    
                    if result.returncode == 0 and result.stdout:
                        urls = [url.strip() for url in result.stdout.split('\n') if url.strip()]
                        all_urls.extend(urls)
                        
                except Exception as e:
                    log_progress(f"{tool_name} failed for {domain}: {e}")
        
        return self.deduplicator.deduplicate_list(all_urls)