from pathlib import Path
from typing import List

from fsm_crawl.agent.generic_url_manager import GenericUrlManager

class FileUrlManager(GenericUrlManager):
    def __init__(self, path: str | Path, limit: int | None = None):
        self.path = Path(path)
        self.urls: List[str] = []
        self._idx = 0

        with self.path.open() as f:
            for line in f:
                if not line.strip():
                    continue
                _, url = line.strip().split(",", 1)
                self.urls.append(url)
                if limit and len(self.urls) >= limit:
                    break

    def get_next_url(self) -> str:
        if self._idx >= len(self.urls):
            raise StopIteration("No more URLs to crawl")

        url = self.urls[self._idx]
        self._idx += 1
        return url

    def get_crawl_size(self) -> int:
        return len(self.urls)
    
    def set_sharding(self, shard_index: int, shard_count: int) -> None:
        """Set sharding for distributed crawling."""
        self.urls = [url for i, url in enumerate(self.urls) if i % shard_count == shard_index]
        self._idx = 0
