from __future__ import annotations

import random
from typing import Any

try:
    from . import cast_context as cast_context_policy
    from . import character_appearance as character_appearance_policy
    from . import generate_prompt_batches as g
except ImportError:  # Allows local smoke tests with top-level imports.
    import cast_context as cast_context_policy
    import character_appearance as character_appearance_policy
    import generate_prompt_batches as g


def couple_context(
    rng: random.Random,
    women_count: int,
    men_count: int,
) -> dict[str, str]:
    primary_subject, subject_phrase, pose, effective_women_count, effective_men_count = cast_context_policy.couple_type_from_counts(
        rng,
        women_count,
        men_count,
        choose=g.choose,
        couple_types=g.COUPLE_TYPES,
    )
    return {
        "subject_type": "couple",
        "subject": primary_subject,
        "subject_phrase": subject_phrase,
        "age": g.choose(rng, g.COUPLE_AGES),
        "body": g.choose(rng, ["slim and average", "curvy and broad", "stocky and curvy", "average and athletic"]),
        "skin": "",
        "hair": "",
        "eyes": "",
        "body_phrase": "",
        "fallback_pose": pose,
        "women_count": str(effective_women_count),
        "men_count": str(effective_men_count),
        "person_count": "2",
    }


def group_context(rng: random.Random, ethnicity: str) -> dict[str, str]:
    eth = "Asian " if ethnicity == "asian" else ""
    return {
        "subject_type": "group",
        "subject": f"mixed {eth}adult group",
        "subject_phrase": f"A mixed {eth}adult group of women and men",
        "age": g.choose(rng, g.GROUP_AGES),
        "body": "diverse",
        "skin": "",
        "hair": "",
        "eyes": "",
        "body_phrase": "diverse adult body types",
    }


def layout_context(subject_type: str) -> dict[str, str]:
    return {
        "subject_type": subject_type,
        "subject": "layout scene",
        "subject_phrase": "Adult layout scene",
        "age": "adult",
        "body": "varied",
        "skin": "",
        "hair": "",
        "eyes": "",
        "body_phrase": "varied adult figures",
    }


def subject_context(
    rng: random.Random,
    subject_type: str,
    ethnicity: str,
    figure: str,
    no_plus_women: bool,
    no_black: bool,
    women_count: int = 1,
    men_count: int = 1,
) -> dict[str, str]:
    if subject_type in ("woman", "man", "single_any"):
        return character_appearance_policy.appearance_for_subject(
            rng,
            subject_type,
            ethnicity,
            figure,
            no_plus_women,
            no_black,
        )

    if subject_type == "configured_cast":
        return cast_context_policy.configured_cast_context(women_count, men_count)

    if subject_type == "couple":
        return couple_context(rng, women_count, men_count)

    if subject_type == "group":
        return group_context(rng, ethnicity)

    return layout_context(subject_type)