#!/usr/bin/env python

# Copyright 2024 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.!
"""Generates ``sentencepiece/__init__.pyi`` from the imported package.

The public Python surface of ``sentencepiece`` is assembled at *import time*
in ``sentencepiece.i`` (``setattr``, ``_add_snake_case``, ``Tokenize = Encode``
etc.), so a static parse of the source cannot see it. This script therefore
imports the built package and introspects the live classes/functions.

Two pieces of information are combined:

* The *method list* is discovered from the runtime objects, so the stub can
  never silently drift when a new public method (e.g. ``ParallelEncode``) or a
  new ``snake_case`` alias is added.
* The *types* come from a small curated overlay below. Anything discovered at
  runtime but missing from the overlay still gets a permissive
  ``(*args, **kwargs) -> Any`` stub (and is reported on stderr) so the stub
  stays complete and the build never breaks.

Usage::

    python tools/gen_stubs.py --output src/sentencepiece/__init__.pyi
    python tools/gen_stubs.py --check   # exit 1 if the committed stub is stale
"""

import argparse
import inspect
import re
import sys

import sentencepiece


HEADER = '''\
# This file is auto-generated by python/tools/gen_stubs.py.
# Do not edit by hand. To update it after changing the public Python API,
# rebuild the package and run:
#     python tools/gen_stubs.py --output src/sentencepiece/__init__.pyi
from typing import Any, IO, Literal, Sequence, overload

__version__: str

_OutType = type[int] | type[str] | type[bytes] | Literal["serialized_proto", "proto", "numpy", "offset_mapping"]
_DecodeOutType = type[str] | type[bytes] | Literal["serialized_proto", "proto", "offset_mapping"]
'''

# Names re-exported / leaked at module scope that are not part of the API.
# Everything genuinely ours has ``__module__ == "sentencepiece"``; this set is
# only used as a defensive backstop for objects that lie about it.
_MODULE_NAME = "sentencepiece"

# Dunder methods worth keeping in the stub. All other dunders are skipped.
ALLOWED_DUNDERS = {"__init__", "__len__", "__getitem__"}

# Class-level attributes that are SWIG/runtime noise.
SKIP_MEMBERS = {"thisown"}

# Order in which classes are emitted. Discovered classes not listed here are
# appended afterwards (sorted), so new public classes still show up.
CLASS_ORDER = [
    "ImmutableSentencePieceText_ImmutableSentencePiece",
    "ImmutableSentencePieceText",
    "ImmutableNBestSentencePieceText",
    "ThreadPool",
    "SentencePieceProcessor",
    "SentencePieceTrainer",
    "SentencePieceNormalizer",
]

# Reusable signature fragments (everything after the method name).
_SPP_INIT = (
    "(self, model_file: str | None = ..., model_proto: bytes | None = ..., "
    "return_type: _OutType | None = ..., add_bos: bool = ..., add_eos: bool = ..., "
    "reverse: bool = ..., emit_unk_piece: bool = ..., enable_sampling: bool = ..., "
    "nbest_size: int = ..., alpha: float = ..., num_threads: int = ..., "
    "out_type: _OutType | None = ...) -> None"
)
_ENCODE = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], return_type: _OutType | None = ..., "
    "add_bos: bool | None = ..., add_eos: bool | None = ..., reverse: bool | None = ..., "
    "emit_unk_piece: bool | None = ..., enable_sampling: bool | None = ..., "
    "nbest_size: int | None = ..., alpha: float | None = ..., num_threads: int | None = ..., "
    "thread_pool: ThreadPool | None = ..., out_type: _OutType | None = ...) -> Any"
)
_PARALLEL_ENCODE = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], return_type: _OutType | None = ..., "
    "add_bos: bool | None = ..., add_eos: bool | None = ..., reverse: bool | None = ..., "
    "emit_unk_piece: bool | None = ..., chunk_len: int | None = ..., "
    "num_threads: int | None = ..., thread_pool: ThreadPool | None = ..., "
    "out_type: _OutType | None = ...) -> Any"
)
_DECODE = (
    "(self, input: int | str | bytes | Sequence[int] | Sequence[str] | Sequence[bytes] "
    "| Sequence[Sequence[int]] | Sequence[Sequence[str]] | Sequence[Sequence[bytes]], "
    "return_type: _DecodeOutType | None = ..., num_threads: int | None = ..., "
    "thread_pool: ThreadPool | None = ..., out_type: _DecodeOutType | None = ..., "
    "return_bytes: bool | None = ...) -> Any"
)
_INPUT_KWARGS = "(self, input: str | bytes | Sequence[str] | Sequence[bytes], **kwargs: Any) -> Any"
_SAMPLE_AS = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], nbest_size: int | None = ..., "
    "alpha: float | None = ..., **kwargs: Any) -> Any"
)
_NBEST_AS = "(self, input: str | bytes | Sequence[str] | Sequence[bytes], nbest_size: int | None = ..., **kwargs: Any) -> Any"
_NBEST = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], return_type: _OutType | None = ..., "
    "add_bos: bool | None = ..., add_eos: bool | None = ..., reverse: bool | None = ..., "
    "emit_unk_piece: bool | None = ..., nbest_size: int | None = ..., "
    "out_type: _OutType | None = ...) -> Any"
)
_SAMPLE_SCORE = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], return_type: _OutType | None = ..., "
    "add_bos: bool | None = ..., add_eos: bool | None = ..., reverse: bool | None = ..., "
    "emit_unk_piece: bool | None = ..., num_samples: int | None = ..., alpha: float | None = ..., "
    "wor: bool | None = ..., include_best: bool | None = ..., out_type: _OutType | None = ...) -> Any"
)
_SAMPLE_SCORE_AS = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes], num_samples: int | None = ..., "
    "alpha: float | None = ..., **kwargs: Any) -> Any"
)
_DECODE_PIECES = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes] | Sequence[Sequence[str]] | Sequence[Sequence[bytes]], "
    "return_type: _DecodeOutType = ..., out_type: _DecodeOutType = ..., **kwargs: Any) -> Any"
)
_DECODE_IDS = (
    "(self, input: int | Sequence[int] | Sequence[Sequence[int]], "
    "return_type: _DecodeOutType = ..., out_type: _DecodeOutType = ..., **kwargs: Any) -> Any"
)
_DECODE_PIECES_SER = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes] | Sequence[Sequence[str]] | Sequence[Sequence[bytes]], "
    "return_type: Literal[\"serialized_proto\"] = ..., out_type: Literal[\"serialized_proto\"] = ..., **kwargs: Any) -> Any"
)
_DECODE_IDS_SER = (
    "(self, input: int | Sequence[int] | Sequence[Sequence[int]], "
    "return_type: Literal[\"serialized_proto\"] = ..., out_type: Literal[\"serialized_proto\"] = ..., **kwargs: Any) -> Any"
)
_DECODE_PIECES_PROTO = (
    "(self, input: str | bytes | Sequence[str] | Sequence[bytes] | Sequence[Sequence[str]] | Sequence[Sequence[bytes]], "
    "**kwargs: Any) -> Any"
)
_DECODE_IDS_PROTO = (
    "(self, input: int | Sequence[int] | Sequence[Sequence[int]], **kwargs: Any) -> Any"
)
_NORMALIZER_INIT = (
    "(self, model_file: str | None = ..., model_proto: bytes | None = ..., "
    "normalizer_spec: bytes | None = ..., "
    "rule_tsv: str | None = ..., rule_name: str | None = ..., "
    "norm_map: Sequence[tuple[str, str]] | None = ..., "
    "add_dummy_prefix: bool = ..., "
    "escape_whitespaces: bool = ..., remove_extra_whitespaces: bool = ...) -> None"
)
_NORMALIZE = "(self, input: str | bytes | Sequence[str] | Sequence[bytes], with_offsets: bool | None = ...) -> Any"

# Curated signatures, keyed by class name then by the *CamelCase* method name.
# snake_case aliases are derived automatically (mirroring _add_snake_case), so
# they never need to be listed here. A list value means @overload.
METHOD_SIGS = {
    "SentencePieceProcessor": {
        "__init__": _SPP_INIT,
        "Init": _SPP_INIT,
        "Load": "(self, model_file: str | None = ..., model_proto: bytes | None = ...) -> bool",
        "LoadFromFile": "(self, arg: str) -> bool",
        "LoadFromSerializedProto": "(self, serialized: bytes) -> bool",
        "SetEncodeExtraOptions": "(self, extra_option: str) -> bool",
        "SetDecodeExtraOptions": "(self, extra_option: str) -> bool",
        "SetVocabulary": "(self, valid_vocab: Sequence[str]) -> bool",
        "ResetVocabulary": "(self) -> bool",
        "LoadVocabulary": "(self, filename: str, threshold: int) -> bool",
        "OverrideNormalizerSpec": "(self, **kwargs: object) -> bool",
        "serialized_model_proto": "(self) -> bytes",
        "GetPieceSize": "(self) -> int",
        # The batch overload uses list[...] (not Sequence[...]) because the
        # runtime only dispatches to the batched path when ``type(arg) is
        # list``; this also avoids a str/Sequence[str] overload overlap.
        "PieceToId": [
            "(self, piece: str) -> int",
            "(self, piece: list[str]) -> list[int]",
        ],
        "IdToPiece": [
            "(self, id: int) -> str",
            "(self, id: list[int]) -> list[str]",
        ],
        "GetScore": [
            "(self, id: int) -> float",
            "(self, id: list[int]) -> list[float]",
        ],
        "IsUnknown": [
            "(self, id: int) -> bool",
            "(self, id: list[int]) -> list[bool]",
        ],
        "IsControl": [
            "(self, id: int) -> bool",
            "(self, id: list[int]) -> list[bool]",
        ],
        "IsUnused": [
            "(self, id: int) -> bool",
            "(self, id: list[int]) -> list[bool]",
        ],
        "IsByte": [
            "(self, id: int) -> bool",
            "(self, id: list[int]) -> list[bool]",
        ],
        "unk_id": "(self) -> int",
        "bos_id": "(self) -> int",
        "eos_id": "(self) -> int",
        "pad_id": "(self) -> int",
        "Encode": _ENCODE,
        "EncodeAsPieces": _INPUT_KWARGS,
        "EncodeAsIds": _INPUT_KWARGS,
        "EncodeAsNumpy": _INPUT_KWARGS,
        "EncodeAsSerializedProto": _INPUT_KWARGS,
        "EncodeAsProto": _INPUT_KWARGS,
        "EncodeAsOffsetMapping": _INPUT_KWARGS,
        "EncodeAsImmutableProto": _INPUT_KWARGS,
        "ParallelEncode": _PARALLEL_ENCODE,
        "ParallelEncodeAsPieces": _INPUT_KWARGS,
        "ParallelEncodeAsIds": _INPUT_KWARGS,
        "ParallelEncodeAsNumpy": _INPUT_KWARGS,
        "ParallelEncodeAsSerializedProto": _INPUT_KWARGS,
        "ParallelEncodeAsProto": _INPUT_KWARGS,
        "ParallelEncodeAsImmutableProto": _INPUT_KWARGS,
        "SampleEncodeAsPieces": _SAMPLE_AS,
        "SampleEncodeAsIds": _SAMPLE_AS,
        "SampleEncodeAsNumpy": _SAMPLE_AS,
        "SampleEncodeAsSerializedProto": _SAMPLE_AS,
        "SampleEncodeAsProto": _SAMPLE_AS,
        "SampleEncodeAsImmutableProto": _SAMPLE_AS,
        "NBestEncode": _NBEST,
        "NBestEncodeAsPieces": _NBEST_AS,
        "NBestEncodeAsIds": _NBEST_AS,
        "NBestEncodeAsNumpy": _NBEST_AS,
        "NBestEncodeAsSerializedProto": _NBEST_AS,
        "NBestEncodeAsProto": _NBEST_AS,
        "NBestEncodeAsImmutableProto": _NBEST_AS,
        "SampleEncodeAndScore": _SAMPLE_SCORE,
        "SampleEncodeAndScoreAsPieces": _SAMPLE_SCORE_AS,
        "SampleEncodeAndScoreAsIds": _SAMPLE_SCORE_AS,
        "SampleEncodeAndScoreAsSerializedProto": _SAMPLE_SCORE_AS,
        "SampleEncodeAndScoreAsProto": _SAMPLE_SCORE_AS,
        "SampleEncodeAndScoreAsImmutableProto": _SAMPLE_SCORE_AS,
        "Decode": _DECODE,
        "DecodePieces": _DECODE_PIECES,
        "DecodeIds": _DECODE_IDS,
        "DecodePiecesAsSerializedProto": _DECODE_PIECES_SER,
        "DecodeIdsAsSerializedProto": _DECODE_IDS_SER,
        "DecodePiecesAsProto": _DECODE_PIECES_PROTO,
        "DecodeIdsAsProto": _DECODE_IDS_PROTO,
        "DecodePiecesAsImmutableProto": _DECODE_PIECES_PROTO,
        "DecodeIdsAsImmutableProto": _DECODE_IDS_PROTO,
        "from_file": "(cls, model_file: str, return_type: _OutType | None = ..., **kwargs: Any) -> SentencePieceProcessor",
        "from_proto": "(cls, model_proto: bytes, return_type: _OutType | None = ..., **kwargs: Any) -> SentencePieceProcessor",
        "CalculateEntropy": "(self, input: str | Sequence[str], alpha: float, num_threads: int | None = ...) -> float | list[float]",
        "Normalize": _NORMALIZE,
        "Tokenize": _ENCODE,
        "Detokenize": _DECODE,
        "piece_size": "(self) -> int",
        "vocab_size": "(self) -> int",
        "__len__": "(self) -> int",
        "__getitem__": "(self, piece: str) -> int",
    },
    "SentencePieceTrainer": {
        "__init__": "(self) -> None",
        "Train": "(arg: str | None = ..., logstream: IO[Any] | None = ..., normalizer: SentencePieceNormalizer | None = ..., **kwargs: Any) -> None",
    },
    "SentencePieceNormalizer": {
        "__init__": _NORMALIZER_INIT,
        "Init": _NORMALIZER_INIT,
        "LoadFromSerializedProto": "(self, serialized: bytes) -> bool",
        "LoadFromSerializedNormalizerSpec": "(self, serialized: bytes) -> bool",
        "LoadFromRuleTSV": "(self, filename: str) -> bool",
        "LoadFromRuleName": "(self, name: str) -> bool",
        "LoadFromFile": "(self, arg: str) -> bool",
        "LoadFromMap": "(self, norm_map: Sequence[tuple[str, str]]) -> bool",
        "Decompile": "(self) -> list[tuple[str, str]]",
        "serialized_model_proto": "(self) -> bytes",
        "serialized_normalizer_spec": "(self) -> bytes",
        "Normalize": _NORMALIZE,
    },
    "ThreadPool": {
        "__init__": "(self, num_threads: int) -> None",
        "num_threads": "(self) -> int",
    },
    "ImmutableSentencePieceText": {
        "SerializeAsString": "(self) -> bytes",
    },
    "ImmutableNBestSentencePieceText": {
        "SerializeAsString": "(self) -> bytes",
    },
}

# Property return types, keyed by class name then property name.
PROPERTY_TYPES = {
    "ImmutableSentencePieceText_ImmutableSentencePiece": {
        "piece": "str",
        "piece_as_bytes": "bytes",
        "surface": "str",
        "surface_as_bytes": "bytes",
        "id": "int",
        "begin": "int",
        "end": "int",
    },
    "ImmutableSentencePieceText": {
        "text": "str",
        "text_as_bytes": "bytes",
        "score": "float",
        "pieces": "Sequence[ImmutableSentencePieceText_ImmutableSentencePiece]",
    },
    "ImmutableNBestSentencePieceText": {
        "nbests": "Sequence[ImmutableSentencePieceText]",
    },
}

# Methods that are static (no ``self``).
STATIC_METHODS = {("SentencePieceTrainer", "Train")}

# Module-level free functions.
FUNCTION_SIGS = {
    "SetRandomGeneratorSeed": "(seed: int) -> None",
    "SetMinLogLevel": "(v: int) -> None",
    "SetNBestTimeout": "(timeout_ms: int) -> None",
    "SetDataDir": "(data_dir: str) -> None",
}


def to_snake(name):
  """Mirror of sentencepiece._add_snake_case for CamelCase -> snake_case."""
  return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower().replace("n_best", "nbest")


def build_sig_map(class_name):
  """Curated signatures expanded with their derived snake_case aliases."""
  sigs = {}
  for camel, sig in METHOD_SIGS.get(class_name, {}).items():
    sigs[camel] = sig
    if re.match(r"^[A-Z]+", camel):
      sigs[to_snake(camel)] = sig
  return sigs


def build_static_set(class_name):
  """Static method names for a class, including derived snake_case aliases."""
  statics = set()
  for cls_name, method in STATIC_METHODS:
    if cls_name != class_name:
      continue
    statics.add(method)
    if re.match(r"^[A-Z]+", method):
      statics.add(to_snake(method))
  return statics


def is_property(cls, name):
  return isinstance(inspect.getattr_static(cls, name, None), property)


def is_classmethod(cls, name):
  return isinstance(inspect.getattr_static(cls, name, None), classmethod)


def public_members(cls):
  """Public (and selected dunder) member names of a class, deduped/sorted."""
  members = set()
  for name in dir(cls):
    if name in SKIP_MEMBERS:
      continue
    if name.startswith("__"):
      if name not in ALLOWED_DUNDERS:
        continue
    elif name.startswith("_"):
      continue
    members.add(name)
  return members


def _method_sort_key(name):
  if name.startswith("__"):
    group = 0
  elif re.match(r"^[A-Z]", name):
    group = 1
  else:
    group = 2
  return (group, name)


def emit_method(name, sig, indent, static=False, is_cls=False):
  pad = " " * indent
  lines = []
  variants = sig if isinstance(sig, list) else [sig]
  overloaded = len(variants) > 1
  for variant in variants:
    if static:
      lines.append(f"{pad}@staticmethod")
    elif is_cls:
      lines.append(f"{pad}@classmethod")
    if overloaded:
      lines.append(f"{pad}@overload")
    lines.append(f"{pad}def {name}{variant}: ...")
  return lines


def emit_class(cls, class_name, uncurated):
  sig_map = build_sig_map(class_name)
  static_set = build_static_set(class_name)
  members = public_members(cls)

  properties = sorted(n for n in members if is_property(cls, n))
  methods = sorted((n for n in members if n not in properties), key=_method_sort_key)

  lines = [f"class {class_name}:"]
  body = []

  prop_types = PROPERTY_TYPES.get(class_name, {})
  for name in properties:
    typ = prop_types.get(name)
    if typ is None:
      uncurated.append(f"{class_name}.{name} (property)")
      typ = "Any"
    body.append("    @property")
    body.append(f"    def {name}(self) -> {typ}: ...")

  for name in methods:
    static = name in static_set
    is_cls = is_classmethod(cls, name)
    sig = sig_map.get(name)
    if sig is None:
      uncurated.append(f"{class_name}.{name}")
      if static:
        sig = "(*args: Any, **kwargs: Any) -> None" if name == "__init__" else "(*args: Any, **kwargs: Any) -> Any"
      elif is_cls:
        sig = "(cls, *args: Any, **kwargs: Any) -> Any"
      else:
        sig = "(self, *args: Any, **kwargs: Any) -> None" if name == "__init__" else "(self, *args: Any, **kwargs: Any) -> Any"
    body.extend(emit_method(name, sig, indent=4, static=static, is_cls=is_cls))

  if not body:
    body.append("    ...")
  lines.extend(body)
  return "\n".join(lines)


def discover_classes(module):
  found = {}
  for name in dir(module):
    if name.startswith("_"):
      continue
    obj = getattr(module, name)
    if isinstance(obj, type) and getattr(obj, "__module__", "").split('.')[0] == _MODULE_NAME:
      found[name] = obj
  return found


def discover_functions(module):
  found = {}
  for name in dir(module):
    if name.startswith("_"):
      continue
    obj = getattr(module, name)
    if isinstance(obj, type):
      continue
    if callable(obj) and getattr(obj, "__module__", "").split('.')[0] == _MODULE_NAME:
      found[name] = obj
  return found


def generate():
  uncurated = []
  classes = discover_classes(sentencepiece)
  functions = discover_functions(sentencepiece)

  ordered_names = [c for c in CLASS_ORDER if c in classes]
  ordered_names += sorted(c for c in classes if c not in CLASS_ORDER)

  blocks = [HEADER.rstrip("\n")]
  for name in ordered_names:
    blocks.append(emit_class(classes[name], name, uncurated))

  func_lines = []
  for name in sorted(functions, key=_method_sort_key):
    sig = FUNCTION_SIGS.get(name)
    if sig is None:
      camel = next(
          (c for c in FUNCTION_SIGS if to_snake(c) == name), None
      )
      sig = FUNCTION_SIGS.get(camel) if camel else None
    if sig is None:
      uncurated.append(name)
      sig = "(*args: Any, **kwargs: Any) -> Any"
    func_lines.append(f"def {name}{sig}: ...")
  if func_lines:
    blocks.append("\n".join(func_lines))

  text = "\n\n\n".join(blocks) + "\n"
  return text, uncurated


def main():
  parser = argparse.ArgumentParser(description=__doc__)
  parser.add_argument("--output", help="Write the stub to this path.")
  parser.add_argument(
      "--check",
      help="Compare against this existing stub and exit 1 if it is stale.",
  )
  args = parser.parse_args()

  text, uncurated = generate()

  if uncurated:
    sys.stderr.write(
        "gen_stubs: emitted permissive stubs for uncurated members "
        "(add typed signatures to METHOD_SIGS/PROPERTY_TYPES to improve):\n"
    )
    for name in uncurated:
      sys.stderr.write(f"  - {name}\n")

  if args.check:
    with open(args.check, encoding="utf-8") as f:
      current = f.read()
    if current != text:
      sys.stderr.write(
          f"\n{args.check} is out of date. Regenerate with:\n"
          "    python tools/gen_stubs.py --output "
          f"{args.check}\n"
      )
      return 1
    return 0

  if args.output:
    with open(args.output, "w", encoding="utf-8") as f:
      f.write(text)
  else:
    sys.stdout.write(text)
  return 0


if __name__ == "__main__":
  sys.exit(main())
