"""
Part 1 entry point: generate images for selected prompts with a chosen
model, one image at a time (no batching, no analysis during generation).

Usage:
    from src.prompts import load_prompt_bank, select_prompts
    from src.providers import PROVIDERS, MockProvider
    from src.generate import run_generation

    bank = load_prompt_bank("data/prompt_bank.csv")
    prompts = select_prompts(bank, categories=["occupation"], n_per_category=5, seed=1)

    provider = MockProvider()  # or PROVIDERS["FX"]()
    manifest_path, df = run_generation(
        prompts=prompts,
        provider=provider,
        images_per_prompt=10,
        output_dir="runs/my_first_audit",
    )
"""
import os

from .manifest import ManifestWriter


def run_generation(prompts, provider, images_per_prompt, output_dir, on_progress=None):
    """
    prompts: list[Prompt] from select_prompts()
    provider: an instance of a ModelProvider subclass
    images_per_prompt: int, 5-100
    output_dir: where images + manifest chunks/final file are written
    on_progress: optional callback(done, total) -- e.g. for a Gradio progress bar

    Returns: (manifest_path, manifest_dataframe)
    """
    if not (5 <= images_per_prompt <= 100):
        raise ValueError("images_per_prompt must be between 5 and 100")

    images_dir = os.path.join(output_dir, "images")
    os.makedirs(images_dir, exist_ok=True)

    writer = ManifestWriter(output_dir)

    total = len(prompts) * images_per_prompt
    done = 0

    for prompt in prompts:
        prompt_dir = os.path.join(images_dir, prompt.category, f"{prompt.prompt_id:03d}")
        os.makedirs(prompt_dir, exist_ok=True)

        for image_index in range(1, images_per_prompt + 1):
            image_id = f"{provider.model_key}_{prompt.category}_{prompt.prompt_id:03d}_{image_index:03d}"

            if writer.is_done(image_id):
                done += 1
                if on_progress:
                    on_progress(done, total)
                continue

            image_bytes = provider.generate(prompt.prompt_text)

            file_path = os.path.join(prompt_dir, f"{image_index:03d}.png")
            with open(file_path, "wb") as f:
                f.write(image_bytes)

            writer.add_row(
                model=provider.model_key,
                category=prompt.category,
                prompt_id=prompt.prompt_id,
                prompt_text=prompt.prompt_text,
                image_index=image_index,
                file_path=file_path,
            )

            done += 1
            if on_progress:
                on_progress(done, total)

    manifest_path, df = writer.finalize()
    return manifest_path, df