import argparse
import random

# --- Data from the Cyberpunk Red Core Rulebook ---

# Difficulty Values (DV) for Architectures
DIFFICULTY_DV = {
    "basic": 6,
    "standard": 8,
    "uncommon": 10,
    "advanced": 12,
}

# Tables for populating floors
LOBBY_TABLE = [
    "Password", "Password", "Password", "Password",
    "Control Node (1 Device)", "Control Node (1 Device)",
    "File", "File",
    "Wisp (Anti-Personnel Black ICE)", "Raven (Anti-Personnel Black ICE)"
]

BASIC_ARCHITECTURE_TABLE = [
    "Password", "Password", "Password", "Password",
    "Control Node (1d6/2 Devices)", "Control Node (1d6/2 Devices)",
    "File", "File",
    "Wisp (Anti-Personnel Black ICE)", "Asp (Anti-Personnel Black ICE)"
]

STANDARD_ARCHITECTURE_TABLE = [
    "Password", "Password", "Password",
    "Control Node (1d6 Devices)", "Control Node (1d6 Devices)",
    "File", "File",
    "Hellhound (Anti-Personnel Black ICE)", "Killer (Anti-Program Black ICE)",
    "Skunk (Anti-Personnel Black ICE)"
]

UNCOMMON_ARCHITECTURE_TABLE = [
    "Password", "Password",
    "Control Node (1d10 Devices)", "Control Node (1d10 Devices)",
    "File", "File",
    "Liche (Anti-Personnel Black ICE)", "Sabertooth (Anti-Program Black ICE)",
    "Giant (Anti-Personnel Black ICE)", "Scorpion (Anti-Personnel Black ICE)"
]

ADVANCED_ARCHITECTURE_TABLE = [
    "Password", "Password",
    "Control Node (1d10 Devices)", "Control Node (1d10 Devices)",
    "File", "File",
    "Kraken (Anti-Personnel Black ICE)", "Dragon (Anti-Program Black ICE)",
    "Giant (Anti-Personnel Black ICE)", "Sabertooth (Anti-Program Black ICE)"
]

# --- Helper Functions ---

def roll(dice_string):
    """Rolls dice based on a string like '1d6' or '3d6'."""
    if not 'd' in dice_string:
        return int(dice_string)
    num_dice, die_type = map(int, dice_string.split('d'))
    return sum(random.randint(1, die_type) for _ in range(num_dice))

def get_table_for_difficulty(difficulty):
    """Returns the correct floor content table for a given difficulty."""
    if difficulty == "basic":
        return BASIC_ARCHITECTURE_TABLE
    elif difficulty == "standard":
        return STANDARD_ARCHITECTURE_TABLE
    elif difficulty == "uncommon":
        return UNCOMMON_ARCHITECTURE_TABLE
    elif difficulty == "advanced":
        return ADVANCED_ARCHITECTURE_TABLE
    else:
        # Default to basic if something goes wrong
        return BASIC_ARCHITECTURE_TABLE

# --- Core Generator Logic ---

def generate_branch(difficulty, num_floors, branch_level):
    """Generates a single branch of the architecture."""
    branch = []
    table = get_table_for_difficulty(difficulty)
    for i in range(num_floors):
        floor_content = random.choice(table)
        branch.append(f"  {'  ' * branch_level}L Floor {i+1}: {floor_content}")
    return branch

def generate_architecture(difficulty, total_floors, num_branches):
    """Generates the entire NET Architecture layout."""
    architecture = []
    dv = DIFFICULTY_DV.get(difficulty, 6)

    # --- Header ---
    architecture.append("=" * 40)
    architecture.append(f"NET ARCHITECTURE - {difficulty.upper()}")
    architecture.append(f"Baseline DV: {dv}")
    architecture.append("=" * 40)
    architecture.append("")

    # --- Lobby Floors ---
    architecture.append("Floor 1 (Lobby): " + random.choice(LOBBY_TABLE))
    architecture.append("Floor 2 (Lobby): " + random.choice(LOBBY_TABLE))

    remaining_floors = total_floors - 2
    floors_per_branch = 0
    if num_branches > 0:
        floors_per_branch = remaining_floors // (num_branches + 1)

    main_shaft_floors = remaining_floors - (floors_per_branch * num_branches)

    # --- Main Shaft ---
    table = get_table_for_difficulty(difficulty)
    for i in range(main_shaft_floors):
        floor_content = random.choice(table)
        architecture.append(f"Floor {i+3}: {floor_content}")

    # --- Branches ---
    if num_branches > 0:
        for i in range(num_branches):
            branch_start_floor = random.randint(3, total_floors - floors_per_branch + 1)
            branch = generate_branch(difficulty, floors_per_branch, 1)
            architecture.insert(branch_start_floor + i, f"-> Branch {i+1} (starts after Floor {branch_start_floor -1})")
            for floor in reversed(branch):
                architecture.insert(branch_start_floor + i + 1, floor)


    return "\n".join(architecture)

# --- Main Execution ---

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate a NET Architecture for Cyberpunk Red.",
        formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument(
        "-d", "--difficulty",
        choices=["basic", "standard", "uncommon", "advanced"],
        default="standard",
        help="The difficulty rating of the architecture. Defaults to 'standard'."
    )
    parser.add_argument(
        "-f", "--floors",
        type=int,
        default=roll("3d6"),
        help="The total number of floors. Defaults to a random 3d6 roll."
    )
    parser.add_argument(
        "-b", "--branches",
        type=int,
        default=0,
        help="The number of branches in the architecture. Defaults to 0."
    )

    args = parser.parse_args()

    # --- Generate and Print ---
    layout = generate_architecture(args.difficulty, args.floors, args.branches)
    print(layout)