from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable

try:
    from . import krea2_pose_variant_catalog
except ImportError:  # Allows local smoke tests with top-level imports.
    import krea2_pose_variant_catalog


MOUTH_EXPRESSION_TERMS = ("mouth", "oral", "tongue", "lips", "gagging", "saliva", "drool")
TOP_VIEW_ORAL_VARIANT = "pov_blowjob_top_down_vertical_shaft"
ORAL_CONTACT_VARIANTS = frozenset(
    (
        TOP_VIEW_ORAL_VARIANT,
        "pov_blowjob_side_profile_oral",
        "pov_blowjob_laying_frontal_oral",
        "pov_blowjob_sitting_upright_oral",
    )
)


@dataclass(frozen=True)
class KreaConfiguredCastRequest:
    row: dict[str, Any]
    detail_level: str
    style_mode: str
    primary: str
    item: str
    scene: str
    expression: str
    composition: str
    source_composition: str
    camera: str
    camera_scene: str
    style: str


@dataclass(frozen=True)
class KreaConfiguredCastPrompt:
    prompt: str
    method: str = "metadata(configured_cast)"

    def as_tuple(self) -> tuple[str, str]:
        return self.prompt, self.method


@dataclass(frozen=True)
class KreaConfiguredCastDependencies:
    clean: Callable[[Any], str]
    sanitize_hardcore_environment_anchors: Callable[[Any], str]
    sanitize_hardcore_axis_values: Callable[[Any], Any]
    sanitize_scene_text_for_cast: Callable[[Any, list[str]], str]
    normalize_hardcore_detail_density: Callable[[Any], str]
    row_action_family: Callable[[Any], str]
    hardcore_action_sentence: Callable[[str, str, str, Any, str, str], str]
    pov_action_phrase: Callable[[str, list[str], str, str, str, Any, str], str]
    pov_labels_from_value: Callable[[Any], list[str]]
    merge_labels: Callable[..., list[str]]
    cast_prose_omit: Callable[[str, list[str]], tuple[str, list[str]]]
    filter_pov_labeled_clauses: Callable[[Any, list[str]], str]
    natural_label_text: Callable[[Any, list[str]], str]
    pov_composition_text: Callable[[Any, list[str]], str]
    pov_camera_phrase: Callable[[list[str]], str]
    expression_phrase: Callable[[Any], str]
    composition_phrase: Callable[..., str]
    paragraph: Callable[[list[str]], str]


def _coworking_action_anchor(action_family: str, scene_text: str, action: str) -> str:
    action_lower = action.lower()
    if "office chair seat and chair arms" in action_lower:
        return ""
    scene_lower = scene_text.lower()
    if not any(term in scene_lower for term in ("coworking", "office", "desk", "laptop", "glass partition")):
        return ""
    if action_family == "climax" and "post-ejaculation open-thigh display" in action_lower:
        return (
            "office chair seat and chair arms frame the lower foreground around her open thighs, "
            "with desk edges, laptop tables, glass partitions, plants, and tall-window depth beside and behind her body"
        )
    if "broad v-frame" in action_lower and "open-thigh frame" in action_lower:
        return (
            "office chair seat and chair arms frame the lower foreground around her hips and raised knees, "
            "with desk edges, laptop tables, glass partitions, plants, and tall-window depth beside and behind her body"
        )
    if action_family != "manual":
        return ""
    return (
        "office chair seat and chair arms frame the lower foreground around her hips, "
        "with desk edges, laptop tables, glass partitions, plants, and tall-window depth beside and behind her body"
    )


def _list_values(value: Any) -> list[str]:
    if isinstance(value, list):
        return [str(item) for item in value if str(item).strip()]
    if isinstance(value, str) and value.strip():
        return [part.strip() for part in value.split(",") if part.strip()]
    return []


def _krea2_variant_keys(row: dict[str, Any]) -> list[str]:
    config = row.get("hardcore_position_config") if isinstance(row.get("hardcore_position_config"), dict) else {}
    axis_values = row.get("item_axis_values") if isinstance(row.get("item_axis_values"), dict) else {}
    return list(dict.fromkeys([*_list_values(config.get("krea2_variant_keys")), *_list_values(axis_values.get("krea2_variant_keys"))]))


def _has_krea2_variant(row: dict[str, Any], key: str) -> bool:
    return key in _krea2_variant_keys(row)


def _has_krea2_oral_contact_variant(row: dict[str, Any]) -> bool:
    return any(key in ORAL_CONTACT_VARIANTS for key in _krea2_variant_keys(row))


def _has_krea2_atlas_variant(row: dict[str, Any]) -> bool:
    return any(krea2_pose_variant_catalog.get_variant(key) for key in _krea2_variant_keys(row))


def _has_krea2_top_down_variant(row: dict[str, Any]) -> bool:
    for key in _krea2_variant_keys(row):
        variant = krea2_pose_variant_catalog.get_variant(key)
        geometry = " ".join(
            [str(variant.get("canonical_geometry") or ""), *[str(cue) for cue in variant.get("prompt_cues") or []]]
        ).lower()
        if any(term in geometry for term in ("top-down", "top view", "top-view", "nadir", "overhead")):
            return True
    return False


def _filter_expression_for_krea2_variant(row: dict[str, Any], expression: Any) -> Any:
    if not _has_krea2_oral_contact_variant(row):
        return expression
    clauses = [clause.strip() for clause in str(expression or "").split(";") if clause.strip()]
    if not clauses:
        return expression
    kept = [
        clause
        for clause in clauses
        if not any(term in clause.lower() for term in MOUTH_EXPRESSION_TERMS)
    ]
    return "; ".join(kept)


def _filter_camera_scene_for_krea2_variant(row: dict[str, Any], camera_scene: Any) -> str:
    text = str(camera_scene or "")
    if _has_krea2_atlas_variant(row):
        return ""
    if (_has_krea2_oral_contact_variant(row) or _has_krea2_top_down_variant(row)) and "eye-level" in text.lower():
        return ""
    return text


def format_configured_cast_result(
    request: KreaConfiguredCastRequest,
    deps: KreaConfiguredCastDependencies,
) -> KreaConfiguredCastPrompt:
    row = request.row
    subject = deps.clean(row.get("subject_phrase") or request.primary or "adult sexual scene")
    cast = deps.clean(row.get("cast_summary"))
    try:
        women_count = int(row.get("women_count") or 0)
        men_count = int(row.get("men_count") or 0)
    except (TypeError, ValueError):
        women_count = men_count = 0
    cast_descriptor_text = deps.clean(row.get("cast_descriptor_text"))
    pov_labels = deps.pov_labels_from_value(row.get("pov_character_labels"))
    camera = request.camera
    if pov_labels:
        camera = ""
    cast_prose, cast_labels = deps.cast_prose_omit(cast_descriptor_text, pov_labels)
    if not cast_labels and women_count == 1 and men_count == 1:
        cast_labels = ["Woman A", "Man A"]
    cast_labels = deps.merge_labels(cast_labels, pov_labels)
    has_atlas_variant = _has_krea2_atlas_variant(row)
    expression = "" if has_atlas_variant else _filter_expression_for_krea2_variant(row, request.expression)
    expression = deps.filter_pov_labeled_clauses(expression, pov_labels)
    expression = deps.natural_label_text(expression, cast_labels)
    composition = deps.sanitize_hardcore_environment_anchors(request.composition)
    source_composition = deps.sanitize_hardcore_environment_anchors(request.source_composition)
    role_graph = deps.sanitize_scene_text_for_cast(
        deps.sanitize_hardcore_environment_anchors(row.get("source_role_graph") or row.get("role_graph")),
        cast_labels,
    )
    item = deps.sanitize_scene_text_for_cast(
        deps.sanitize_hardcore_environment_anchors(request.item),
        cast_labels,
    )
    role_graph = deps.natural_label_text(role_graph, cast_labels)
    item = deps.natural_label_text(item, cast_labels)
    axis_values = deps.sanitize_hardcore_axis_values(row.get("item_axis_values"))
    detail_density = deps.normalize_hardcore_detail_density(row.get("hardcore_detail_density"))
    action_family = deps.row_action_family(row)
    action = deps.hardcore_action_sentence(
        role_graph,
        item,
        source_composition,
        axis_values,
        detail_density,
        action_family,
    )
    action = deps.pov_action_phrase(
        action,
        pov_labels,
        role_graph,
        item,
        source_composition,
        axis_values,
        detail_density,
    )
    scene_anchor = _coworking_action_anchor(
        action_family,
        " ".join(part for part in (request.scene, request.camera_scene, composition, source_composition) if part),
        action,
    )
    camera_scene = _filter_camera_scene_for_krea2_variant(row, request.camera_scene)
    output_composition = "" if has_atlas_variant else deps.pov_composition_text(composition, pov_labels)
    parts = [
        action,
        scene_anchor,
        deps.pov_camera_phrase(pov_labels),
        cast_prose,
        f"A consensual explicit adult scene with {subject}" if not action else "",
        f"The cast includes {cast}" if cast and not cast_prose and not (women_count == 1 and men_count == 1) else "",
        f"The setting is {request.scene}" if request.scene else "",
        camera_scene,
        deps.expression_phrase(expression),
        deps.composition_phrase(output_composition, action, "The image is framed as", detail_density),
        camera,
        "" if has_atlas_variant else request.style if request.detail_level != "concise" else "",
    ]
    return KreaConfiguredCastPrompt(deps.paragraph(parts))


def format_configured_cast(
    request: KreaConfiguredCastRequest,
    deps: KreaConfiguredCastDependencies,
) -> tuple[str, str]:
    return format_configured_cast_result(request, deps).as_tuple()
