from pathlib import Path
from typing import Any, List, Optional

output_file = Path(__file__).parent.resolve() / "chia_rs.pyi"

def transform_type(m: str) -> str:
    n, t = m.split(":")
    if "List[" in t:
        t = t.replace("List[", "Sequence[")
    elif "bytes32" == t.strip():
        t = " bytes"
    return f"{n}:{t}"

def print_class(f: Any, name: str, members: List[str], extra: Optional[str] = None):

    # f-strings don't allow backslashes, which makes it a bit tricky to
    # manipulate strings with newlines
    nl = '\n'
    f.write(f"""
class {name}:
    {(nl + '    ').join(members)}{nl + extra if extra else ''}
    def __init__(
        self,
        {(',' + nl + '        ').join(map(transform_type, members))}
    ) -> None: ...
    def __hash__(self) -> int: ...
    def __str__(self) -> str: ...
    def __repr__(self) -> str: ...
    def __richcmp__(self) -> Any: ...
    def __deepcopy__(self) -> {name}: ...
    def __copy__(self) -> {name}: ...
    @staticmethod
    def from_bytes(bytes) -> {name}: ...
    @staticmethod
    def parse_rust(ReadableBuffer) -> Tuple[{name}, int]: ...
    def to_bytes(self) -> bytes: ...
    def __bytes__(self) -> bytes: ...
    def get_hash(self) -> bytes32: ...
    def to_json_dict(self) -> Dict[str, Any]: ...
    @staticmethod
    def from_json_dict(o: Dict[str, Any]) -> {name}: ...
""")

with open(output_file, "w") as f:
    f.write("""
#
# this file is generated by generate_type_stubs.py
#

from typing import List, Optional, Sequence, Tuple
from chia.types.blockchain_format.sized_bytes import bytes32

ReadableBuffer = Union[bytes, bytearray, memoryview]

def compute_merkle_set_root(items: Sequence[bytes]) -> bytes: ...

def run_generator(
    program: bytes, args: bytes, max_cost: int, flags: int
) -> Tuple[Optional[int], Optional[SpendBundleConditions]]: ...

COND_ARGS_NIL: int = ...
NO_UNKNOWN_CONDS: int = ...
STRICT_ARGS_COUNT: int = ...
MEMPOOL_MODE: int = ...

NO_NEG_DIV: int = ...
NO_UNKNOWN_OPS: int = ...

def run_chia_program(
    program: bytes, args: bytes, max_cost: int, flags: int
) -> Pair[int, LazyNode]: ...

class LazyNode:
    def pair() -> Optional[Tuple[LazyNode, LazyNode]]: ...
    def atom() -> bytes: ...

def serialized_length(program: bytes) -> int: ...
def tree_hash(program: ReadableBuffer) -> bytes32: ...
def get_puzzle_and_solution_for_coin(program: bytes, args: bytes, max_cost: int, find_parent: bytes32, find_amount: int, find_ph: bytes32) -> Tuple[bytes, bytes]: ...
""")

    print_class(f, "Spend", [
        "coin_id: bytes",
        "puzzle_hash: bytes",
        "height_relative: Optional[int]",
        "seconds_relative: int",
        "create_coin: List[Tuple[bytes, int, Optional[bytes]]]",
        "agg_sig_me: List[Tuple[bytes, bytes]]",
    ])

    print_class(f, "SpendBundleConditions", [
        "spends: List[Spend]",
        "reserve_fee: int",
        "height_absolute: int",
        "seconds_absolute: int",
        "agg_sig_unsafe: List[Tuple[bytes, bytes]]",
        "cost: int",
    ])

    print_class(f, "Coin", [
        "parent_coin_info: bytes32",
        "puzzle_hash: bytes32",
        "amount: int",
    ], "    def name(self) -> bytes32: ...")

    print_class(f, "CoinState", [
        "coin: Coin",
        "spent_height: Optional[int]",
        "created_height: Optional[int]",
    ])

    print_class(f, "RespondToPhUpdates", [
        "puzzle_hashes: List[bytes32]",
        "min_height: int",
        "coin_states: List[CoinState]",
    ])
