#!/usr/bin/env python3
"""Generate engine/quanpin/autocorrect_table.h from quanpin_utils.cpp.

The intact pinyin list is parsed directly from the engine source, so the
generated table can never drift from the runtime syllable list.

The table is split in four so each autocorrection type can be toggled
independently (quanpin.autocorrect_transposition / autocorrect_neighbor,
both default off):
  - transposition: single adjacent-letter swap, sahng -> shang
  - neighbor:      single QWERTY neighbor-key substitution, shabg -> shang
  - deletion:      single dropped letter, shng -> shang
  - insertion:     single inserted QWERTY-neighbor letter, sjhang -> shang

A wrong key may map to several syllables (e.g. ahan -> {shan, zhan}). Such
ambiguous variants are kept instead of dropped: the runtime disambiguates
with k-best cuts and word frequency, the probability-weighted approach of
google patent CN 101133411 B. A (wrong, correct) pair produced by two rules
stays in both tables so each per-type toggle remains effective on its own.

Variant length is capped at >= 3: every 2-letter string belongs to the
jianpin (abbreviated pinyin) space, where "wj" means w+j and must never be
rewritten into "wu" by the correction table.

Usage:
    python scripts/generate_quanpin_autocorrect.py          # regenerate the header
    python scripts/generate_quanpin_autocorrect.py --check  # verify the header is up to date (CI)
"""

from __future__ import annotations

import pathlib
import re
import sys

ENGINE_DIR = pathlib.Path(__file__).resolve().parents[2] / "engine"
UTILS_PATH = ENGINE_DIR / "quanpin" / "quanpin_utils.cpp"
OUT_PATH = ENGINE_DIR / "quanpin" / "autocorrect_table.h"

# QWERTY letter-key neighbors (finger-movement range, letter keys only).
QWERTY_NEIGHBORS = {
    "q": ["w", "a"],
    "w": ["q", "e", "s", "a"],
    "e": ["w", "r", "s", "d"],
    "r": ["e", "t", "d", "f"],
    "t": ["r", "y", "f", "g"],
    "y": ["t", "u", "g", "h"],
    "u": ["y", "i", "h", "j"],
    "i": ["u", "o", "j", "k"],
    "o": ["i", "p", "k", "l"],
    "p": ["o", "l"],
    "a": ["q", "w", "s", "z"],
    "s": ["a", "w", "d", "e", "x", "z"],
    "d": ["s", "e", "r", "f", "c", "x"],
    "f": ["d", "r", "t", "g", "v", "c"],
    "g": ["f", "t", "y", "h", "b", "v"],
    "h": ["g", "y", "u", "j", "n", "b"],
    "j": ["h", "u", "i", "k", "m", "n"],
    "k": ["j", "i", "o", "l", "m"],
    "l": ["k", "o", "p"],
    "z": ["a", "s", "x"],
    "x": ["z", "s", "d", "c"],
    "c": ["x", "d", "f", "v"],
    "v": ["c", "f", "g", "b"],
    "b": ["v", "g", "h", "n"],
    "m": ["n", "j", "k"],
    "n": ["b", "h", "j", "m"],
}

HEADER_TEMPLATE = """\
// Generated by scripts/generate_quanpin_autocorrect.py -- DO NOT EDIT.
// Quanpin typing autocorrection tables: wrong spelling -> correct syllable.
// Rules: single adjacent-letter transposition, one QWERTY neighbor-key
// substitution, one dropped letter, or one inserted letter (a repeat of an
// adjacent key or one of its QWERTY neighbors), applied to a legal syllable.
// Every key is a non-legal pinyin string of at least three letters (2-letter
// strings belong to the jianpin space). One key may map to several syllables:
// ambiguous variants are kept for runtime disambiguation via k-best cuts and
// word frequency (probability-weighted correction, cf. CN 101133411 B).
// Each table backs an independently toggleable correction type, and a pair
// produced by two rules stays in both tables.
//
// Correct syllables are stored once in kCorrectSyllables and referenced by
// 16-bit index from every Entry instead of repeating the same ~420 strings
// across 17k+ rows. The snapshot is re-extracted from intact_pinyin_list()
// in quanpin_utils.cpp on every generator run, so the indices can never
// drift from the runtime syllable list; --check fails on a stale header.
#pragma once

#include <cstddef>
#include <cstdint>
#include <string_view>

namespace quanpin::autocorrect
{{

// Every legal syllable, in intact_pinyin_list() order.
inline constexpr std::string_view kCorrectSyllables[] = {{
{syllable_entries}
}};

struct Entry
{{
    const char *wrong;
    std::uint16_t correct; // index into kCorrectSyllables
}};

// Adjacent-letter transpositions: sahng -> shang.
inline constexpr Entry kTranspositionEntries[] = {{
{transposition_entries}
}};

// QWERTY neighbor-key substitutions: shabg -> shang (n -> b).
inline constexpr Entry kNeighborEntries[] = {{
{neighbor_entries}
}};

// Dropped letters: shng -> shang (missing a).
inline constexpr Entry kDeletionEntries[] = {{
{deletion_entries}
}};

// Inserted letters: sjhang -> shang (j is a QWERTY neighbor of h).
inline constexpr Entry kInsertionEntries[] = {{
{insertion_entries}
}};

inline constexpr std::size_t kTranspositionCount = sizeof(kTranspositionEntries) / sizeof(kTranspositionEntries[0]);
inline constexpr std::size_t kNeighborCount = sizeof(kNeighborEntries) / sizeof(kNeighborEntries[0]);
inline constexpr std::size_t kDeletionCount = sizeof(kDeletionEntries) / sizeof(kDeletionEntries[0]);
inline constexpr std::size_t kInsertionCount = sizeof(kInsertionEntries) / sizeof(kInsertionEntries[0]);

}} // namespace quanpin::autocorrect
"""


def extract_intact_pinyin_list() -> list[str]:
    text = UTILS_PATH.read_text(encoding="utf-8")
    match = re.search(r"intact_pinyin_list\(\)[\s\S]*?kList = \{([\s\S]*?)\};", text)
    if not match:
        raise SystemExit("ERROR: cannot find intact_pinyin_list in " + str(UTILS_PATH))
    return re.findall(r'"([a-z]+)"', match.group(1))


def is_jqxy_v_alias(syllable: str) -> bool:
    # Mirrors quanpin::is_jqxy_v_alias: jv/jve/xvan/... only alias the u spelling.
    # Nobody types them on purpose, so typos are corrected to the canonical
    # syllable only; deriving keys from the alias would add noise like nvan -> jvan.
    return len(syllable) >= 2 and syllable[0] in "jqxy" and syllable[1] == "v"


def build_variants(
    syllables: list[str], make_variants, min_length: int
) -> dict[str, set[str]]:
    """Run one rule path over all syllables, keeping ambiguous variants.

    A wrong string that maps to several syllables keeps every mapping: the
    runtime disambiguates with k-best cuts and word frequency (probability-
    weighted correction, cf. CN 101133411 B) instead of dropping the key at
    generation time.
    """
    valid = set(syllables)
    mapping: dict[str, set[str]] = {}

    def add_variant(variant: str, correct: str) -> None:
        if (
            variant == correct
            or variant in valid
            or len(variant) < min_length
            or len(variant) > 6
            or "'" in variant
        ):
            return
        mapping.setdefault(variant, set()).add(correct)

    for syllable in syllables:
        if len(syllable) < 2 or is_jqxy_v_alias(syllable):
            continue
        for variant in make_variants(syllable):
            add_variant(variant, syllable)

    return mapping


def insertion_variants(syllable: str):
    # sjhang -> shang: insert one letter at any position. The inserted letter
    # must look like a real fat-finger: one of the QWERTY neighbors of an
    # adjacent key. Unconstrained insertion would generate ~60k mostly-
    # unrealistic keys, so position anchors are start (right neighbor only),
    # end (left neighbor only), and middle (either side). A double tap (repeat
    # of the adjacent key itself) is deliberately NOT generated: it is rare
    # enough that the widened false-positive surface is not worth it. The <=6
    # length cap means 6-letter syllables (zhuang) get no insertion coverage
    # (their variants would be 7 letters, beyond every legal syllable) --
    # deliberate, re-evaluate only if evaluation shows a real gap for 6-letter
    # syllables.
    for position in range(len(syllable) + 1):
        candidates: set[str] = set()
        if position > 0:
            left = syllable[position - 1]
            candidates.update(QWERTY_NEIGHBORS.get(left, ()))
        if position < len(syllable):
            right = syllable[position]
            candidates.update(QWERTY_NEIGHBORS.get(right, ()))
        for letter in candidates:
            yield syllable[:position] + letter + syllable[position:]


def build_entries(
    syllables: list[str],
) -> tuple[dict[str, set[str]], dict[str, set[str]], dict[str, set[str]], dict[str, set[str]]]:
    def transposition_variants(syllable: str):
        # sahng -> shang: swap each adjacent letter pair.
        return (
            syllable[:i] + syllable[i + 1] + syllable[i] + syllable[i + 2 :]
            for i in range(len(syllable) - 1)
        )

    def neighbor_variants(syllable: str):
        # shabg -> shang: replace each letter with a QWERTY neighbor.
        return (
            syllable[:i] + neighbor + syllable[i + 1 :]
            for i, ch in enumerate(syllable)
            for neighbor in QWERTY_NEIGHBORS.get(ch, ())
        )

    def deletion_variants(syllable: str):
        # shng -> shang: drop each single letter.
        return (syllable[:i] + syllable[i + 1 :] for i in range(len(syllable)))

    # 2-letter strings are jianpin space (wj/wh/wk/wl...), never correction keys.
    transposition = build_variants(syllables, transposition_variants, min_length=3)
    neighbor = build_variants(syllables, neighbor_variants, min_length=3)
    deletion = build_variants(syllables, deletion_variants, min_length=3)
    insertion = build_variants(syllables, insertion_variants, min_length=3)

    # Cross-type conflicts are intentionally kept in every table that produces
    # them: the runtime indexes each table per correction type with its own
    # edge weight, so a pair shared by two rules is what keeps the per-type
    # toggles independently effective.
    return transposition, neighbor, deletion, insertion


# engine/.clang-format pins ColumnLimit: 120; the generator must emit the
# clang-format-stable bytes itself because CI gates the header twice
# (generator --check and scripts/format.sh --check) on the same file.
COLUMN_LIMIT = 120


def render_grid(cells: list[str]) -> str:
    """Compact grid mirroring clang-format's column-aligned brace-list layout
    (verified byte-identical against the pinned formatter on all tables):

    - cell i lands in column i % columns;
    - each column is as wide as its own longest cell plus one separator
      space, except the trailing column which needs no separator;
    - the column count is the widest one whose full rows stay inside the
      column limit, scanning downward.

    Both the per-column widths and the column count must be computed, not
    fixed: clang-format tries the widest packing first, so the short index
    cells reflow to a different column count than the string cells they
    replace (7/6 columns where the old string cells gave 5).
    """
    if not cells:
        return ""
    for columns in range(min(len(cells), COLUMN_LIMIT - 4), 1, -1):
        slots = [max(len(cell) for cell in cells[col::columns]) + 1 for col in range(columns - 1)]
        slots.append(max(len(cell) for cell in cells[columns - 1 :: columns]))
        if 4 + sum(slots) > COLUMN_LIMIT:
            continue
        lines = []
        for i in range(0, len(cells), columns):
            chunk = cells[i : i + columns]
            lines.append("    " + "".join(cell.ljust(slots[j]) for j, cell in enumerate(chunk)).rstrip())
        return "\n".join(lines)
    # Single-column fallback: one cell per line, nothing to align.
    return "\n".join("    " + cell for cell in cells)


def render_entry_lines(entries: dict[str, set[str]], syllable_index: dict[str, int]) -> str:
    # Sort by the correct string, not by its index, so the emitted row order
    # -- the runtime's deterministic "table order" tie-break -- stays exactly
    # what the string-keyed tables produced before the index encoding.
    pairs = sorted(
        (wrong, correct) for wrong, corrects in entries.items() for correct in corrects
    )
    cells = []
    for wrong, correct in pairs:
        index = syllable_index.get(correct)
        if index is None:
            raise SystemExit(
                f"ERROR: correct syllable {correct!r} is missing from the extracted intact pinyin list"
            )
        cells.append(f'{{"{wrong}", {index}}},')
    return render_grid(cells)


def render_header(
    syllables: list[str],
    transposition: dict[str, set[str]],
    neighbor: dict[str, set[str]],
    deletion: dict[str, set[str]],
    insertion: dict[str, set[str]],
) -> str:
    syllable_index = {syllable: i for i, syllable in enumerate(syllables)}
    return HEADER_TEMPLATE.format(
        syllable_entries=render_grid([f'"{syllable}",' for syllable in syllables]),
        transposition_entries=render_entry_lines(transposition, syllable_index),
        neighbor_entries=render_entry_lines(neighbor, syllable_index),
        deletion_entries=render_entry_lines(deletion, syllable_index),
        insertion_entries=render_entry_lines(insertion, syllable_index),
    )


def main() -> int:
    syllables = extract_intact_pinyin_list()
    if not syllables:
        raise SystemExit("ERROR: extracted empty pinyin list")
    transposition, neighbor, deletion, insertion = build_entries(syllables)
    rendered = render_header(syllables, transposition, neighbor, deletion, insertion)

    if "--check" in sys.argv[1:]:
        if not OUT_PATH.exists():
            raise SystemExit(f"ERROR: {OUT_PATH} is missing")
        if OUT_PATH.read_text(encoding="utf-8") != rendered:
            raise SystemExit(f"ERROR: {OUT_PATH} is stale; rerun generate_quanpin_autocorrect.py")
        print(
            f"OK: {OUT_PATH} is up to date ({len(syllables)} syllables + {len(transposition)} transposition + "
            f"{len(neighbor)} neighbor + {len(deletion)} deletion + {len(insertion)} insertion keys)"
        )
        return 0

    # newline="\n" 钉住字节级输出：CI 的 --check 在 Linux 上写 LF，Windows 的
    # 默认换行转换会把文件改成 CRLF，两端产物必须一致。
    OUT_PATH.write_text(rendered, encoding="utf-8", newline="\n")
    print(
        f"Wrote {OUT_PATH} with {len(syllables)} syllables + {len(transposition)} transposition + "
        f"{len(neighbor)} neighbor + {len(deletion)} deletion + {len(insertion)} insertion keys"
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())