import argparse
import random
import string
from pathlib import Path

from tqdm.auto import tqdm


def some_str(N: int = None, with_space=0) -> str:
    N = N or random.randint(10, 100)
    candidates = string.ascii_lowercase + string.digits
    if with_space:
        candidates += " " * int(with_space)
    return "".join(random.choices(candidates, k=N))
    return str(uuid.uuid4())


def letters_only_str(N: int = None) -> str:
    N = N or random.randint(10, 100)
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=N))
    return str(uuid.uuid4())


def generate_heading():
    has_id = random.choice([True, False])
    level = random.randint(1, 6)

    s = "#" * level + " " * random.randint(1, 20) + some_str()
    if has_id:
        id_spacing = random.randint(1, 100)

        s += " " * id_spacing + "  :" + some_str(random.randint(1, 4))

    return s


def generate_params():
    n_params = random.choice([0, 15])
    s = ""
    for param in range(n_params):
        is_kwarg = random.choice([True, False])
        value_type = random.choice(["str", "int", "float"])
        value = "v"
        if value_type == "str":
            value = some_str(random.randint(1, 20))
            if random.choice([True, False]):
                value = f'"{value}"'
        elif value_type == "int":
            value = random.randint(0, 1000)
        elif value_type == "float":
            value = random.randint(0, 1000)
            value += 1 / random.randint(1, 1000)

        s += " " * random.randint(1, 5)
        if is_kwarg:
            s += f"{some_str(random.randint(1, 10))}={value}"
        else:
            s += str(value)
    return s


def generate_metaline():
    s = "."
    has_inline_text = random.choice([True, False])
    has_params = random.choice([True, False])
    blockType = letters_only_str(random.randint(1, 5))

    s += blockType + " " * random.randint(1, 10)

    if has_params:
        s += generate_params()
        if has_inline_text:
            s += " .. "
    else:
        s += some_str(random.randint(10, 100), with_space=True)

    if has_inline_text:
        t = " " + some_str(with_space=3)
        if random.choice([True, False]):
            t = f'"{t}"'
        s += t

    s += " " * random.randint(0, 10)

    return s


def generate_segment_content():
    n_lines = random.randint(0, 100)
    s = ""
    for _ in range(n_lines):
        empty = random.choice([True, False, False, False, False, False])
        if empty:
            s += "\n"
            continue
        s += (
            some_str(random.randint(0, 100), with_space=10)
            + " " * random.randint(0, 10)
            + "\n"
        )
    return s


def generate_segment_line():
    s = ".."
    s += " " * random.randint(1, 10)
    if random.choice([True, False]):
        s += generate_params()
        s += " " * random.randint(1, 10)
    return s


class Stop(Exception): ...


KB = 1024
MB = 1024 * KB
GB = 1024 * MB


def parse_size(size_str: str) -> int:
    """Parse a size string like '8GB', '512MB', '1024KB' into bytes."""
    size_str = size_str.strip().upper()
    if size_str.endswith("GB"):
        return int(float(size_str[:-2]) * GB)
    elif size_str.endswith("MB"):
        return int(float(size_str[:-2]) * MB)
    elif size_str.endswith("KB"):
        return int(float(size_str[:-2]) * KB)
    elif size_str.endswith("B"):
        return int(size_str[:-1])
    else:
        return int(size_str)


def main():
    parser = argparse.ArgumentParser(
        description="Generate a large wasot/markdown file for testing"
    )
    parser.add_argument(
        "-o", "--output",
        type=Path,
        default=Path("tmp/big_input.md"),
        help="Output file path (default: tmp/big_input.md)"
    )
    parser.add_argument(
        "-s", "--size",
        type=str,
        default="8GB",
        help="Target file size, e.g. 8GB, 512MB, 1024KB (default: 8GB)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility (default: 42)"
    )
    args = parser.parse_args()

    max_bytes = parse_size(args.size)
    output_path = args.output

    # Create parent directories if they don't exist
    output_path.parent.mkdir(parents=True, exist_ok=True)

    random.seed(args.seed)

    total_bytes = 0

    with open(output_path, "w") as file:
        with tqdm(total=max_bytes, unit="B", unit_scale=True) as pbar:

            def p(s):
                nonlocal total_bytes
                length = len(s.encode("utf-8"))
                pbar.update(length)
                total_bytes += length
                if total_bytes > max_bytes:
                    raise Stop()
                file.write(s + "\n")

            try:
                p("\n" * random.randint(1, 10))
                while True:
                    p(generate_heading())
                    p("\n" * random.randint(1, 5))

                    n_blocks = random.randint(1, 30)
                    for _ in range(n_blocks):
                        p(generate_metaline())
                        if random.choice([True, False]):
                            p(generate_segment_content())
                        for _ in range(random.randint(0, 4)):
                            p(generate_segment_line())
                            p(generate_segment_content())
            except Stop:
                pass

    print(f"Generated {output_path} ({total_bytes:,} bytes)")


if __name__ == "__main__":
    main()
