"""Every notation of the same game must give the same tokens.

    pip install . chess && python tests/test_notations.py      # hand-written cases + 2,000 random games
python-chess writes each random game in 7 notations; the Rust tokenizer must map all of them to the tokens of
the true moves.
"""
import random, subprocess, sys
import chess
import chsm8_tok as T

FIG = {"K": "♔", "Q": "♕", "R": "♖", "B": "♗", "N": "♘"}

# One game with en passant (3. exf6), both castles, captures and an underpromotion with capture (8. gxh8=N).
CASES = {
    "SAN numbered": "1. e4 d5 2. e5 f5 3. exf6 Nc6 4. fxg7 Nf6 5. Nf3 Bf5 6. Be2 Qd7 7. O-O O-O-O 8. gxh8=N",
    "SAN unnumbered": "e4 d5 e5 f5 exf6 Nc6 fxg7 Nf6 Nf3 Bf5 Be2 Qd7 O-O O-O-O gxh8=N",
    "SAN 1.e4 style, 0-0, e.p., no '='": "1.e4 d5 2.e5 f5 3.exf6 e.p. Nc6 4.fxg7 Nf6 5.Nf3 Bf5 6.Be2 Qd7 "
                                         "7.0-0 0-0-0 8.gxh8N",
    "SAN black-move numbers, (N)": "1. e4 1... d5 2. e5 2... f5 3. exf6 3... Nc6 4. fxg7 4... Nf6 5. Nf3 "
                                   "5... Bf5 6. Be2 6... Qd7 7. O-O 7... O-O-O 8. gxh8(N)",
    "SAN /N promotion, lowercase castling": "e4 d5 e5 f5 exf6 Nc6 fxg7 Nf6 Nf3 Bf5 Be2 Qd7 o-o o-o-o gxh8/N",
    "LAN": "1. e2-e4 d7-d5 2. e4-e5 f7-f5 3. e5xf6 Nb8-c6 4. f6xg7 Ng8-f6 5. Ng1-f3 Bc8-f5 6. Bf1-e2 Qd8-d7 "
           "7. O-O O-O-O 8. g7xh8=N",
    "LAN, no separators, castling as king move": "e2e4 d7d5 e4e5 f7f5 e5xf6 Nb8c6 f6xg7 Ng8f6 Ng1f3 Bc8f5 "
           "Bf1e2 Qd8d7 Ke1-g1 Ke8-c8 g7xh8=N",
    "UCI": "e2e4 d7d5 e4e5 f7f5 e5f6 b8c6 f6g7 g8f6 g1f3 c8f5 f1e2 d8d7 e1g1 e8c8 g7h8n",
    "UCI, uppercase promo, castling as king-takes-rook": "e2e4 d7d5 e4e5 f7f5 e5f6 b8c6 f6g7 g8f6 g1f3 c8f5 "
           "f1e2 d8d7 e1h1 e8a8 g7h8N",
    "figurine SAN": "1. e4 d5 2. e5 f5 3. exf6 ♘c6 4. fxg7 ♘f6 5. ♘f3 ♗f5 6. ♗e2 ♕d7 7. O-O O-O-O 8. gxh8=♘",
    "PGN with headers, comments, variations, NAGs, glyphs": '[Event "test"]\n[White "A"]\n[Black "B"]\n\n'
           "1. e4 {king pawn} d5 $1 (1... e5 2. Nf3 (2. f4)) 2. e5 f5!? 3. exf6 Nc6 ; comment\n 4. fxg7 Nf6 "
           "5. Nf3 Bf5 6. Be2 Qd7 7. O-O O-O-O 8. gxh8=N+ 1-0",
    "mixed notations": "1. e4 d7-d5 2. e4e5 f5 3. e5xf6 b8c6 4. fxg7 Ng8-f6 5. g1f3 Bf5 6. Bf1-e2 d8d7 "
                       "7. 0-0 e8c8 8. g7h8n",
}

EP = {  # en passant in every notation
    "SAN": "1. e4 a6 2. e5 d5 3. exd6",
    "SAN e.p.": "1. e4 a6 2. e5 d5 3. exd6 e.p.",
    "LAN": "1. e2-e4 a7-a6 2. e4-e5 d7-d5 3. e5xd6",
    "UCI": "e2e4 a7a6 e4e5 d7d5 e5d6",
}


def check(group, rust=None):
    toks = {k: T.tokenize(v) for k, v in group.items()}
    ref = next(iter(toks.values()))
    for k, t in toks.items():
        assert t == ref, f"{k}: {t} != {ref}"
        if rust:
            r = run_rust(rust, group[k])
            assert r == ref, f"rust {k}: {r} != {ref}"
    return ref


def run_rust(binary, text):
    out = subprocess.run([binary], input=text, capture_output=True, text=True, check=True).stdout
    return [int(x) for x in out.split()]


def render(moves, style):
    """Write a move list in one notation with python-chess."""
    b, out = chess.Board(), []
    for i, m in enumerate(moves):
        num = f"{i // 2 + 1}. " if i % 2 == 0 else ""
        if style == "uci":
            s = m.uci()
        elif style == "lan":
            s = num + b.lan(m)
        elif style == "san":
            s = num + b.san(m)
        elif style == "san_bare":
            s = b.san(m).rstrip("+#")
        elif style == "lan_bare":
            s = b.lan(m).replace("-", "").rstrip("+#")
        elif style == "figurine":
            s = num + "".join(FIG.get(c, c) for c in b.san(m))
        elif style == "mixed":
            s = random.choice([m.uci(), b.lan(m), b.san(m)])
        out.append(s)
        b.push(m)
    return " ".join(out)


def random_game(rng, max_plies=200):
    b, moves = chess.Board(), []
    while not b.is_game_over() and len(moves) < max_plies:
        legal = list(b.legal_moves)
        promos = [m for m in legal if m.promotion]
        m = rng.choice(promos) if promos and rng.random() < .5 else rng.choice(legal)
        moves.append(m); b.push(m)
    return moves


if __name__ == "__main__":
    rust = sys.argv[sys.argv.index("--rust") + 1] if "--rust" in sys.argv else None
    ref = check(CASES, rust)
    assert T.token_to_uci(ref[-1]) == "g7h8n"
    check(EP, rust)
    print(f"hand-written: {len(CASES)} notations of one game + {len(EP)} en-passant forms agree")
    print("  UCI :", " ".join(T.to_uci(CASES["SAN numbered"])[:8]), "...")
    print("  tok :", ref[:8], "...")

    rng = random.Random(0)
    styles = ["uci", "lan", "san", "san_bare", "lan_bare", "figurine", "mixed"]
    n = plies = 0
    for g in range(2000):
        moves = random_game(rng)
        want = [m.from_square | (m.to_square << 6) | ({None: 0, 2: 1, 3: 2, 4: 3, 5: 4}[m.promotion] << 12) for m in moves]
        group = {s: render(moves, s) for s in styles}
        for s, text in group.items():
            assert T.tokenize(text) == want, (g, s)
            if rust and g < 300:
                assert run_rust(rust, text) == want, ("rust", g, s)
        n += 1; plies += len(moves)
    print(f"random games: {n} games, {plies} plies, {len(styles)} notations each: all identical tokens"
          + (" (Rust CLI checked on 300 games)" if rust else ""))