#!/usr/bin/env python3
"""
Generate unicode_data.hpp and unicode_data.cpp from Python's unicodedata module.
Usage: python gen_unicode_data.py [output_dir]
       Default output_dir: ../../transformers/llm/engine/src/
"""

import sys
import os
import unicodedata

# Unicode General Category -> enum value mapping
CATEGORIES = [
    "Cn",  # 0  Not Assigned
    "Lu",  # 1  Uppercase Letter
    "Ll",  # 2  Lowercase Letter
    "Lt",  # 3  Titlecase Letter
    "Lm",  # 4  Modifier Letter
    "Lo",  # 5  Other Letter
    "Mn",  # 6  Nonspacing Mark
    "Mc",  # 7  Spacing Mark
    "Me",  # 8  Enclosing Mark
    "Nd",  # 9  Decimal Number
    "Nl",  # 10 Letter Number
    "No",  # 11 Other Number
    "Pc",  # 12 Connector Punctuation
    "Pd",  # 13 Dash Punctuation
    "Ps",  # 14 Open Punctuation
    "Pe",  # 15 Close Punctuation
    "Pi",  # 16 Initial Punctuation
    "Pf",  # 17 Final Punctuation
    "Po",  # 18 Other Punctuation
    "Sm",  # 19 Math Symbol
    "Sc",  # 20 Currency Symbol
    "Sk",  # 21 Modifier Symbol
    "So",  # 22 Other Symbol
    "Zs",  # 23 Space Separator
    "Zl",  # 24 Line Separator
    "Zp",  # 25 Paragraph Separator
    "Cc",  # 26 Control
    "Cf",  # 27 Format
    "Cs",  # 28 Surrogate
    "Co",  # 29 Private Use
]

CAT_TO_IDX = {c: i for i, c in enumerate(CATEGORIES)}

MAX_CP = 0x110000


def get_category_ranges():
    """Build run-length encoded category ranges."""
    ranges = []
    prev_cat = -1
    range_start = 0
    for cp in range(MAX_CP):
        try:
            cat = unicodedata.category(chr(cp))
        except (ValueError, OverflowError):
            cat = "Cn"
        idx = CAT_TO_IDX.get(cat, 0)
        if idx != prev_cat:
            if prev_cat >= 0:
                ranges.append((range_start, cp - 1, prev_cat))
            range_start = cp
            prev_cat = idx
    if prev_cat >= 0:
        ranges.append((range_start, MAX_CP - 1, prev_cat))
    return ranges


def get_tolower_map():
    """Build toLower mapping (only entries where lower != original)."""
    mapping = []
    for cp in range(MAX_CP):
        try:
            ch = chr(cp)
            lc = ch.lower()
            if len(lc) == 1:
                lcp = ord(lc)
                if lcp != cp:
                    mapping.append((cp, lcp))
        except (ValueError, OverflowError):
            pass
    return mapping



def write_hpp(out_dir):
    path = os.path.join(out_dir, "unicode_data.hpp")
    with open(path, 'w') as f:
        f.write("// Auto-generated by gen_unicode_data.py - DO NOT EDIT\n")
        f.write("#ifndef UNICODE_DATA_HPP\n")
        f.write("#define UNICODE_DATA_HPP\n\n")
        f.write("#include <cstdint>\n")
        f.write("#include <cstddef>\n\n")
        f.write("namespace MNN {\nnamespace Unicode {\n\n")

        f.write("enum Category : uint8_t {\n")
        for i, c in enumerate(CATEGORIES):
            f.write(f"    CAT_{c} = {i},\n")
        f.write(f"    CAT_COUNT = {len(CATEGORIES)}\n")
        f.write("};\n\n")

        f.write("struct CategoryRange {\n    uint32_t start;\n    uint32_t end;\n    uint8_t category;\n};\n\n")
        f.write("struct CaseMapping {\n    uint32_t from;\n    uint32_t to;\n};\n\n")

        f.write("extern const CategoryRange kCategoryRanges[];\n")
        f.write("extern const size_t kCategoryRangesSize;\n\n")
        f.write("extern const CaseMapping kToLower[];\n")
        f.write("extern const size_t kToLowerSize;\n\n")

        f.write("} // namespace Unicode\n} // namespace MNN\n\n")
        f.write("#endif // UNICODE_DATA_HPP\n")
    print(f"  wrote {path}")



def write_cpp(out_dir):
    print("  generating category ranges...")
    cat_ranges = get_category_ranges()
    print(f"    {len(cat_ranges)} ranges")

    print("  generating toLower map...")
    tolower = get_tolower_map()
    print(f"    {len(tolower)} entries")

    path = os.path.join(out_dir, "unicode_data.cpp")
    with open(path, 'w') as f:
        f.write("// Auto-generated by gen_unicode_data.py - DO NOT EDIT\n")
        f.write('#include "unicode_data.hpp"\n\n')
        f.write("namespace MNN {\nnamespace Unicode {\n\n")

        # Category ranges
        f.write(f"const CategoryRange kCategoryRanges[] = {{\n")
        for start, end, cat in cat_ranges:
            f.write(f"    {{0x{start:X}, 0x{end:X}, {cat}}},\n")
        f.write("};\n")
        f.write(f"const size_t kCategoryRangesSize = {len(cat_ranges)};\n\n")

        # toLower
        f.write(f"const CaseMapping kToLower[] = {{\n")
        for frm, to in tolower:
            f.write(f"    {{0x{frm:X}, 0x{to:X}}},\n")
        f.write("};\n")
        f.write(f"const size_t kToLowerSize = {len(tolower)};\n\n")

        f.write("} // namespace Unicode\n} // namespace MNN\n")
    print(f"  wrote {path}")


def main():
    out_dir = sys.argv[1] if len(sys.argv) > 1 else os.path.join(os.path.dirname(__file__), "..", "..", "transformers", "llm", "engine", "src")
    out_dir = os.path.abspath(out_dir)
    os.makedirs(out_dir, exist_ok=True)
    print(f"Output directory: {out_dir}")
    write_hpp(out_dir)
    write_cpp(out_dir)
    print("Done.")


if __name__ == "__main__":
    main()