from fsm_crawl.agent.file_url_manager import FileUrlManager
from fsm_crawl.agent.generic_agent import GenericAgent
from fsm_crawl.agent.generic_url_manager import GenericUrlManager
from fsm_crawl.agent.probability_based_agent import ProbabilityBasedAgent
from fsm_crawl.browser.engine.generic_browser_engine import GenericBrowserEngine
from fsm_crawl.browser.engine.playwright_engine import PlaywrightEngine
from fsm_crawl.experiment.generic_expirement import GenericExperiment
from fsm_crawl.monitoring.callback_pipeline import GenericCallbackPipeline
from fsm_crawl.blocking.labeler import Labeler
from fsm_crawl.monitoring.response_logging_pipeline import RequestResponseLoggingPipeline


class ExplorativeCrawlExpirement(GenericExperiment):
    def __init__(self, path: str, prefix: str, headless: bool = True, shard_index: int = 0, shard_count: int = 1, num_tabs: int = 10):
        self.agent = ProbabilityBasedAgent()
        self.labeler = Labeler()
        self.pipeline = RequestResponseLoggingPipeline(self.labeler, output_dir="crawl"+f"_shard{shard_index}", filename_prefix=prefix+f"_shard{shard_index}")
        self.url_manager = FileUrlManager(path=path)
        self.url_manager.set_sharding(shard_index=shard_index, shard_count=shard_count)
        self.engine = PlaywrightEngine(headless=headless, num_tabs=num_tabs)
    
    def get_url_manager(self) -> GenericUrlManager:
        return self.url_manager
    
    def get_agent(self) -> GenericAgent:
        return self.agent
    
    def get_engine(self) -> GenericBrowserEngine:
        return self.engine
    
    def get_pipeline(self) -> GenericCallbackPipeline:
        return self.pipeline
        
    def get_name(self):
        return "ExplorativeCrawl"
    
    def get_max_depth(self) -> int:
        return 8
    
    def has_consent_pass(self) -> bool:
        return True  
    
    def get_expirement_split(self) -> None | int:
        return 999