from __future__ import annotations

import random
from typing import Any

from data.schema import CLASSES, Fact


def add_rule(record: dict[str, Any], seed: int, held_out: bool = False) -> dict[str, Any]:
    rng = random.Random((seed * 9_999_991) + int(record["scene_id"].split("-")[-1]))
    target = rng.choice(record["objects"])
    color = target["color"]
    color_index = ("red", "blue", "green", "yellow").index(color)
    class_index = (color_index + (1 if held_out else 0)) % len(CLASSES)
    class_name = CLASSES[class_index]
    rule = Fact(color, "implies_class", class_name)
    result = dict(record)
    result.update(
        {
            "rules_text": f"Every {color} object belongs to class {class_name}.",
            "text_facts": [rule.to_dict()],
            "question": f"Which object belongs to class {class_name}?",
            "answer": target["object_id"],
            "combination": f"{color}->{class_name}",
        }
    )
    return result