#!/usr/bin/env python3
"""Generate deterministic JAX states for TypeScript engine parity tests."""

from __future__ import annotations

import json
from pathlib import Path

import jax.numpy as jnp
import numpy as np

from lost_cities_jax.engine import legal_action_mask, reset_from_order, step
from lost_cities_jax.obs import observation

OUTPUT = Path(__file__).resolve().parents[1] / "web" / "src" / "game" / "parity-fixture.json"


def state_json(state) -> dict:
    return {
        "deckOrder": np.asarray(state.deck_order).astype(int).tolist(),
        "drawPtr": int(state.draw_ptr),
        "cardLoc": np.asarray(state.card_loc).astype(int).tolist(),
        "handPublic": np.asarray(state.hand_public).astype(bool).tolist(),
        "colTop": np.asarray(state.col_top).astype(int).tolist(),
        "colHandshakes": np.asarray(state.col_hs).astype(int).tolist(),
        "colLength": np.asarray(state.col_len).astype(int).tolist(),
        "piles": [
            np.asarray(state.pile[color, : int(state.pile_len[color])]).astype(int).tolist()
            for color in range(5)
        ],
        "toMove": int(state.to_move),
        "stepCount": int(state.step_count),
        "done": bool(state.done),
    }


def main() -> None:
    rng = np.random.default_rng(20260713)
    order = rng.permutation(60).astype(np.int8)
    state = reset_from_order(jnp.asarray(order))
    rows = []
    for index in range(24):
        mask = np.asarray(legal_action_mask(state), dtype=bool)
        rows.append(
            {
                "index": index,
                "state": state_json(state),
                "legalMask": mask.tolist(),
                "observationP0": np.asarray(observation(state, jnp.int32(0))).tolist(),
                "observationP1": np.asarray(observation(state, jnp.int32(1))).tolist(),
            }
        )
        legal = np.flatnonzero(mask)
        action = int(legal[(index * 17 + 3) % len(legal)])
        rows[-1]["action"] = action
        state, _, _ = step(state, jnp.int32(action))
        if bool(state.done):
            break
    OUTPUT.write_text(json.dumps({"format": "jax-web-parity-v1", "rows": rows}) + "\n")
    print(f"wrote {OUTPUT} ({len(rows)} states)")


if __name__ == "__main__":
    main()