import pandas as pd
import random
from typing import Dict, Optional


class URLManager:
    """
    Provides starting URLs and next URLs for a simulated user crawl.
    Only includes English ('en') categories.
    Supports weighted stochastic selection based on user personality.
    """

    def __init__(self, csv_path: str = "blocking/curlie_filtered.csv", seed: Optional[int] = None):
        if seed is not None:
            random.seed(seed)

        # Load dataset
        self.df = pd.read_csv(csv_path)
        self.df["label"] = self.df["label"].fillna("")
        self.df["top_category"] = self.df["label"].apply(self._get_top_category)

        # Filter only English categories
        self.df = self.df[self.df["top_category"].notna()]

        # Map category -> list of URLs
        self.category_map = {}
        for cat, group in self.df.groupby("top_category"):
            self.category_map[cat] = list(group["url"].unique())

        # Flatten all URLs
        self.all_urls = list(self.df["url"].unique())

    @staticmethod
    def _get_top_category(label: str) -> Optional[str]:
        """
        Extract the first-level category for English URLs only.
        Returns None if not English.
        Example:
        - /en/Computers/Hardware/Components/Motherboards -> Computers
        - /zh_TW/休閒/寵物 -> None
        """
        parts = label.strip("/").split("/")
        if len(parts) >= 2 and parts[0].lower() == "en":
            return parts[1]  # first segment after "en"
        return None

    def get_random_url(self, personality: dict = None) -> str:
        if not personality:
            return random.choice(self.all_urls)

        categories = list(personality.keys())
        weights = [personality[cat] for cat in categories]
        total = sum(weights)
        weights = [w / total for w in weights]

        # Select category first using personality weights
        chosen_category = random.choices(categories, weights=weights, k=1)[0]

        urls_in_cat = self.category_map.get(chosen_category, [])
        if not urls_in_cat:
            return random.choice(self.all_urls)
        return random.choice(urls_in_cat)

    def get_multiple_urls(self, personality: Optional[Dict[str, float]] = None, n: int = 1) -> list[str]:
        return [self.get_random_url(personality) for _ in range(n)]

    def list_categories(self) -> list[str]:
        return sorted(self.category_map.keys())

    def get_urls_for_category(self, category: str) -> list[str]:
        return self.category_map.get(category, [])
