from playwright.async_api import async_playwright, Page, Browser
from fsm_crawl.browser.engine.generic_browser_engine import GenericBrowserEngine
from fsm_crawl.experiment.generic_expirement import GenericExperiment
from tqdm import tqdm
import asyncio

class BrowserManager:
    _ERROR_MESSAGE = "Browser not started. Use 'with BrowserManager(...) as bm:'"
    _ERROR_EXPERIMENT_MISSING_MESSAGE = "Cannot start crawl without experiment"
    def __init__(self, engine: GenericBrowserEngine):
        self.experiment: GenericExperiment = None
        self.engine = engine
        self.current_depth = 0
        
    async def __aenter__(self):
        self.engine.install()
        await self.engine.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.engine.stop()
        if self.experiment.get_pipeline():
            await self.experiment.get_pipeline().on_crawl_end()
            
    async def start_experiment(self, expirement: GenericExperiment, tqdm_position: int = 0):
        if not self.engine.is_ready():
            raise RuntimeError(self._ERROR_MESSAGE)
        if not expirement:
            raise RuntimeError(self._ERROR_EXPERIMENT_MISSING_MESSAGE)
        self.experiment = expirement
        
        if self.experiment.get_pipeline():
            self.engine.on_response(self.experiment.get_pipeline().on_response)
            await self.engine.on_request(self.experiment.get_pipeline().on_request)
            
        if self.experiment.has_consent_pass():
            await self.engine.start_consent_task()
        
        # Get all URLs to crawl
        all_urls = []
        for _ in range(self.experiment.get_url_manager().get_crawl_size()):
            try:
                url = self.experiment.get_url_manager().get_next_url()
                if url:
                    all_urls.append(url)
            except Exception:
                continue
        
        # Crawl all URLs in batches using parallel tabs
        num_tabs = len(self.engine.pages) if hasattr(self.engine, 'pages') else 1
        with tqdm(total=len(all_urls), desc=self.experiment.get_name(), position=tqdm_position, leave=True) as pbar:
            for batch_idx in range(0, len(all_urls), num_tabs):
                batch = all_urls[batch_idx:batch_idx + num_tabs]
                # Process batch in parallel across tabs
                tasks = []
                for tab_idx, url in enumerate(batch):
                    page = self.engine.pages[tab_idx] if hasattr(self.engine, 'pages') else None
                    task = asyncio.create_task(self._perform_agent_iteration_on_page(url, page))
                    tasks.append(task)
                
                # Wait for entire batch to complete
                await asyncio.gather(*tasks, return_exceptions=True)
                pbar.update(len(batch))
  
        if self.experiment.has_consent_pass():
             await self.engine.stop_consent_task()
                       
    async def _perform_agent_iteration_on_page(self, url, page=None):
        """Perform a single iteration on a specific page."""
        try:
            links = await self.engine.goto(url, page=page)

            scroll_amounts = self.experiment.get_agent().scroll()
            dwell_time = self.experiment.get_agent().dwell()

            scroll_task = asyncio.create_task(self.engine.perform_scroll(scroll_amounts, page=page))
            choose_task = asyncio.to_thread(self.experiment.get_agent().choose_next, links, 0)
            next_link, _, = await asyncio.gather(choose_task, scroll_task)
            await asyncio.sleep(dwell_time)
            return next_link
        except Exception as e:
            return None
                       
    async def _perform_agent_iteration(self, url):
        try:
            links = await self.engine.goto(url)

            scroll_amounts = self.experiment.get_agent().scroll()
            dwell_time = self.experiment.get_agent().dwell()

            scroll_task = asyncio.create_task(self.engine.perform_scroll(scroll_amounts))
            choose_task = asyncio.to_thread(self.experiment.get_agent().choose_next, links, self.current_depth)
            next_link, _, = await asyncio.gather(choose_task, scroll_task)
            await asyncio.sleep(dwell_time)
            self.current_depth += 1
            return next_link
        except Exception as e:
            return None
                    
