#!/usr/bin/env python3
import os
import random
import string

OUTPUT_DIR = "/tmp"
DECORATORS = [
    "# @pipeline",
    "# @fallible",
    "# @parallel",
    "# @export",
    "# @timeout(60)",
    "# @retry(3, 5)",
    "# @if(condition)",
    "# @after(dep)",
    "# @pipe(step1, step2)",
    "# @health(/health)",
    "# @loop(10)",
]
FUNCTION_BODIES = [
    "echo 'Processing...';",
    "local var=$(date +%s);",
    "if [ -f /tmp/test ]; then echo 'exists'; fi;",
    "for i in $(seq 1 10); do echo $i; done;",
    "case $1 in start) echo 'starting';; stop) echo 'stopping';; esac;",
    "read -r line < /dev/stdin;",
    "export PATH=$PATH:/usr/local/bin;",
    "set -e; set -u;",
    "trap 'echo error' ERR;",
    "cd /tmp || exit 1;",
    "local result=$(grep -r 'pattern' /var/log/ 2>/dev/null);",
    'for f in /tmp/*.tmp; do rm -f "$f"; done',
    "if [[ $var -gt 100 ]]; then echo 'large'; fi;",
    'while read -r line; do echo "$line"; done < /etc/passwd;',
    'select opt in a b c; do echo "Selected $opt"; break; done;',
]


def random_string(length=8):
    return "".join(random.choices(string.ascii_letters + string.digits, k=length))


def generate_function(body_lines):
    lines = []
    func_name = f"func_{random_string(6)}"
    for _ in range(random.randint(0, 3)):
        lines.append(random.choice(DECORATORS))
    lines.append(f"{func_name}() {{")
    for _ in range(body_lines):
        lines.append("    " + random.choice(FUNCTION_BODIES))
    lines.append("}")
    lines.append("")
    return lines


def generate_script(target_lines, num_functions, target_bytes):
    lines = []
    total_lines = 0
    body_per_func = max(1, (target_lines - num_functions * 3) // num_functions)

    for _ in range(num_functions):
        func_lines = generate_function(body_per_func)
        lines.extend(func_lines)
        total_lines += len(func_lines)

    while (
        total_lines < target_lines
        or len("\n".join(lines).encode("utf-8")) < target_bytes
    ):
        lines.append(random.choice(FUNCTION_BODIES))
        total_lines += 1

    return "\n".join(lines)


def main():
    sizes = [
        ("small", 1000, 10, 50_000),
        ("medium", 5000, 50, 500_000),
        ("large", 10000, 100, 950_000),
    ]
    for name, lines, funcs, target_bytes in sizes:
        script = generate_script(lines, funcs, target_bytes)
        filepath = os.path.join(OUTPUT_DIR, f"bench_{name}.sh")
        with open(filepath, "w") as f:
            f.write(script)
        actual_bytes = len(script.encode("utf-8"))
        actual_lines = len(script.splitlines())
        print(f"Generated {filepath}: lines={actual_lines}, bytes={actual_bytes}")


if __name__ == "__main__":
    main()