scripts/two_user_flow_tree.py

#!/usr/bin/env python3
"""Bounded exhaustive two-user flow tree simulator.

The action space is intentionally finite: a small catalog, two users, one scenario,
one pair relationship, and every meaningful mutation available from each reached
state inside a max-depth bound. Each edge records an expected result and the real
result from the backend, then validates invariants for the post-state.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sqlite3
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

from sqlmodel import Session

from backend.core import Backend
from models import Kink, SimilarityEdge

KINK_ID = "oral"
OTHER_KINK_ID = "impact"
SCENARIO_ID = "oral_scene"
ROLE_ID = "giver"

PLAY_CATEGORIES: tuple[tuple[str, str, list[str]], ...] = (
    ("love_together", "love", ["together"]),
    ("like_to_me", "like", ["to_me"]),
    ("like_by_me", "like", ["by_me"]),
    ("curious_both_directed", "curious", ["to_me", "by_me"]),
    ("hidden", "not_interested", []),
    ("no_go", "hard_no", []),
)
ADD_PLAY_CATEGORIES: tuple[tuple[str, str, list[str]], ...] = (
    ("like_together", "like", ["together"]),
    ("no_go", "hard_no", []),
)
SCENARIO_CATEGORIES: tuple[tuple[str, str, list[str]], ...] = (
    ("love_together", "love", ["together"]),
    ("like_to_me", "like", ["to_me"]),
    ("hidden", "not_interested", []),
    ("no_go", "hard_no", []),
)
ADD_SCENARIO_CATEGORIES: tuple[tuple[str, str, list[str]], ...] = (
    ("like_together", "like", ["together"]),
)


@dataclass(frozen=True)
class Actor:
    label: str
    user_id: str


@dataclass
class Action:
    label: str
    expected_ok: bool
    apply: Callable[[Backend], tuple[bool, str]]


@dataclass
class Node:
    id: int
    db_path: Path
    depth: int
    path: list[str]
    state_hash: str


def _serialize(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"))


def _state_hash(state: dict[str, Any]) -> str:
    return hashlib.sha256(_serialize(state).encode()).hexdigest()[:16]


def seed_backend(path: Path, *, seed_plays_per_user: int = 0) -> dict[str, str]:
    os.environ.setdefault("KINK_SKIP_HEAVY_WARM", "1")
    b = Backend(path)
    with Session(b.engine) as session:
        for kid, name in [
            (KINK_ID, "Oral Sex"),
            (OTHER_KINK_ID, "Spanking"),
            (SCENARIO_ID, "Slow Oral Sex Scene"),
            (ROLE_ID, "Giver"),
        ]:
            session.add(
                Kink(
                    id=kid,
                    name=name,
                    cluster="fetlife_fetish",
                    short_definition=f"{name} fixture",
                    notes="Kinksters: 1000",
                )
            )
        bulk_count = max(0, seed_plays_per_user * 2)
        for idx in range(bulk_count):
            kid = f"bulk_{idx:03d}"
            session.add(
                Kink(
                    id=kid,
                    name=f"Bulk Flow Kink {idx:03d}",
                    cluster="fetlife_fetish",
                    short_definition=f"Bulk flow fixture {idx:03d}",
                    notes="Kinksters: 1000",
                )
            )
        session.add(
            SimilarityEdge(
                id="sim_oral_impact",
                left_kink_id=KINK_ID,
                right_kink_id=OTHER_KINK_ID,
                similarity_type="catalog",
                score=0.8,
                method="two_user_tree",
            )
        )
        session.commit()
    b._invalidate_catalog_cache(force=True)
    b.upsert_scenario_parent(SCENARIO_ID, KINK_ID, 0.95, evidence={"source": "two_user_tree"})
    left = b.create_user()
    right = b.create_user()
    ids = {"A": left.id, "B": right.id}
    if seed_plays_per_user > 0:
        overlap_count = seed_plays_per_user // 2
        for actor_label, user in (("A", left), ("B", right)):
            start = 0 if actor_label == "A" else overlap_count
            for offset in range(seed_plays_per_user):
                idx = start + offset
                rating = ("love", "like", "curious")[idx % 3]
                directions = (["together"], ["to_me"], ["by_me"], ["to_me", "by_me"])[idx % 4]
                b.save_play_preference(user.id, f"bulk_{idx:03d}", rating, list(directions))
    b.engine.dispose()
    return ids


def clone_db(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if dst.exists():
        dst.unlink()
    with sqlite3.connect(src) as source:
        source.execute("PRAGMA wal_checkpoint(FULL)")
        with sqlite3.connect(dst) as target:
            source.backup(target)


def actors(ids: dict[str, str]) -> tuple[Actor, Actor]:
    return Actor("A", ids["A"]), Actor("B", ids["B"])


def other_actor(actor: Actor, ids: dict[str, str]) -> Actor:
    a, b = actors(ids)
    return b if actor.label == "A" else a


def user_state(b: Backend, actor: Actor) -> dict[str, Any]:
    user = b.get_user(actor.user_id)
    assert user is not None
    return user


def linked(user: dict[str, Any], other: Actor) -> bool:
    return other.user_id in user.get("partners", [])


def pair_group_ids(b: Backend, ids: dict[str, str]) -> list[str]:
    a, b_actor = actors(ids)
    groups = b.list_partner_groups(a.user_id)
    return [
        g["id"]
        for g in groups
        if g["name"] == "Together" and set(g.get("participant_ids", [])) == {a.user_id, b_actor.user_id}
    ]


def visible_group_ids(b: Backend, actor: Actor) -> list[str]:
    return [g["id"] for g in b.list_partner_groups(actor.user_id)]


def state_snapshot(b: Backend, ids: dict[str, str]) -> dict[str, Any]:
    out: dict[str, Any] = {"users": {}, "groups": {}}
    for actor in actors(ids):
        user = user_state(b, actor)
        out["users"][actor.label] = {
            "plays": user.get("plays", {}),
            "scenario_preferences": user.get("scenario_preferences", {}),
            "roles": user.get("roles", []),
            "preferences": user.get("preferences", {}),
            "partners": [("A" if pid == ids["A"] else "B" if pid == ids["B"] else pid) for pid in user.get("partners", [])],
            "incoming": [("A" if pid == ids["A"] else "B" if pid == ids["B"] else pid) for pid in user.get("incoming_partner_requests", [])],
            "outgoing": [("A" if pid == ids["A"] else "B" if pid == ids["B"] else pid) for pid in user.get("outgoing_partner_requests", [])],
        }
        out["groups"][actor.label] = [
            {
                "id": g["id"],
                "name": g["name"],
                "owner": "A" if g["owner_user_id"] == ids["A"] else "B" if g["owner_user_id"] == ids["B"] else g["owner_user_id"],
                "members": sorted("A" if uid == ids["A"] else "B" if uid == ids["B"] else uid for uid in g.get("member_ids", [])),
                "participants": sorted("A" if uid == ids["A"] else "B" if uid == ids["B"] else uid for uid in g.get("participant_ids", [])),
                "my_share": bool(g.get("my_share_full_list")),
                "sharing": sorted("A" if uid == ids["A"] else "B" if uid == ids["B"] else uid for uid in g.get("sharing_member_ids", [])),
            }
            for g in b.list_partner_groups(actor.user_id)
        ]
    return out


def expected_direct_matches(b: Backend, ids: dict[str, str]) -> set[str]:
    a, b_actor = actors(ids)
    ua = user_state(b, a)
    ub = user_state(b, b_actor)
    matches: set[str] = set()
    for kink_id, left in ua.get("plays", {}).items():
        right = ub.get("plays", {}).get(kink_id)
        kink = b.get_kink(kink_id)
        if not right or not kink:
            continue
        if left["interest_state"] not in {"love", "like", "curious"}:
            continue
        if right["interest_state"] not in {"love", "like", "curious"}:
            continue
        if b._directions_compatible(left.get("directions", []), right.get("directions", []), kink):
            matches.add(kink_id)
    return matches


def validate_invariants(b: Backend, ids: dict[str, str], *, check_recommendations: bool = False) -> list[str]:
    errors: list[str] = []
    a, b_actor = actors(ids)
    for actor in (a, b_actor):
        user = user_state(b, actor)
        board = b.play_board(actor.user_id)
        for kink_id, state in user.get("plays", {}).items():
            buckets = {name for name, items in board.items() if any(item["id"] == kink_id for item in items)}
            expected: set[str]
            if state["interest_state"] == "hard_no":
                expected = {"no_go"}
            elif state["interest_state"] == "not_interested":
                expected = {"hidden"}
            else:
                expected = set(state.get("directions") or ["together"])
            if buckets != expected:
                errors.append(f"{actor.label}:{kink_id}:board={sorted(buckets)} expected={sorted(expected)}")

        for group in b.list_partner_groups(actor.user_id):
            participants = set(group.get("participant_ids", []))
            if actor.user_id not in participants:
                errors.append(f"{actor.label}:{group['id']}:visible-but-not-participant")
            if not b.can_access_partner_group(actor.user_id, group["id"]):
                errors.append(f"{actor.label}:{group['id']}:visible-but-inaccessible")
            if group["owner_user_id"] in group.get("member_ids", []):
                errors.append(f"{actor.label}:{group['id']}:owner-counted-as-member")
            if group["member_count"] != len(group.get("member_ids", [])):
                errors.append(f"{actor.label}:{group['id']}:bad-member-count")

    if linked(user_state(b, a), b_actor) and linked(user_state(b, b_actor), a):
        pair_groups = pair_group_ids(b, ids)
        if len(pair_groups) != 1:
            errors.append(f"pair-group-count={len(pair_groups)}")
        direct = {item["id"] for item in b.shared_play_list(a.user_id, b_actor.user_id)["direct_matches"]}
        expected = expected_direct_matches(b, ids)
        if direct != expected:
            errors.append(f"direct-overlap={sorted(direct)} expected={sorted(expected)}")
        for gid in set(visible_group_ids(b, a)) & set(visible_group_ids(b, b_actor)):
            left = b.shared_play_list_for_group(a.user_id, gid)
            right = b.shared_play_list_for_group(b_actor.user_id, gid)
            if left != right:
                errors.append(f"{gid}:group-overlap-differs-by-viewer")

    for actor in (a, b_actor):
        for group in b.list_partner_groups(actor.user_id):
            gid = group["id"]
            for viewer in (a, b_actor):
                for target in (a, b_actor):
                    board = b.partner_full_board(viewer.user_id, gid, target.user_id)
                    target_view = next((g for g in b.list_partner_groups(viewer.user_id) if g["id"] == gid), None)
                    viewer_can_access = b.can_access_partner_group(viewer.user_id, gid)
                    target_shared = bool(target_view and target.user_id in target_view.get("sharing_member_ids", []))
                    if viewer.user_id == target.user_id and viewer_can_access:
                        target_shared = bool(target_view and target_view.get("my_share_full_list"))
                    if target_shared and board is None:
                        errors.append(f"{viewer.label}->{target.label}:{gid}:shared-board-missing")
                    if not target_shared and board is not None and viewer.user_id != target.user_id:
                        errors.append(f"{viewer.label}->{target.label}:{gid}:unshared-board-visible")

    if check_recommendations:
        for actor in (a, b_actor):
            user = user_state(b, actor)
            play_ids = set(user.get("plays", {}))
            rec_ids = {item["kink"]["id"] for item in b.recommend(actor.user_id, limit=8)}
            prompt_ids = {item["kink"]["id"] for item in b.prompt_candidates(actor.user_id, limit=8)}
            if rec_ids & play_ids:
                errors.append(f"{actor.label}:recs-include-rated={sorted(rec_ids & play_ids)}")
            if prompt_ids & play_ids:
                errors.append(f"{actor.label}:prompts-include-rated={sorted(prompt_ids & play_ids)}")
    return errors


def make_actions(b: Backend, ids: dict[str, str]) -> list[Action]:
    out: list[Action] = []
    for actor in actors(ids):
        other = other_actor(actor, ids)
        user = user_state(b, actor)
        other_user = user_state(b, other)
        is_linked = linked(user, other)
        has_outgoing = other.user_id in user.get("outgoing_partner_requests", [])
        has_incoming = other.user_id in user.get("incoming_partner_requests", [])
        def req(backend: Backend, a=actor, o=other) -> tuple[bool, str]:
            return backend.request_partner_link(a.user_id, o.user_id)

        out.append(Action(f"{actor.label}.request_link.{other.label}", not is_linked, req))

        def accept(backend: Backend, a=actor, o=other) -> tuple[bool, str]:
            return backend.accept_partner_request(a.user_id, o.user_id), "accepted"

        out.append(Action(f"{actor.label}.accept_request_from.{other.label}", has_incoming, accept))

        def decline(backend: Backend, a=actor, o=other) -> tuple[bool, str]:
            return backend.decline_partner_request(a.user_id, o.user_id), "declined"

        out.append(Action(f"{actor.label}.decline_request_from.{other.label}", has_incoming, decline))

        def cancel(backend: Backend, a=actor, o=other) -> tuple[bool, str]:
            return backend.cancel_partner_request(a.user_id, o.user_id), "cancelled"

        out.append(Action(f"{actor.label}.cancel_request_to.{other.label}", has_outgoing, cancel))

        play_categories = ADD_PLAY_CATEGORIES if KINK_ID not in user.get("plays", {}) else PLAY_CATEGORIES
        for cat, rating, directions in play_categories:
            current = user.get("plays", {}).get(KINK_ID)
            if current and current.get("interest_state") == rating and sorted(current.get("directions", [])) == sorted(directions):
                continue

            def save_play(backend: Backend, a=actor, r=rating, d=directions) -> tuple[bool, str]:
                backend.save_play_preference(a.user_id, KINK_ID, r, list(d))
                return True, "saved"

            verb = "add" if KINK_ID not in user.get("plays", {}) else "change"
            out.append(Action(f"{actor.label}.{verb}_play.{KINK_ID}.{cat}", True, save_play))

        if KINK_ID in user.get("plays", {}):
            def remove_play(backend: Backend, a=actor) -> tuple[bool, str]:
                backend.remove_play_preference(a.user_id, KINK_ID)
                return True, "removed"

            out.append(Action(f"{actor.label}.remove_play.{KINK_ID}", True, remove_play))

        scenario_categories = (
            ADD_SCENARIO_CATEGORIES
            if SCENARIO_ID not in user.get("scenario_preferences", {})
            else SCENARIO_CATEGORIES
        )
        for cat, rating, directions in scenario_categories:
            current = user.get("scenario_preferences", {}).get(SCENARIO_ID)
            if current and current.get("interest_state") == rating and sorted(current.get("directions", [])) == sorted(directions):
                continue

            def save_scenario(backend: Backend, a=actor, r=rating, d=directions) -> tuple[bool, str]:
                return backend.save_scenario_preference(a.user_id, KINK_ID, SCENARIO_ID, r, list(d)), "saved"

            verb = "add" if SCENARIO_ID not in user.get("scenario_preferences", {}) else "change"
            out.append(Action(f"{actor.label}.{verb}_scenario.{SCENARIO_ID}.{cat}", True, save_scenario))

        if SCENARIO_ID in user.get("scenario_preferences", {}):
            def remove_scenario(backend: Backend, a=actor) -> tuple[bool, str]:
                return backend.remove_scenario_preference(a.user_id, SCENARIO_ID), "removed"

            out.append(Action(f"{actor.label}.remove_scenario.{SCENARIO_ID}", True, remove_scenario))

        role_selected = ROLE_ID in user.get("roles", [])

        def save_role(backend: Backend, a=actor, selected=not role_selected) -> tuple[bool, str]:
            backend.save_role_preference(a.user_id, ROLE_ID, selected)
            return True, "saved"

        out.append(Action(f"{actor.label}.{'remove' if role_selected else 'add'}_role.{ROLE_ID}", True, save_role))

        show_images = bool(user.get("preferences", {}).get("show_images"))

        def save_pref(backend: Backend, a=actor, show=not show_images) -> tuple[bool, str]:
            return backend.save_preferences(a.user_id, show_images=show) is not None, "saved"

        out.append(Action(f"{actor.label}.set_show_images.{str(not show_images).lower()}", True, save_pref))

        if is_linked:
            def create_group(backend: Backend, a=actor, o=other) -> tuple[bool, str]:
                return backend.create_partner_group(a.user_id, "Custom", [o.user_id]) is not None, "created"

            out.append(Action(f"{actor.label}.create_group.Custom.with.{other.label}", True, create_group))

        for gid in visible_group_ids(b, actor):
            for share in (True, False):
                group = next(g for g in b.list_partner_groups(actor.user_id) if g["id"] == gid)
                if bool(group.get("my_share_full_list")) == share:
                    continue

                def toggle(backend: Backend, a=actor, g=gid, s=share) -> tuple[bool, str]:
                    return backend.toggle_share_full_list(a.user_id, g, s), "toggled"

                out.append(Action(f"{actor.label}.share.{gid}.{str(share).lower()}", True, toggle))

            group = next(g for g in b.list_partner_groups(actor.user_id) if g["id"] == gid)
            if group["owner_user_id"] == actor.user_id and other.user_id in group.get("member_ids", []):
                def remove_member(backend: Backend, a=actor, g=gid, o=other) -> tuple[bool, str]:
                    return backend.remove_partner_group_member(a.user_id, g, o.user_id), "removed"

                out.append(Action(f"{actor.label}.remove_member.{other.label}.from.{gid}", True, remove_member))
            if group["owner_user_id"] != actor.user_id and other.user_id in group.get("participant_ids", []):
                def invalid_remove(backend: Backend, a=actor, g=gid, o=other) -> tuple[bool, str]:
                    return backend.remove_partner_group_member(a.user_id, g, o.user_id), "removed"

                out.append(Action(f"{actor.label}.non_owner_remove_member.{other.label}.from.{gid}", False, invalid_remove))

    # Reverse pending request means request_link links immediately, not just pending.
    for action in out:
        if ".request_link." in action.label:
            actor_label = action.label.split(".")[0]
            actor = Actor(actor_label, ids[actor_label])
            other = other_actor(actor, ids)
            if actor.user_id in user_state(b, other).get("outgoing_partner_requests", []):
                action.expected_ok = True
    return out


def run_simulation(
    max_depth: int,
    workdir: Path | None = None,
    *,
    check_recommendations: bool = False,
    seed_plays_per_user: int = 0,
) -> dict[str, Any]:
    temp_ctx = None
    if workdir is None:
        temp_ctx = tempfile.TemporaryDirectory(prefix="kink_two_user_tree_")
        workdir = Path(temp_ctx.name)
    else:
        workdir.mkdir(parents=True, exist_ok=True)
    try:
        base_db = workdir / "node_0.db"
        ids = seed_backend(base_db, seed_plays_per_user=seed_plays_per_user)
        b0 = Backend(base_db)
        root_state = state_snapshot(b0, ids)
        root = Node(0, base_db, 0, [], _state_hash(root_state))
        b0.engine.dispose()
        queue: list[Node] = [root]
        seen: set[tuple[int, str]] = {(0, root.state_hash)}
        nodes = [{"id": root.id, "depth": 0, "parent": None, "action": None, "state_hash": root.state_hash, "path": []}]
        edges: list[dict[str, Any]] = []
        next_id = 1
        expanded = 0
        while queue:
            node = queue.pop(0)
            if node.depth >= max_depth:
                continue
            b = Backend(node.db_path)
            actions = make_actions(b, ids)
            b.engine.dispose()
            expanded += 1
            for action in actions:
                child_db = workdir / f"node_{next_id}.db"
                clone_db(node.db_path, child_db)
                child_backend = Backend(child_db)
                actual_ok, detail = action.apply(child_backend)
                invariant_errors = validate_invariants(
                    child_backend,
                    ids,
                    check_recommendations=check_recommendations,
                )
                child_state = state_snapshot(child_backend, ids)
                child_hash = _state_hash(child_state)
                child_path = [*node.path, action.label]
                edge = {
                    "from": node.id,
                    "to": next_id,
                    "depth": node.depth + 1,
                    "action": action.label,
                    "expected_ok": action.expected_ok,
                    "actual_ok": actual_ok,
                    "detail": detail,
                    "state_hash": child_hash,
                    "invariants_ok": not invariant_errors,
                    "invariant_errors": invariant_errors,
                    "path": child_path,
                }
                edges.append(edge)
                if actual_ok != action.expected_ok or invariant_errors:
                    child_backend.engine.dispose()
                    raise AssertionError(_serialize({"failure": edge, "state": child_state}))
                if (node.depth + 1, child_hash) not in seen:
                    seen.add((node.depth + 1, child_hash))
                    nodes.append(
                        {
                            "id": next_id,
                            "depth": node.depth + 1,
                            "parent": node.id,
                            "action": action.label,
                            "state_hash": child_hash,
                            "path": child_path,
                        }
                    )
                    queue.append(Node(next_id, child_db, node.depth + 1, child_path, child_hash))
                child_backend.engine.dispose()
                next_id += 1
        return {
            "max_depth": max_depth,
            "users": ids,
            "nodes": nodes,
            "edges": edges,
            "expanded_nodes": expanded,
            "unique_states": len(seen),
            "checked_edges": len(edges),
            "check_recommendations": check_recommendations,
            "seed_plays_per_user": seed_plays_per_user,
        }
    finally:
        if temp_ctx is not None:
            temp_ctx.cleanup()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-depth", type=int, default=2)
    parser.add_argument("--json-out", type=Path, default=None)
    parser.add_argument("--check-recommendations", action="store_true")
    parser.add_argument("--seed-plays-per-user", type=int, default=0)
    args = parser.parse_args()
    report = run_simulation(
        max_depth=args.max_depth,
        check_recommendations=args.check_recommendations,
        seed_plays_per_user=args.seed_plays_per_user,
    )
    summary = {
        "max_depth": report["max_depth"],
        "expanded_nodes": report["expanded_nodes"],
        "unique_states": report["unique_states"],
        "checked_edges": report["checked_edges"],
        "check_recommendations": report["check_recommendations"],
        "seed_plays_per_user": report["seed_plays_per_user"],
    }
    if args.json_out:
        args.json_out.parent.mkdir(parents=True, exist_ok=True)
        args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
        summary["json_out"] = str(args.json_out)
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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