#!/usr/bin/env python3
"""Toolset distributions for batch data-generation runs.

A distribution maps toolset names to the % chance each is enabled for a prompt
(sampled independently, so several toolsets can be active at once).
"""

from typing import Dict, List, Optional
import random
from toolsets import validate_toolset


def _dist(description: str, **toolsets: int) -> Dict[str, object]:
    return {"description": description, "toolsets": toolsets}


DISTRIBUTIONS = {
    "default": _dist("All available tools, all the time",
                     web=100, vision=100, image_gen=100, terminal=100, file=100, browser=100),
    "image_gen": _dist("Heavy focus on image generation with vision and web support",
                       image_gen=90, vision=90, web=55, terminal=45),
    "research": _dist("Web research with vision analysis and reasoning",
                      web=90, browser=70, vision=50, terminal=10),
    "science": _dist("Scientific research with web, terminal, file, and browser capabilities",
                     web=94, terminal=94, file=94, vision=65, browser=50, image_gen=15),
    "development": _dist("Terminal, file tools, and reasoning with occasional web lookup",
                         terminal=80, file=80, web=30, vision=10),
    "safe": _dist("All tools except terminal for safety",
                  web=80, browser=70, vision=60, image_gen=60),
    "balanced": _dist("Equal probability of all toolsets",
                      web=50, vision=50, image_gen=50, terminal=50, file=50, browser=50),
    "minimal": _dist("Only web tools for basic research", web=100),
    "terminal_only": _dist("Terminal and file tools for code execution tasks", terminal=100, file=100),
    "terminal_web": _dist("Terminal and file tools with web search for documentation lookup",
                          terminal=100, file=100, web=100),
    "creative": _dist("Image generation and vision analysis focus", image_gen=90, vision=90, web=30),
    "reasoning": _dist("Heavy research/reasoning distribution with minimal other tools",
                       web=90, file=60, terminal=20),
    "browser_use": _dist("Full browser-based web interaction with search, vision, and page control",
                         browser=100, web=80, vision=70),
    "browser_only": _dist("Only browser automation tools for pure web interaction tasks", browser=100),
    # browser-use-tasks.jsonl: the browser toolset includes web_search since Google blocks direct browser searches
    "browser_tasks": _dist(
        "Browser-focused distribution (browser toolset includes web_search for finding URLs since Google blocks direct browser searches)",
        browser=97, vision=12, terminal=15,
    ),
    # nous-terminal-tasks.jsonl
    "terminal_tasks": _dist(
        "Terminal-focused distribution with high terminal/file availability, occasional other tools",
        terminal=97, file=97, web=97, browser=75, vision=50, image_gen=10,
    ),
    # mixed-browser-terminal-tasks.jsonl
    "mixed_tasks": _dist(
        "Mixed distribution with high browser, terminal, and file availability for complex tasks",
        browser=92, terminal=92, file=92, web=35, vision=15, image_gen=15,
    ),
}


def get_distribution(name: str) -> Optional[Dict[str, any]]:
    """Distribution definition (description + toolsets), or None if unknown."""
    return DISTRIBUTIONS.get(name)


def list_distributions() -> Dict[str, Dict]:
    return DISTRIBUTIONS.copy()


def sample_toolsets_from_distribution(distribution_name: str) -> List[str]:
    """Sample toolset names, each included independently with its % probability.

    Falls back to the highest-probability toolset when nothing was rolled.
    Raises ValueError for an unknown distribution.
    """
    dist = get_distribution(distribution_name)
    if not dist:
        raise ValueError(f"Unknown distribution: {distribution_name}")

    selected_toolsets = []
    for toolset_name, probability in dist["toolsets"].items():
        if not validate_toolset(toolset_name):
            print(f"⚠️  Warning: Toolset '{toolset_name}' in distribution '{distribution_name}' is not valid")
            continue
        if random.random() * 100 < probability:
            selected_toolsets.append(toolset_name)

    if not selected_toolsets and dist["toolsets"]:
        highest_prob_toolset = max(dist["toolsets"].items(), key=lambda x: x[1])[0]
        if validate_toolset(highest_prob_toolset):
            selected_toolsets.append(highest_prob_toolset)

    return selected_toolsets


def validate_distribution(distribution_name: str) -> bool:
    return distribution_name in DISTRIBUTIONS


def print_distribution_info(distribution_name: str) -> None:
    """Print a distribution's description and toolset probabilities (highest first)."""
    dist = get_distribution(distribution_name)
    if not dist:
        print(f"❌ Unknown distribution: {distribution_name}")
        return

    print(f"\n📊 Distribution: {distribution_name}")
    print(f"   Description: {dist['description']}")
    print("   Toolsets:")
    for toolset, prob in sorted(dist["toolsets"].items(), key=lambda x: x[1], reverse=True):
        print(f"     • {toolset:15} : {prob:3}% chance")