from __future__ import annotations

import random
from collections import deque
from dataclasses import dataclass, field
from typing import Any

from .models import Entity, HiddenProgram, Query, TransitionRule, WorldEvent, WorldSkeleton


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


@dataclass
class _World:
    locations: dict[str, str]
    order_level: dict[str, int]
    set_level: dict[str, int]
    above: dict[tuple[str, str], str] = field(default_factory=dict)
    subsets: dict[tuple[str, str], str] = field(default_factory=dict)
    members: dict[tuple[str, str], str] = field(default_factory=dict)
    disjoint: dict[frozenset[str], str] = field(default_factory=dict)
    connected: set[frozenset[str]] = field(default_factory=set)
    numbers: dict[str, int] = field(default_factory=dict)
    flags: dict[tuple[str, str], bool] = field(default_factory=dict)


def resolve_length(value: str | int) -> int:
    """Resolve required proof length, independently of event-log or text length."""
    if isinstance(value, int):
        result = value
    elif value in LENGTH_PRESETS:
        result = LENGTH_PRESETS[value]
    else:
        try:
            result = int(value)
        except ValueError as exc:
            raise ValueError(f"length must be an integer or one of: {', '.join(LENGTH_PRESETS)}") from exc
    if not 3 <= result <= 10:
        raise ValueError("length must be between 3 and 10")
    return result


def _program() -> HiddenProgram:
    return HiddenProgram(
        types=("Entity", "Set", "Region", "Number", "Boolean", "State"),
        relations=(
            "EQUAL",
            "NOT_EQUAL",
            "ABOVE",
            "BELOW",
            "MEMBER",
            "SUBSET",
            "DISJOINT",
            "ADJACENT",
            "CONNECTED",
            "CAUSES",
            "LOCATED_IN",
        ),
        operators=(
            "NOT",
            "AND",
            "OR",
            "IMPLIES",
            "ADD",
            "SUBTRACT",
            "COMPARE",
            "UNION",
            "INTERSECTION",
            "MOVE",
            "SWAP",
            "SET",
            "COMPOSE",
        ),
        quantifiers=("ALL", "EXISTS"),
        structures=("CHAIN", "BRANCH", "MERGE", "CYCLE", "NEST", "CONDITION", "SEQUENCE"),
        inference_rules=(
            "ABOVE(x,y) AND ABOVE(y,z) IMPLIES ABOVE(x,z)",
            "ABOVE(x,y) IMPLIES BELOW(y,x)",
            "MEMBER(x,A) AND SUBSET(A,B) IMPLIES MEMBER(x,B)",
            "MEMBER(x,A) AND DISJOINT(A,B) IMPLIES NOT MEMBER(x,B)",
            "CONNECTED is symmetric",
            "Numeric updates compose in temporal order",
        ),
        transition_rules=(
            TransitionRule(
                "move",
                "SELECT_ENTITY",
                ("target region differs from current region",),
                0.72,
                ("MOVE(entity, region)",),
            ),
            TransitionRule(
                "encounter-order",
                "COLOCATED_PAIR",
                ("entities share a region", "latent order levels are adjacent"),
                0.78,
                ("SET ABOVE(higher, lower)", "SET BELOW(lower, higher)"),
            ),
            TransitionRule(
                "membership",
                "ENTITY_SET_CONTACT",
                ("set is at base level",),
                0.68,
                ("SET MEMBER(entity, set)",),
            ),
            TransitionRule(
                "subset",
                "SET_PAIR_CONTACT",
                ("set levels are adjacent",),
                0.76,
                ("SET SUBSET(lower, higher)",),
            ),
            TransitionRule(
                "disjoint",
                "SET_PAIR_CONTACT",
                ("sets have equal level", "sets are distinct"),
                0.64,
                ("SET DISJOINT(left, right)",),
            ),
            TransitionRule(
                "numeric-update",
                "COUNTER_PULSE",
                (),
                0.82,
                ("ADD or SUBTRACT(counter, amount)",),
            ),
            TransitionRule(
                "connect",
                "ENTITY_PAIR_CONTACT",
                (),
                0.55,
                ("SET CONNECTED(left, right)",),
            ),
            TransitionRule(
                "state-change",
                "ENTITY_SIGNAL",
                (),
                0.58,
                ("SET Boolean state",),
            ),
        ),
    )


def _entities(entity_count: int) -> tuple[list[Entity], dict[str, list[str]]]:
    if entity_count < 48:
        raise ValueError("entity_count must be at least 48")
    region_count = max(4, entity_count // 32)
    set_count = max(24, entity_count // 5)
    counter_count = max(8, entity_count // 16)
    agent_count = entity_count - region_count - set_count - counter_count
    groups = {
        "Entity": [f"E{i:03d}" for i in range(agent_count)],
        "Set": [f"S{i:03d}" for i in range(set_count)],
        "Region": [f"R{i:03d}" for i in range(region_count)],
        "Number": [f"N{i:03d}" for i in range(counter_count)],
    }
    entities = [Entity(identifier, kind) for kind, values in groups.items() for identifier in values]
    return entities, groups


def _initial_world(
    groups: dict[str, list[str]], rng: random.Random, proof_length: int
) -> tuple[_World, dict[str, Any]]:
    agents, sets, regions, counters = (
        groups["Entity"],
        groups["Set"],
        groups["Region"],
        groups["Number"],
    )
    order_levels = max(proof_length + 2, 8)
    set_levels = max(proof_length + 1, 8)
    locations = {agent: rng.choice(regions) for agent in agents}
    world = _World(
        locations=locations,
        order_level={agent: index % order_levels for index, agent in enumerate(agents)},
        set_level={item: index % set_levels for index, item in enumerate(sets)},
        numbers={counter: rng.randint(10, 90) for counter in counters},
    )
    public = [
        {"id": f"I{index:04d}", "predicate": "LOCATED_IN", "arguments": [agent, region]}
        for index, (agent, region) in enumerate(locations.items())
    ]
    offset = len(public)
    public.extend(
        {
            "id": f"I{offset + index:04d}",
            "predicate": "INITIAL_VALUE",
            "arguments": [counter, value],
        }
        for index, (counter, value) in enumerate(world.numbers.items())
    )
    hidden = {
        "order_level": world.order_level,
        "set_level": world.set_level,
        "initial_locations": dict(locations),
        "initial_numbers": dict(world.numbers),
    }
    return world, {"public_observations": public, "hidden_state": hidden}


def _event(
    events: list[WorldEvent],
    time: int,
    kind: str,
    participants: tuple[str, ...],
    **arguments: Any,
) -> WorldEvent:
    item = WorldEvent(f"V{len(events):05d}", time, kind, participants, arguments)
    events.append(item)
    return item


def _simulate(
    world: _World,
    groups: dict[str, list[str]],
    steps: int,
    rng: random.Random,
) -> list[WorldEvent]:
    agents, sets, regions, counters = (
        groups["Entity"],
        groups["Set"],
        groups["Region"],
        groups["Number"],
    )
    events: list[WorldEvent] = []
    kinds = ("ENCOUNTER", "MOVE", "MEMBER", "SUBSET", "DISJOINT", "NUMBER", "CONNECT", "STATE")
    weights = (30, 12, 11, 15, 7, 11, 8, 6)
    probabilities = {
        "ENCOUNTER": 0.78,
        "MOVE": 0.72,
        "MEMBER": 0.68,
        "SUBSET": 0.76,
        "DISJOINT": 0.64,
        "NUMBER": 0.82,
        "CONNECT": 0.55,
        "STATE": 0.58,
    }

    for time in range(1, steps + 1):
        kind = rng.choices(kinds, weights=weights, k=1)[0]
        if rng.random() >= probabilities[kind]:
            continue

        if kind == "MOVE":
            agent = rng.choice(agents)
            choices = [region for region in regions if region != world.locations[agent]]
            destination = rng.choice(choices)
            origin = world.locations[agent]
            world.locations[agent] = destination
            _event(events, time, "MOVE", (agent, origin, destination), origin=origin, destination=destination)

        elif kind == "ENCOUNTER":
            lower_level = rng.randrange(max(world.order_level.values()))
            lower = [item for item in agents if world.order_level[item] == lower_level]
            higher = [item for item in agents if world.order_level[item] == lower_level + 1]
            if not lower or not higher:
                continue
            low, high = rng.choice(lower), rng.choice(higher)
            if world.locations[low] != world.locations[high] or (high, low) in world.above:
                continue
            item = _event(
                events,
                time,
                "RELATE",
                (high, low),
                predicate="ABOVE",
                subject=high,
                object=low,
                region=world.locations[high],
            )
            world.above[(high, low)] = item.id

        elif kind == "MEMBER":
            agent = rng.choice(agents)
            base_sets = [item for item in sets if world.set_level[item] == 0]
            container = rng.choice(base_sets)
            if (agent, container) in world.members:
                continue
            item = _event(
                events,
                time,
                "RELATE",
                (agent, container),
                predicate="MEMBER",
                subject=agent,
                object=container,
            )
            world.members[(agent, container)] = item.id

        elif kind == "SUBSET":
            lower_level = rng.randrange(max(world.set_level.values()))
            lower = [item for item in sets if world.set_level[item] == lower_level]
            higher = [item for item in sets if world.set_level[item] == lower_level + 1]
            child, parent = rng.choice(lower), rng.choice(higher)
            if (child, parent) in world.subsets or any(left == child for left, _right in world.subsets):
                continue
            item = _event(
                events,
                time,
                "RELATE",
                (child, parent),
                predicate="SUBSET",
                subject=child,
                object=parent,
            )
            world.subsets[(child, parent)] = item.id

        elif kind == "DISJOINT":
            level = rng.randrange(max(world.set_level.values()) + 1)
            same_level = [item for item in sets if world.set_level[item] == level]
            if len(same_level) < 2:
                continue
            left, right = rng.sample(same_level, 2)
            key = frozenset((left, right))
            if key in world.disjoint:
                continue
            item = _event(
                events,
                time,
                "RELATE",
                (left, right),
                predicate="DISJOINT",
                subject=left,
                object=right,
            )
            world.disjoint[key] = item.id

        elif kind == "NUMBER":
            counter = rng.choice(counters)
            delta = rng.choice((-7, -5, -3, 2, 4, 6))
            before = world.numbers[counter]
            after = max(0, before + delta)
            world.numbers[counter] = after
            _event(
                events,
                time,
                "UPDATE_NUMBER",
                (counter,),
                operation="ADD" if delta >= 0 else "SUBTRACT",
                amount=abs(delta),
                before=before,
                after=after,
            )

        elif kind == "CONNECT":
            left, right = rng.sample(agents, 2)
            key = frozenset((left, right))
            if key in world.connected:
                continue
            world.connected.add(key)
            _event(
                events,
                time,
                "RELATE",
                (left, right),
                predicate="CONNECTED",
                subject=left,
                object=right,
            )

        else:
            agent = rng.choice(agents)
            flag = rng.choice(("ACTIVE", "SEALED", "VERIFIED"))
            before = world.flags.get((agent, flag), False)
            after = not before
            world.flags[(agent, flag)] = after
            _event(
                events,
                time,
                "SET_STATE",
                (agent,),
                state=flag,
                before=before,
                after=after,
            )
    return events


def _shortest_path(
    edges: dict[tuple[str, str], str], minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
    adjacency: dict[str, list[tuple[str, str]]] = {}
    for (left, right), event_id in edges.items():
        adjacency.setdefault(left, []).append((right, event_id))
    candidates: list[tuple[str, str, tuple[str, ...], tuple[str, ...]]] = []
    for source in adjacency:
        queue = deque([(source, (source,), ())])
        seen = {source}
        while queue:
            node, nodes, event_ids = queue.popleft()
            if len(event_ids) >= minimum:
                candidates.append((source, node, nodes, event_ids))
                continue
            for target, event_id in adjacency.get(node, []):
                if target not in seen:
                    seen.add(target)
                    queue.append((target, nodes + (target,), event_ids + (event_id,)))
    if not candidates:
        return None
    candidates.sort(key=lambda item: (len(item[3]), item[0], item[1]))
    return candidates[0]


def _membership_path(
    world: _World, minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
    subset_path = _shortest_path(world.subsets, minimum - 1)
    if subset_path is None:
        return None
    candidates = []
    for (entity, base_set), member_event in world.members.items():
        path = _shortest_path_from(world.subsets, base_set, minimum - 1)
        if path:
            target, nodes, event_ids = path
            candidates.append((entity, target, (member_event,) + event_ids, (entity,) + nodes))
    if not candidates:
        return None
    candidates.sort(key=lambda item: (len(item[2]), item[0], item[1]))
    entity, target, event_ids, nodes = candidates[0]
    return entity, target, nodes, event_ids


def _membership_negative(
    world: _World, minimum: int
) -> tuple[str, str, tuple[str, ...], tuple[str, ...]] | None:
    adjacency: dict[str, list[tuple[str, str]]] = {}
    for (left, right), event_id in world.subsets.items():
        adjacency.setdefault(left, []).append((right, event_id))
    for (entity, base_set), member_event in sorted(world.members.items()):
        reachable = {base_set}
        queue = deque([(base_set, (base_set,), ())])
        paths: list[tuple[str, tuple[str, ...], tuple[str, ...]]] = []
        while queue:
            node, nodes, event_ids = queue.popleft()
            paths.append((node, nodes, event_ids))
            for target, event_id in adjacency.get(node, []):
                if target not in reachable:
                    reachable.add(target)
                    queue.append((target, nodes + (target,), event_ids + (event_id,)))
        for known_set, nodes, subset_events in paths:
            if 1 + len(subset_events) < minimum - 1:
                continue
            for pair, disjoint_event in world.disjoint.items():
                if known_set not in pair:
                    continue
                other = next(item for item in pair if item != known_set)
                if other in reachable:
                    continue
                proof_ids = (member_event,) + subset_events + (disjoint_event,)
                proof_nodes = (entity,) + nodes + (other,)
                return entity, other, proof_nodes, proof_ids
    return None


def _shortest_path_from(
    edges: dict[tuple[str, str], str], source: str, minimum: int
) -> tuple[str, tuple[str, ...], tuple[str, ...]] | None:
    adjacency: dict[str, list[tuple[str, str]]] = {}
    for (left, right), event_id in edges.items():
        adjacency.setdefault(left, []).append((right, event_id))
    queue = deque([(source, (source,), ())])
    seen = {source}
    while queue:
        node, nodes, event_ids = queue.popleft()
        if len(event_ids) >= minimum:
            return node, nodes, event_ids
        for target, event_id in adjacency.get(node, []):
            if target not in seen:
                seen.add(target)
                queue.append((target, nodes + (target,), event_ids + (event_id,)))
    return None


def _choose_query(
    index: int, complexity: int, world: _World, events: list[WorldEvent]
) -> tuple[Query, str, tuple[str, ...], tuple[str, ...], str] | None:
    if index % 2 == 0:
        path = _shortest_path(world.above, complexity)
        if path is None:
            return None
        source, target, nodes, proof_ids = path
        positive = index % 4 == 0
        arguments = (source, target) if positive else (target, source)
        gold = "YES" if positive else "NO"
        facts = tuple(
            f"ABOVE({nodes[position]},{nodes[position + 1]})"
            for position in range(len(nodes) - 1)
        )
        if not positive:
            facts += (f"ABOVE({source},{target}) therefore NOT ABOVE({target},{source})",)
        return Query("ABOVE", arguments, events[-1].time), gold, proof_ids, facts, "spatial_order"

    if index % 4 == 1:
        negative = _membership_negative(world, complexity)
        if negative is None:
            return None
        entity, target, nodes, proof_ids = negative
        known_set, disjoint_set = nodes[-2], nodes[-1]
        facts = [f"MEMBER({entity},{nodes[1]})"]
        facts.extend(
            f"SUBSET({nodes[position]},{nodes[position + 1]})"
            for position in range(1, len(nodes) - 2)
        )
        facts.extend(
            (
                f"DISJOINT({known_set},{disjoint_set})",
                f"therefore NOT MEMBER({entity},{target})",
            )
        )
        return (
            Query("MEMBER", (entity, target), events[-1].time),
            "NO",
            proof_ids,
            tuple(facts),
            "membership_disjointness",
        )

    membership = _membership_path(world, complexity)
    if membership is None:
        return None
    entity, target, nodes, proof_ids = membership
    facts = [f"MEMBER({entity},{nodes[1]})"]
    facts.extend(
        f"SUBSET({nodes[position]},{nodes[position + 1]})"
        for position in range(1, len(nodes) - 1)
    )
    return (
        Query("MEMBER", (entity, target), events[-1].time),
        "YES",
        proof_ids,
        tuple(facts),
        "membership_inheritance",
    )


def _final_state(world: _World) -> dict[str, Any]:
    return {
        "locations": dict(world.locations),
        "above": [[left, right] for left, right in sorted(world.above)],
        "subsets": [[left, right] for left, right in sorted(world.subsets)],
        "members": [[entity, container] for entity, container in sorted(world.members)],
        "disjoint": [sorted(pair) for pair in sorted(world.disjoint, key=lambda item: sorted(item))],
        "connected": [sorted(pair) for pair in sorted(world.connected, key=lambda item: sorted(item))],
        "numbers": dict(world.numbers),
        "flags": [
            {"entity": entity, "state": state, "value": value}
            for (entity, state), value in sorted(world.flags.items())
        ],
    }


def _generate_one(
    index: int,
    complexity: int,
    entity_count: int,
    simulation_steps: int,
    rng: random.Random,
) -> WorldSkeleton:
    entities, groups = _entities(entity_count)
    world, initial_state = _initial_world(groups, rng, complexity)
    events = _simulate(world, groups, simulation_steps, rng)
    if not events:
        raise RuntimeError("simulation produced no events")
    selected = _choose_query(index, complexity, world, events)
    if selected is None:
        raise RuntimeError("simulation did not produce a sufficiently deep proof")
    query, gold, proof_ids, proof_facts, query_family = selected
    direct = any(
        event.kind == "RELATE"
        and event.arguments.get("predicate") == query.predicate
        and (event.arguments.get("subject"), event.arguments.get("object")) == query.arguments
        for event in events
    )
    if direct:
        raise RuntimeError("query was directly stated in the event history")
    checks = {
        "direct_query_event_absent": True,
        "proof_length": len(proof_ids),
        "minimum_proof_length": complexity,
        "event_count": len(events),
        "proof_event_ratio": len(proof_ids) / len(events),
        "query_family": query_family,
        "weak_baselines": {
            "direct_fact": False,
            "proof_shorter_than_requested": len(proof_ids) < complexity,
            "proof_is_large_fraction_of_history": len(proof_ids) / len(events) >= 0.05,
        },
    }
    return WorldSkeleton(
        id=f"world-{index:04d}",
        family="dynamic_multi_logic_world",
        complexity=complexity,
        entity_count=entity_count,
        simulation_steps=simulation_steps,
        hidden_program=_program(),
        entities=tuple(entities),
        initial_state=initial_state,
        events=tuple(events),
        final_state=_final_state(world),
        query=query,
        answer_options=ANSWER_OPTIONS,
        gold_answer=gold,
        proof_event_ids=proof_ids,
        proof_facts=proof_facts,
        generator_checks=checks,
    )


def generate_skeletons(
    count: int,
    length: str | int,
    seed: int = 0,
    entity_count: int = 128,
    simulation_steps: int = 1000,
) -> list[WorldSkeleton]:
    if count < 1:
        raise ValueError("count must be positive")
    if simulation_steps < 100:
        raise ValueError("simulation_steps must be at least 100")
    complexity = resolve_length(length)
    rng = random.Random(seed)
    generated: list[WorldSkeleton] = []
    for index in range(count):
        for _attempt in range(60):
            try:
                generated.append(
                    _generate_one(index, complexity, entity_count, simulation_steps, rng)
                )
                break
            except RuntimeError:
                continue
        else:
            raise RuntimeError("could not generate a world with the requested proof complexity")
    return generated