"""
Script to generate Python typecheck test files for all possible ABI type string argument combinations.

Usage:
    python scripts/generate_overload_tests.py [--impl codec|abi|both] [--suite fixed|var|both] [--sample-1-of-x N] [--seed S] [--lengths L1,L2,...]

This will overwrite test files in:
- tests/typecheck/codec/decode/fixed/len_1/overload_test_data_*.py (tuple-based)
- tests/typecheck/codec/decode/fixed/len_2/overload_test_data_*.py
- tests/typecheck/codec/decode/fixed/len_3/overload_test_data_*.py
- tests/typecheck/codec/decode/variable/len_1/overload_test_data_*.py (iterable-based)
- tests/typecheck/codec/decode/variable/len_2/overload_test_data_*.py
- tests/typecheck/codec/decode/variable/len_3/overload_test_data_*.py
- tests/typecheck/abi/decode/fixed/len_1/overload_test_data_*.py (tuple-based, abi.decode)
- tests/typecheck/abi/decode/variable/len_1/overload_test_data_*.py (iterable-based, abi.decode)
...

Test files are chunked in 10,000-case intervals.

Key features:
- Streams generation: cases are written chunk-by-chunk, not all held in memory.
- Progress bar (tqdm) measures chunk files, not individual cases.
- For TupleXXXIntTypeStr, only 10 representative values are sampled (ending with int8, int16, int64, int128, int256, and their uint variants).
- Each chunk's lines are collected and written in a single operation for efficiency.
- Covers: all int, uint, bytes, bool, string, arrays, and all possible combinations up to length 3, including fallback/unknown types in all permutations.
- If you add new ABI type strings, rerun this script to update the test files.

New options:
- --suite: Restrict to 'fixed', 'var', or 'both' (default: both)
- --sample-1-of-x N: Only generate a random 1/N of the chunk files for each length (default: 1, i.e., all files)
- --seed S: Seed for random sampling (default: 42). Controls which chunk files are selected when sampling, for reproducibility.
- --lengths L1,L2,...: Only generate test data for the specified lengths (e.g., --lengths 1,2,3)
"""

import argparse
import itertools
import math
import pathlib
import random
import re
import sys
from typing import (
    Any,
    Dict,
    Final,
    ForwardRef,
    Iterator,
    List,
    Literal,
    Set,
    Tuple,
    Union,
    get_args,
    get_origin,
)

from tqdm import (
    tqdm,
)
from typing_extensions import (
    LiteralString,
)

MAX_LEN_FIXED: Final = 3
MAX_LEN_VARIABLE: Final = 4
CHUNK_SIZE: Final = 10_000
FILENAME_PREFIX: Final = "overload_test_data_"

ANY_ALIAS: Final = "A"
ASSERT_TYPE_ALIAS: Final = "ck"
BOOL_ALIAS: Final = "bo"
BYTES_ALIAS: Final = "by"
DATA_ALIAS: Final = "D"
FUNC_ALIAS: Final = "f"
HEXADDRESS_ALIAS: Final = "HA"
INT_ALIAS: Final = "i"
STR_ALIAS: Final = "s"
TUPLE_ALIAS: Final = "T"
TUPLE_ANY_ALIAS: Final = "TA"
UNION_ALIAS: Final = "U"


HEADER_CODEC = f"""# AUTOGENERATED: Run scripts/generate_overload_tests.py to update.
# Requires mypy extras
# To run: pip install mypy && mypy tests/typecheck

# alias some things so we can minimize test data size in source control
from typing import Any as {ANY_ALIAS}
from typing import Tuple as {TUPLE_ALIAS}
from typing import Union as {UNION_ALIAS}
from typing_extensions import assert_type as {ASSERT_TYPE_ALIAS}
from eth_typing import HexAddress as {HEXADDRESS_ALIAS}
from faster_eth_abi.codec import ABIDecoder
from faster_eth_abi.registry import ABIRegistry

# more aliases to save a bunch more chars
{BOOL_ALIAS} = bool
{BYTES_ALIAS} = bytes
{INT_ALIAS} = int
{STR_ALIAS} = str

{TUPLE_ANY_ALIAS} = T[A, ...]

THA = T[HA, ...]
Tbo = T[bo, ...]
Tby = T[by, ...]
Ti = T[i, ...]
Ts = T[s, ...]
TUa_bo = T[U[HA, bo], ...]
TUa_by = T[U[HA, by], ...]
TUa_i = T[U[HA, i], ...]
TUa_s = T[U[HA, s], ...]
TUbo_by = T[U[bo, by], ...]
TUbo_i = T[U[bo, i], ...]
TUbo_s = T[U[bo, s], ...]
TUby_i = T[U[by, i], ...]
TUby_s = T[U[by, s], ...]
TUi_s = T[U[i, s], ...]
TUa_bo_by = T[U[bo, by, HA], ...]
TUbo_by_i = T[U[bo, by, i], ...]
TUbo_by_s = T[U[bo, by, s], ...]
TUa_bo_i = T[U[bo, i, HA], ...]
TUa_by_i = T[U[by, i, HA], ...]
TUby_i_s = T[U[by, i, s], ...]

decoder = ABIDecoder(ABIRegistry())

# the callable we are checking
{FUNC_ALIAS} = decoder.decode

# input data
{DATA_ALIAS} = b"\\x00" * 32
"""

HEADER_ABI = f"""# AUTOGENERATED: Run scripts/generate_overload_tests.py to update.
# Requires mypy extras
# To run: pip install mypy && mypy tests/typecheck

# alias some things so we can minimize test data size in source control
from typing import Any as {ANY_ALIAS}
from typing import Tuple as {TUPLE_ALIAS}
from typing import Union as {UNION_ALIAS}
from typing_extensions import assert_type as {ASSERT_TYPE_ALIAS}
from eth_typing import HexAddress as {HEXADDRESS_ALIAS}
from faster_eth_abi import abi

# more aliases to save a bunch more chars
{BOOL_ALIAS} = bool
{BYTES_ALIAS} = bytes
{INT_ALIAS} = int
{STR_ALIAS} = str

{TUPLE_ANY_ALIAS} = T[A, ...]

THA = T[HA, ...]
Tbo = T[bo, ...]
Tby = T[by, ...]
Ti = T[i, ...]
Ts = T[s, ...]
TUa_bo = T[U[HA, bo], ...]
TUa_by = T[U[HA, by], ...]
TUa_i = T[U[HA, i], ...]
TUa_s = T[U[HA, s], ...]
TUbo_by = T[U[bo, by], ...]
TUbo_i = T[U[bo, i], ...]
TUbo_s = T[U[bo, s], ...]
TUby_i = T[U[by, i], ...]
TUby_s = T[U[by, s], ...]
TUi_s = T[U[i, s], ...]
TUa_bo_by = T[U[bo, by, HA], ...]
TUbo_by_i = T[U[bo, by, i], ...]
TUbo_by_s = T[U[bo, by, s], ...]
TUa_bo_i = T[U[bo, i, HA], ...]
TUa_by_i = T[U[by, i, HA], ...]
TUby_i_s = T[U[by, i, s], ...]

# the callable we are checking
{FUNC_ALIAS} = abi.decode

# input data
{DATA_ALIAS} = b"\\x00" * 32
"""

RETURN_TYPE_MAP = {
    "address": HEXADDRESS_ALIAS,
    "address[]": ANY_ALIAS,  # f"{TUPLE_ALIAS}[{HEXADDRESS_ALIAS}, ...]",
    "bytes": BYTES_ALIAS,
    "bytes[]": ANY_ALIAS,  # f"{TUPLE_ALIAS}[{BYTES_ALIAS}, ...]",
    "string": STR_ALIAS,
    "string[]": ANY_ALIAS,  # f"{TUPLE_ALIAS}[{STR_ALIAS}, ...]",
    "int": INT_ALIAS,
    "bool": BOOL_ALIAS,
    "bool[]": ANY_ALIAS,  # f"{TUPLE_ALIAS}[{BOOL_ALIAS}, ...]",
    "?": ANY_ALIAS,
}
for i in range(8, 257, 8):
    RETURN_TYPE_MAP[f"int{i}"] = INT_ALIAS
    RETURN_TYPE_MAP[f"uint{i}"] = INT_ALIAS
    RETURN_TYPE_MAP[f"int{i}[]"] = f"{TUPLE_ALIAS}[{INT_ALIAS}, ...]"
    RETURN_TYPE_MAP[f"uint{i}[]"] = f"{TUPLE_ALIAS}[{INT_ALIAS}, ...]"
for i in range(1, 33):
    RETURN_TYPE_MAP[f"bytes{i}"] = BYTES_ALIAS
    RETURN_TYPE_MAP[f"bytes{i}[]"] = f"{TUPLE_ALIAS}[{BYTES_ALIAS}, ...]"

VAR_TUPLE_ALIASES = {
    # this shouldnt be needed but just in case
    "T[A, ...]": TUPLE_ANY_ALIAS,
    # single type
    "T[HA, ...]": "THA",
    "T[bo, ...]": "Tbo",
    "T[by, ...]": "Tby",
    "T[i, ...]": "Ti",
    "T[s, ...]": "Ts",
    # TODO: add some more later, this is good for now
    "T[U[HA, bo], ...]": "TUa_bo",
    "T[U[HA, by], ...]": "TUa_by",
    "T[U[HA, i], ...]": "TUa_i",
    "T[U[HA, s], ...]": "TUa_s",
    "T[U[bo, by], ...]": "TUbo_by",
    "T[U[bo, i], ...]": "TUbo_i",
    "T[U[bo, s], ...]": "TUbo_s",
    "T[U[by, i], ...]": "TUby_i",
    "T[U[by, s], ...]": "TUby_s",
    "T[U[i, str], ...]": "TUi_s",
    "T[U[bo, by, HA], ...]": "TUa_bo_by",
    "T[U[bo, by, i], ...]": "TUbo_by_i",
    "T[U[bo, by, s], ...]": "TUbo_by_s",
    "T[U[bo, i, HA], ...]": "TUa_bo_i",
    "T[U[by, i, HA], ...]": "TUa_by_i",
    "T[U[by, i, s], ...]": "TUby_i_s",
}


def build_alias_map() -> Dict[str, Any]:
    sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
    import faster_eth_abi.typing

    alias_map = {}
    for name in dir(faster_eth_abi.typing):
        if not name.startswith("_"):
            obj = getattr(faster_eth_abi.typing, name)
            alias_map[name] = obj
    return alias_map


def all_integers_in_set(s: str, allowed: Set[int]) -> bool:
    """
    Parses the input string for all integer values (including negative numbers)
    and returns True if all found integers are members of the allowed set.
    """
    found = {int(x) for x in re.findall(r"-?\d+", s)}
    return found.issubset(allowed)


def extract_all_literals(typ, alias_map, alias_path=None):
    if alias_path is None:
        alias_path = []
    if typ is str:
        raise ValueError(
            f"Encountered plain 'str' type in alias path: {' -> '.join(alias_path)}."
        )
    origin = get_origin(typ)
    if origin is Union:
        for arg in get_args(typ):
            yield from extract_all_literals(arg, alias_map, alias_path)
        return
    elif origin is Literal:
        if alias_path[-1] in ("BytesArrayTypeStr", "IntArrayTypeStr"):
            # TODO: we need to implement overloads for these, for now we only have the literals defined
            print(f"skipping {alias_path}")
            return
        if alias_path[-1].startswith("Tuple") and alias_path[-1].endswith("IntTypeStr"):
            # for now, we will exclude these. We need to implement them
            # in the actual overloads, which takes a new script. stay tuned.
            return
            for value in get_args(typ):
                if all_integers_in_set(value, {8, 16, 64, 128, 256}):
                    if value not in RETURN_TYPE_MAP:
                        parentheses_stripped = value[1:-1]
                        types = parentheses_stripped.split(",")
                        RETURN_TYPE_MAP[
                            value
                        ] = f"{TUPLE_ALIAS}[{', '.join(RETURN_TYPE_MAP[t] for t in types)}]"
                    yield value
            return
        else:
            yield from get_args(typ)
    elif isinstance(typ, ForwardRef):
        ref_name = typ.__forward_arg__
        if ref_name in alias_map:
            yield from extract_all_literals(
                alias_map[ref_name], alias_map, alias_path + [ref_name]
            )
            return
        raise RuntimeError(
            f"ForwardRef {ref_name!r} not found in alias map (path: {' -> '.join(alias_path)})"
        )
    else:
        raise RuntimeError(
            f"Type {typ!r} is not a Union, Literal, ForwardRef, or str (origin={origin!r}, path: {' -> '.join(alias_path)})"
        )


def expand_literal_values() -> Iterator[LiteralString]:
    alias_map = build_alias_map()
    alias_names = [
        "AddressTypeStr",
        "BytesTypeStr",
        "StringTypeStr",
        "IntegerTypeStr",
        "BoolTypeStr",
        "ArrayTypeStr",
        "TupleTypeStr",
        "DecodesToIntTypeStr",
    ]
    for name in sorted(alias_names):
        if name == "DecodesToIntTypeStr":
            from faster_eth_abi.codec import (
                DecodesToIntTypeStr,
            )

            typ = DecodesToIntTypeStr
        else:
            typ = alias_map[name]
        yield from map(str, extract_all_literals(typ, alias_map, [name]))


ALL_LITERALS: Final = sorted(set(expand_literal_values())) + ["?"]
NUM_LITERALS = len(ALL_LITERALS)
print(f"Found {NUM_LITERALS} ABI type string literals (including fallback).")


def get_expected_type_tuple(types: Tuple[str, ...]) -> str:
    py_types = [ANY_ALIAS if t == "?" else RETURN_TYPE_MAP[t] for t in types]
    return f"{TUPLE_ALIAS}[{', '.join(py_types)}]"


def get_expected_type_iterable(types: Tuple[str, ...]) -> str:
    py_types = [ANY_ALIAS if t == "?" else RETURN_TYPE_MAP[t] for t in types]
    unique_types = sorted(set(py_types))
    if ANY_ALIAS in unique_types:
        expected = TUPLE_ANY_ALIAS
    elif len(unique_types) == 1:
        expected = f"{TUPLE_ALIAS}[{unique_types[0]}, ...]"
    else:
        expected = f"{TUPLE_ALIAS}[{UNION_ALIAS}[{', '.join(unique_types)}], ...]"
    return VAR_TUPLE_ALIASES.get(expected) or expected


def compute_total_chunks(L: int) -> int:
    total_cases = NUM_LITERALS**L
    return math.ceil(total_cases / CHUNK_SIZE)


def stream_cases_and_write_files(
    mode: Literal["fixed", "variable"],
    impl: Literal["abi", "codec"],
    sample_1_of_x: int = 1,
    seed: int = 42,
    lengths: List[int] = None,
) -> int:
    FOLDER = "fixed" if mode == "fixed" else "variable"
    if lengths is None:
        MAX_LEN = MAX_LEN_FIXED if mode == "fixed" else MAX_LEN_VARIABLE
        lengths = list(range(1, MAX_LEN + 1))

    case_counter = 0

    for L in lengths:
        total_chunks = compute_total_chunks(L)
        chunk_indices = list(range(1, total_chunks + 1))
        if sample_1_of_x > 1:
            random.seed(seed)
            num_to_pick = max(1, total_chunks // sample_1_of_x)
            chunk_indices = sorted(random.sample(chunk_indices, num_to_pick))

        with tqdm(
            total=len(chunk_indices),
            desc=f"Streaming {mode} chunks for {impl} (len_{L})",
        ) as progress:
            PATH_WITHOUT_NUMBER = (
                f"tests/typecheck/{impl}/decode/{FOLDER}/len_{L}/{FILENAME_PREFIX}"
            )
            chunk_idx = 1
            chunk_lines = []
            file_counter = 0

            def format_case(typestr: str, expected: str) -> str:
                nonlocal case_counter, file_counter
                case = f"{ASSERT_TYPE_ALIAS}({FUNC_ALIAS}({typestr}, {DATA_ALIAS}), {expected})  # {mode} case {case_counter}\n"
                case_counter += 1
                file_counter += 1
                return case

            if mode == "fixed":

                def addline(case: Tuple[str, ...]) -> Iterator[str]:
                    typestr = (
                        f"('{case[0]}',)"
                        if len(case) == 1
                        else f"({', '.join(map(repr, case))})"
                    )
                    expected = get_expected_type_tuple(case)
                    chunk_lines.append(format_case(typestr, expected))
                    # chunk_lines.append(format_case("ABICodec.decode", f"codec, {typestr}", expected)
                    # chunk_lines.append(format_case("ABIDecoder.decode", f"codec, {typestr}", expected)
                    # chunk_lines.append(format_case("ABIDecoder.decode", f"decoder, {typestr}", expected)

            else:

                def addline(case: Tuple[str, ...]) -> Iterator[str]:
                    typestr = f"[{', '.join(map(repr, case))}]"
                    expected = get_expected_type_iterable(case)
                    chunk_lines.append(format_case(typestr, expected))
                    # chunk_lines.append(format_case("ABICodec.decode", f"codec, {typestr}", expected)
                    # chunk_lines.append(format_case("ABIDecoder.decode", f"codec, {typestr}", expected)
                    # chunk_lines.append(format_case("ABIDecoder.decode", f"decoder, {typestr}", expected)

            def write_file() -> None:
                nonlocal chunk_idx
                if sample_1_of_x == 1 or chunk_idx in chunk_indices:
                    path = pathlib.Path(f"{PATH_WITHOUT_NUMBER}{chunk_idx:05d}.py")
                    path.parent.mkdir(parents=True, exist_ok=True)
                    with open(path, "w") as f:
                        f.write(HEADER_ABI if impl == "abi" else HEADER_CODEC)
                        f.writelines(chunk_lines)
                    progress.update(1)
                chunk_idx += 1
                chunk_lines.clear()

            for combo in itertools.product(ALL_LITERALS, repeat=L):
                if file_counter and file_counter % CHUNK_SIZE == 0:
                    write_file()
                addline(combo)

            if chunk_lines:
                write_file()

            print(f"finished generating cases for length {L}")

    return case_counter


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--impl", choices=["codec", "abi", "both"], default="both")
    parser.add_argument(
        "--suite",
        choices=["fixed", "variable", "both"],
        default="both",
        help="Restrict to 'fixed', 'var', or 'both' test suites",
    )
    parser.add_argument(
        "--sample-1-of-x",
        type=int,
        default=1,
        help="Only generate a random 1/N of the chunk files for each length (default: 1, i.e., all files)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Seed for random sampling (default: 42). Controls which chunk files are selected when sampling, for reproducibility.",
    )
    parser.add_argument(
        "--lengths",
        type=str,
        help="Comma-separated list of lengths to generate (e.g., 1,2,3,4)",
    )
    args = parser.parse_args()

    impls = ["codec", "abi"] if args.impl == "both" else [args.impl]
    modes = ["fixed", "variable"] if args.suite == "both" else [args.suite]
    lengths = None
    if args.lengths:
        lengths = [int(x) for x in args.lengths.split(",") if x.strip()]
    for impl in impls:
        for mode in modes:
            print(
                f"Streaming {'tuple' if mode == 'fixed' else 'iterable'}-based test cases for {impl}..."
            )
            stream_cases_and_write_files(
                mode=mode,
                impl=impl,
                sample_1_of_x=args.sample_1_of_x,
                seed=args.seed,
                lengths=lengths,
            )

    print("Test generation complete.")


if __name__ == "__main__":
    main()
