from __future__ import annotations

import random
from typing import Any

from data.schema import CLASSES, Fact, Rule


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]
    text_fact = Fact(target["object_id"], "color", color)
    rule = Rule("var_x", "color", color, "class", class_name)
    result = dict(record)
    result.update(
        {
            "rules_text": (
                f"{target['object_id']} is {color}. "
                f"Every {color} object belongs to class {class_name}."
            ),
            "text_facts": [text_fact.to_dict()],
            "text_rules": [rule.to_dict()],
            "question": f"Which object belongs to class {class_name}?",
            "answer": target["object_id"],
            "combination": f"{color}->{class_name}",
        }
    )
    return result