from __future__ import annotations

import random
from dataclasses import dataclass
from itertools import combinations
from typing import Any

from .models import BehaviorCase, BehaviorSkeleton, Condition, HiddenRule, HiddenSystem


LENGTH_PRESETS = {"short": 4, "medium": 6, "long": 8}
ANSWER_OPTIONS = ("PRESENT", "ABSENT")


@dataclass(frozen=True)
class _State:
    present: bool = False
    invalidated: bool = False
    flags: frozenset[str] = frozenset()


def resolve_length(value: str | int) -> int:
    """Resolve the number of distinct events in the held-out query."""
    if isinstance(value, int):
        result = value
    elif value in LENGTH_PRESETS:
        result = LENGTH_PRESETS[value]
    else:
        try:
            result = int(value)
        except ValueError as exc:
            choices = ", ".join(LENGTH_PRESETS)
            raise ValueError(f"length must be an integer or one of: {choices}") from exc
    if not 4 <= result <= 10:
        raise ValueError("length must be between 4 and 10")
    return result


def _condition_holds(condition: Condition, attributes: dict[str, Any], state: _State) -> bool:
    if condition.source == "attribute":
        actual = attributes.get(condition.key)
    elif condition.source == "state":
        actual = condition.key in state.flags
    else:
        raise ValueError(f"unknown condition source: {condition.source}")

    operations = {
        "EQ": lambda left, right: left == right,
        "LE": lambda left, right: left <= right,
        "GT": lambda left, right: left > right,
        "IS_SET": lambda left, right: left is bool(right),
    }
    try:
        return operations[condition.operator](actual, condition.value)
    except KeyError as exc:
        raise ValueError(f"unknown condition operator: {condition.operator}") from exc


def _apply_action(state: _State, rule: HiddenRule) -> _State:
    if rule.action == "SET_FLAG":
        return _State(state.present, state.invalidated, state.flags | {rule.target})
    if rule.action == "ADD_RECORD":
        if state.invalidated:
            return state
        return _State(True, False, state.flags)
    if rule.action == "INVALIDATE_RECORD":
        return _State(False, True, state.flags)
    raise ValueError(f"unknown action: {rule.action}")


def exact_present_probability(
    system: HiddenSystem, attributes: dict[str, Any], events: tuple[str, ...]
) -> float:
    """Enumerate all stochastic branches while respecting event and state order."""
    initial = _State(present=system.initial_record == "PRESENT")
    distribution: dict[_State, float] = {initial: 1.0}
    for event in events:
        for rule in (item for item in system.rules if item.on_event == event):
            next_distribution: dict[_State, float] = {}
            for state, mass in distribution.items():
                if not all(_condition_holds(item, attributes, state) for item in rule.conditions):
                    next_distribution[state] = next_distribution.get(state, 0.0) + mass
                    continue
                inactive_mass = mass * (1.0 - rule.probability)
                active_state = _apply_action(state, rule)
                next_distribution[state] = next_distribution.get(state, 0.0) + inactive_mass
                next_distribution[active_state] = (
                    next_distribution.get(active_state, 0.0) + mass * rule.probability
                )
            distribution = next_distribution
    return sum(mass for state, mass in distribution.items() if state.present)


def _simulate_once(
    system: HiddenSystem,
    attributes: dict[str, Any],
    events: tuple[str, ...],
    rng: random.Random,
) -> bool:
    state = _State(present=system.initial_record == "PRESENT")
    for event in events:
        for rule in (item for item in system.rules if item.on_event == event):
            if all(_condition_holds(item, attributes, state) for item in rule.conditions):
                if rng.random() < rule.probability:
                    state = _apply_action(state, rule)
    return state.present


def _observe_case(
    case_id: str,
    system: HiddenSystem,
    attributes: dict[str, Any],
    events: tuple[str, ...],
    trials: int,
    rng: random.Random,
) -> BehaviorCase:
    present = sum(_simulate_once(system, attributes, events, rng) for _ in range(trials))
    return BehaviorCase(
        id=case_id,
        attributes=attributes,
        events=events,
        trials=trials,
        observed_present=present,
        observed_absent=trials - present,
    )


def _build_system(index: int, complexity: int, rng: random.Random) -> HiddenSystem:
    """Build a finite stochastic program with a gated rule and veto conflicts."""
    present_favored = index % 2 == 0
    add_probability = rng.choice((0.85, 0.9, 0.95)) if present_favored else rng.choice((0.65, 0.75))
    veto_probability = rng.choice((0.15, 0.25)) if present_favored else rng.choice((0.7, 0.8, 0.9))
    marker_probability = rng.choice((0.6, 0.75, 0.9))
    gated_action = "ADD_RECORD" if present_favored else "INVALIDATE_RECORD"
    gated_probability = rng.choice((0.6, 0.75, 0.85))

    rules = [
        HiddenRule("R0", "E0", (), add_probability, "ADD_RECORD", "record"),
        HiddenRule(
            "R1",
            "E1",
            (Condition("attribute", "load", "LE", 4),),
            veto_probability,
            "INVALIDATE_RECORD",
            "record",
        ),
        HiddenRule("R2", "E2", (), marker_probability, "SET_FLAG", "marked"),
        HiddenRule(
            "R3",
            "E3",
            (Condition("state", "marked", "IS_SET", True),),
            gated_probability,
            gated_action,
            "record",
        ),
    ]
    for number in range(4, complexity):
        action = "ADD_RECORD" if (number + index) % 3 else "INVALIDATE_RECORD"
        probability = rng.choice((0.2, 0.35, 0.5))
        condition = Condition("attribute", "class", "EQ", "priority")
        rules.append(
            HiddenRule(f"R{number}", f"E{number}", (condition,), probability, action, "record")
        )
    return HiddenSystem(
        initial_record="ABSENT",
        conflict_policy="PERSISTENT_INVALIDATION_VETO",
        rules=tuple(rules),
    )


def _support_event_sets(complexity: int) -> list[tuple[str, ...]]:
    events = [f"E{number}" for number in range(complexity)]
    cases: list[tuple[str, ...]] = [(event,) for event in events]
    cases.extend((("E2", "E3"), ("E3", "E2"), ("E0", "E1")))
    cases.extend(tuple(pair) for pair in combinations(events, 2) if tuple(pair) not in cases)
    return cases


def _generate_one(index: int, complexity: int, rng: random.Random) -> BehaviorSkeleton:
    system = _build_system(index, complexity, rng)
    low = {"load": 3, "class": "priority"}
    high = {"load": 8, "class": "standard"}
    query_attributes = dict(low)
    query_events = tuple(f"E{number}" for number in range(complexity))
    probability = exact_present_probability(system, query_attributes, query_events)
    if abs(probability - 0.5) < 0.08:
        raise RuntimeError("generated query has an insufficient probability margin")
    gold = "PRESENT" if probability > 0.5 else "ABSENT"

    trials = 48 if complexity <= 4 else 64
    support: list[BehaviorCase] = []
    for case_number, events in enumerate(_support_event_sets(complexity)):
        attributes = low if case_number % 3 else high
        support.append(
            _observe_case(
                f"history-{case_number:03d}", system, dict(attributes), events, trials, rng
            )
        )
    # Keeping events fixed while changing load makes the veto guard observable.
    next_id = len(support)
    support.append(
        _observe_case(f"history-{next_id:03d}", system, dict(low), ("E0", "E1"), trials, rng)
    )
    support.append(
        _observe_case(
            f"history-{next_id + 1:03d}", system, dict(high), ("E0", "E1"), trials, rng
        )
    )
    support.append(
        _observe_case(
            f"history-{next_id + 2:03d}",
            system,
            dict(low),
            ("E0", "E2", "E3"),
            trials,
            rng,
        )
    )

    query = BehaviorCase(
        id="query",
        attributes=query_attributes,
        events=query_events,
        trials=1,
        observed_present=None,
        observed_absent=None,
    )
    trace = tuple(
        [f"query events: {' -> '.join(query_events)}"]
        + [
            f"{rule.id}: on {rule.on_event}, p={rule.probability}, {rule.action}({rule.target})"
            for rule in system.rules
        ]
        + [f"exact P(PRESENT)={probability:.8f}", f"gold={gold}"]
    )
    return BehaviorSkeleton(
        id=f"dynamic-{index:04d}",
        family="stochastic_behavioral_system",
        complexity=complexity,
        event_vocabulary=tuple(f"E{number}" for number in range(complexity)),
        attribute_vocabulary={"load": "integer", "class": "category"},
        observation_protocol=(
            "Each trial begins with an absent record and no internal flags.",
            "Events occur from left to right; the same hidden system is used in every trial.",
            "PRESENT and ABSENT counts are final observations over independent repeated trials.",
            "For the query, choose which final outcome is more likely; an exact probability is not required.",
        ),
        hidden_system=system,
        support_cases=tuple(support),
        query_case=query,
        answer_options=ANSWER_OPTIONS,
        gold_answer=gold,
        gold_present_probability=probability,
        latent_trace=trace,
    )


def generate_skeletons(count: int, length: str | int, seed: int = 0) -> list[BehaviorSkeleton]:
    if count < 1:
        raise ValueError("count must be positive")
    complexity = resolve_length(length)
    rng = random.Random(seed)
    generated: list[BehaviorSkeleton] = []
    for index in range(count):
        for _attempt in range(100):
            try:
                generated.append(_generate_one(index, complexity, rng))
                break
            except RuntimeError as exc:
                if "insufficient probability margin" not in str(exc):
                    raise
        else:
            raise RuntimeError("could not generate a query with a stable probability margin")
    return generated