#!/usr/bin/env python3
"""Smoke-test core prompt routes without importing ComfyUI.

The checks here are intentionally lightweight invariants, not golden prompt
snapshots. They prove that representative rows still carry structured metadata
and that the Krea2, SDXL, and caption formatter paths consume metadata instead
of silently falling back to raw text parsing.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable


ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import caption_naturalizer  # noqa: E402
import krea_formatter  # noqa: E402
import prompt_builder as pb  # noqa: E402
import sdxl_formatter  # noqa: E402


Trigger = "sxcppnl7"
SdxlTrigger = "mythp0rt"


@dataclass
class SmokeReport:
    passed: list[str] = field(default_factory=list)
    failed: list[str] = field(default_factory=list)

    def ok(self, name: str) -> None:
        self.passed.append(name)
        print(f"PASS {name}")

    def fail(self, name: str, message: str) -> None:
        detail = f"{name}: {message}"
        self.failed.append(detail)
        print(f"FAIL {detail}")


def _clean_key(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", " ", str(value or "").lower()).strip()


def _json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=True, sort_keys=True)


def _expect(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def _expect_text(name: str, value: Any, min_len: int = 8) -> str:
    text = str(value or "").strip()
    _expect(len(text) >= min_len, f"{name} is empty or too short")
    _expect("None" not in text, f"{name} leaked None")
    _expect("  " not in text, f"{name} has repeated spaces")
    _expect(" ," not in text and " ." not in text, f"{name} has bad punctuation spacing")
    return text


def _expect_no_duplicate_comma_items(name: str, value: Any) -> None:
    items = [_clean_key(part) for part in str(value or "").split(",")]
    items = [part for part in items if part]
    duplicates = sorted({part for part in items if items.count(part) > 1})
    _expect(not duplicates, f"{name} has duplicate comma items: {duplicates[:5]}")


def _trigger_count(text: str, trigger: str) -> int:
    return len(re.findall(rf"(?<![a-z0-9_]){re.escape(trigger)}(?![a-z0-9_])", text, flags=re.IGNORECASE))


def _expect_trigger_once(name: str, value: Any, trigger: str) -> None:
    text = str(value or "")
    count = _trigger_count(text, trigger)
    _expect(count == 1, f"{name} should contain trigger {trigger!r} exactly once, got {count}")


def _expect_row_base(row: dict[str, Any], name: str) -> None:
    _expect(isinstance(row, dict), f"{name} did not return a metadata row")
    _expect_text(f"{name}.prompt", row.get("prompt"), 20)
    _expect_text(f"{name}.negative_prompt", row.get("negative_prompt"), 8)
    _expect_no_duplicate_comma_items(f"{name}.negative_prompt", row.get("negative_prompt"))
    _expect(json.loads(_json(row)) == row, f"{name} is not JSON-stable")


def _expect_custom_row(row: dict[str, Any], name: str) -> None:
    _expect_row_base(row, name)
    _expect(row.get("source") == "json_category", f"{name}.source should be json_category")
    _expect_text(f"{name}.item", row.get("item"), 8)
    _expect_text(f"{name}.scene_text", row.get("scene_text"), 8)
    _expect_text(f"{name}.composition", row.get("composition"), 8)
    _expect_text(f"{name}.role_graph", row.get("source_role_graph") or row.get("role_graph"), 8)
    _expect(isinstance(row.get("item_axis_values"), dict), f"{name}.item_axis_values missing")


def _expect_formatter_outputs(row: dict[str, Any], name: str, *, target: str = "auto") -> None:
    metadata = _json(row)
    krea = krea_formatter.format_krea2_prompt("", metadata_json=metadata, target=target)
    _expect("metadata" in krea.get("method", ""), f"{name}.krea did not use metadata: {krea.get('method')}")
    _expect_text(f"{name}.krea_prompt", krea.get("krea_prompt"), 20)
    _expect_no_duplicate_comma_items(f"{name}.krea_negative", krea.get("negative_prompt"))

    sdxl = sdxl_formatter.format_sdxl_prompt(
        "",
        metadata_json=metadata,
        target=target,
        trigger=SdxlTrigger,
        prepend_trigger=True,
    )
    _expect("metadata" in sdxl.get("method", ""), f"{name}.sdxl did not use metadata: {sdxl.get('method')}")
    _expect_text(f"{name}.sdxl_prompt", sdxl.get("sdxl_prompt"), 20)
    _expect_trigger_once(f"{name}.sdxl_prompt", sdxl.get("sdxl_prompt"), SdxlTrigger)
    _expect_no_duplicate_comma_items(f"{name}.sdxl_negative", sdxl.get("negative_prompt"))

    caption, method = caption_naturalizer.naturalize_caption(
        "",
        metadata_json=metadata,
        trigger=Trigger,
        include_trigger=True,
    )
    _expect("metadata" in method, f"{name}.caption did not use metadata: {method}")
    _expect_text(f"{name}.caption", caption, 20)
    _expect_trigger_once(f"{name}.caption", caption, Trigger)


def _character_cast(*, pov_man: bool = False) -> str:
    cast = pb.build_character_slot_json(
        subject_type="woman",
        label="A",
        age="25-year-old adult",
        ethnicity="western_european",
        figure="balanced",
        body="slim",
        descriptor_detail="full",
        expression_intensity=0.65,
        softcore_expression_intensity=0.45,
        hardcore_expression_intensity=0.85,
    )["character_cast"]
    return pb.build_character_slot_json(
        subject_type="man",
        label="A",
        age="40-year-old adult",
        ethnicity="western_european",
        figure="balanced",
        body="average",
        descriptor_detail="compact",
        expression_intensity=0.55,
        softcore_expression_intensity=0.35,
        hardcore_expression_intensity=0.75,
        presence_mode="pov" if pov_man else "visible",
        character_cast=cast,
    )["character_cast"]


def _action_filter(focus: str) -> str:
    kwargs = {
        "allow_toys": False,
        "allow_double": False,
        "allow_penetration": focus in ("penetration_only", "keep_pool"),
        "allow_foreplay": focus in ("foreplay_only", "keep_pool"),
        "allow_interaction": focus in ("interaction_only", "keep_pool"),
        "allow_manual": focus in ("manual_only", "keep_pool"),
        "allow_oral": focus in ("oral_only", "keep_pool"),
        "allow_outercourse": focus in ("outercourse_only", "keep_pool"),
        "allow_anal": focus in ("anal_only", "keep_pool"),
        "allow_climax": focus in ("climax_only", "keep_pool"),
    }
    return pb.build_hardcore_action_filter_json(focus=focus, **kwargs)


def _prompt_row(
    *,
    name: str,
    category: str,
    subcategory: str,
    seed: int,
    character_cast: str = "",
    women_count: int = 1,
    men_count: int = 1,
    hardcore_position_config: str = "",
) -> dict[str, Any]:
    row = pb.build_prompt(
        category=category,
        subcategory=subcategory,
        row_number=1,
        start_index=1,
        seed=seed,
        clothing="random",
        ethnicity="any",
        poses="random",
        backside_bias=0.35,
        figure="random",
        no_plus_women=False,
        no_black=False,
        minimal_clothing_ratio=0.5,
        standard_pose_ratio=0.5,
        trigger=Trigger,
        prepend_trigger_to_prompt=True,
        extra_positive="",
        extra_negative="",
        character_cast=character_cast,
        women_count=women_count,
        men_count=men_count,
        expression_enabled=True,
        expression_intensity=0.6,
        hardcore_position_config=hardcore_position_config,
    )
    _expect_row_base(row, name)
    return row


def smoke_builtin_single() -> None:
    row = _prompt_row(name="builtin_single_woman", category="woman", subcategory="random", seed=1001, men_count=0)
    _expect(row.get("source") == "built_in_generator", "builtin row should come from built-in generator")
    _expect_trigger_once("builtin_single_woman.prompt", row.get("prompt"), Trigger)
    _expect_formatter_outputs(row, "builtin_single_woman", target="single")


def smoke_hardcore_category_routes() -> None:
    cast = _character_cast()
    cases = [
        ("hardcore_penetration", "Penetrative sex", "penetration_only"),
        ("hardcore_oral", "Oral sex", "oral_only"),
        ("hardcore_manual", "Manual stimulation", "manual_only"),
        ("hardcore_outercourse", "Outercourse and genital teasing", "outercourse_only"),
        ("hardcore_foreplay", "Foreplay and teasing", "foreplay_only"),
        ("hardcore_aftercare", "Aftercare and cleanup", "interaction_only"),
    ]
    for index, (name, subcategory, focus) in enumerate(cases, start=1101):
        row = _prompt_row(
            name=name,
            category="Hardcore sexual poses",
            subcategory=subcategory,
            seed=index,
            character_cast=cast,
            women_count=1,
            men_count=1,
            hardcore_position_config=_action_filter(focus),
        )
        _expect_custom_row(row, name)
        _expect(row.get("subject_type") == "configured_cast", f"{name} should use configured cast")
        _expect_formatter_outputs(row, name, target="single")


def _insta_options(**overrides: Any) -> str:
    options = pb.build_insta_of_options_json(
        softcore_cast="same_as_hardcore",
        hardcore_cast="couple",
        hardcore_women_count=1,
        hardcore_men_count=1,
        softcore_level="lingerie_tease",
        hardcore_level="hardcore",
        platform_style="hybrid",
        continuity="same_creator_same_room",
        hardcore_clothing_continuity="explicit_nude",
        softcore_camera_mode="standard",
        hardcore_camera_mode="standard",
        camera_detail="compact",
        hardcore_detail_density="balanced",
    )
    data = json.loads(options)
    data.update(overrides)
    return _json(data)


def _expect_pair(pair: dict[str, Any], name: str) -> None:
    _expect(pair.get("mode") == "Insta/OF", f"{name}.mode should be Insta/OF")
    _expect_row_base(pair.get("softcore_row") or {}, f"{name}.softcore_row")
    _expect_custom_row(pair.get("hardcore_row") or {}, f"{name}.hardcore_row")
    _expect_text(f"{name}.softcore_prompt", pair.get("softcore_prompt"), 20)
    _expect_text(f"{name}.hardcore_prompt", pair.get("hardcore_prompt"), 20)
    _expect_no_duplicate_comma_items(f"{name}.softcore_negative", pair.get("softcore_negative_prompt"))
    _expect_no_duplicate_comma_items(f"{name}.hardcore_negative", pair.get("hardcore_negative_prompt"))
    _expect_formatter_outputs(pair, name, target="softcore")
    _expect_formatter_outputs(pair, f"{name}.hardcore", target="hardcore")


def smoke_insta_pair() -> None:
    pair = pb.build_insta_of_pair(
        row_number=1,
        start_index=1,
        seed=2101,
        ethnicity="any",
        figure="random",
        no_plus_women=False,
        no_black=False,
        trigger=Trigger,
        prepend_trigger_to_prompt=True,
        options_json=_insta_options(),
        character_cast=_character_cast(),
        hardcore_position_config=_action_filter("penetration_only"),
    )
    _expect_pair(pair, "insta_pair_same_cast")
    _expect(pair["softcore_row"].get("scene_text") == pair["hardcore_row"].get("scene_text"), "pair scene continuity broke")


def smoke_insta_pair_pov() -> None:
    pair = pb.build_insta_of_pair(
        row_number=1,
        start_index=1,
        seed=2201,
        ethnicity="any",
        figure="random",
        no_plus_women=False,
        no_black=False,
        trigger=Trigger,
        prepend_trigger_to_prompt=True,
        options_json=_insta_options(),
        character_cast=_character_cast(pov_man=True),
        hardcore_position_config=_action_filter("oral_only"),
    )
    _expect_pair(pair, "insta_pair_pov_man")
    pov_labels = pair.get("pov_character_labels") or []
    _expect("Man A" in pov_labels, "pair POV labels should include Man A")
    hard_row = pair.get("hardcore_row") or {}
    _expect("Man A" in (hard_row.get("pov_character_labels") or []), "hard row POV labels should include Man A")
    krea = krea_formatter.format_krea2_prompt("", metadata_json=_json(pair), target="hardcore")
    prompt = krea.get("krea_prompt") or ""
    _expect("viewer" in prompt.lower(), "POV Krea prompt should mention viewer perspective")


def smoke_no_expression_fallback() -> None:
    cast = pb.build_character_slot_json(
        subject_type="woman",
        label="A",
        age="25-year-old adult",
        ethnicity="western_european",
        body="slim",
        descriptor_detail="full",
        expression_enabled=False,
    )["character_cast"]
    row = _prompt_row(
        name="hardcore_expression_disabled",
        category="Hardcore sexual poses",
        subcategory="Penetrative sex",
        seed=2301,
        character_cast=cast,
        women_count=1,
        men_count=1,
        hardcore_position_config=_action_filter("penetration_only"),
    )
    _expect_custom_row(row, "hardcore_expression_disabled")
    _expect(not row.get("expression"), "expression should stay disabled without fallback")
    _expect("Facial expressions:" not in row.get("prompt", ""), "disabled expression leaked into prompt")
    _expect_formatter_outputs(row, "hardcore_expression_disabled", target="single")


SMOKE_CASES: list[tuple[str, Callable[[], None]]] = [
    ("builtin_single_woman", smoke_builtin_single),
    ("hardcore_category_routes", smoke_hardcore_category_routes),
    ("insta_pair_same_cast", smoke_insta_pair),
    ("insta_pair_pov_man", smoke_insta_pair_pov),
    ("expression_disabled", smoke_no_expression_fallback),
]


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--case",
        choices=[name for name, _func in SMOKE_CASES],
        action="append",
        help="Run only the named smoke case. Can be passed multiple times.",
    )
    args = parser.parse_args(argv)
    selected = set(args.case or [])
    report = SmokeReport()
    for name, func in SMOKE_CASES:
        if selected and name not in selected:
            continue
        try:
            func()
        except Exception as exc:  # noqa: BLE001 - report all smoke failures uniformly.
            report.fail(name, str(exc))
        else:
            report.ok(name)
    print(f"\nSummary: {len(report.passed)} passed, {len(report.failed)} failed")
    if report.failed:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())