#!/usr/bin/env python3
"""
Interactive wizard for configuring the LLVM Instrumentor pass.

This script helps users create custom instrumentation configurations by:
1. Generating a default config file using opt
2. Presenting available instrumentation options interactively
3. Allowing users to enable/disable specific instrumentation opportunities
4. Saving the customized configuration to a JSON file
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
from typing import Dict, List, Any, Optional, Tuple


class InstrumentorConfigWizard:
    def __init__(self, opt_path: str = None):
        """Initialize the wizard with the path to opt."""
        self.opt_path = opt_path or self.find_opt()
        self.config = {}
        self.enabled_opportunities = set()
        self.same_pre_post = True
        self.navigation_stack = []

    def find_opt(self) -> str:
        """Find the opt binary in the build directory."""
        # Try common locations relative to this script
        script_dir = os.path.dirname(os.path.abspath(__file__))
        repo_root = os.path.dirname(os.path.dirname(script_dir))

        # Check build/bin/opt
        opt_candidates = [
            os.path.join(repo_root, "build", "bin", "opt"),
            os.path.join(repo_root, "build", "Debug", "bin", "opt"),
            os.path.join(repo_root, "build", "Release", "bin", "opt"),
            "opt",  # Try system PATH
        ]

        for candidate in opt_candidates:
            if os.path.exists(candidate):
                return candidate
            # Check if it's in PATH
            try:
                subprocess.run(
                    [candidate, "--version"], capture_output=True, check=True, timeout=5
                )
                return candidate
            except (
                subprocess.CalledProcessError,
                FileNotFoundError,
                subprocess.TimeoutExpired,
            ):
                continue

        raise FileNotFoundError(
            "Could not find 'opt' binary. Please specify the path using --opt-path"
        )

    def generate_default_config(self) -> Dict[str, Any]:
        """Generate a default configuration by running opt."""
        print(f"Generating default configuration using: {self.opt_path}")

        # Create a minimal LLVM IR module to trigger config generation
        minimal_ir = """
define i32 @main() {
  %1 = alloca i32
  store i32 0, ptr %1
  %2 = load i32, ptr %1
  ret i32 %2
}
"""

        with tempfile.TemporaryDirectory() as tmpdir:
            ir_file = os.path.join(tmpdir, "input.ll")
            config_file = os.path.join(tmpdir, "config.json")

            # Write minimal IR
            with open(ir_file, "w") as f:
                f.write(minimal_ir)

            # Run opt with instrumentor to generate config
            try:
                cmd = [
                    self.opt_path,
                    "-passes=instrumentor",
                    f"-instrumentor-write-config-file={config_file}",
                    "-disable-output",
                    ir_file,
                ]

                result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)

                if result.returncode != 0:
                    print(
                        f"Warning: opt returned non-zero exit code: {result.returncode}"
                    )
                    if result.stderr:
                        print(f"stderr: {result.stderr}")

                # Read the generated config
                if not os.path.exists(config_file):
                    raise FileNotFoundError(
                        f"Config file was not generated at {config_file}"
                    )

                with open(config_file, "r") as f:
                    config = json.load(f)

                print("✓ Default configuration generated successfully\n")
                return config

            except subprocess.TimeoutExpired:
                raise RuntimeError("opt command timed out")
            except Exception as e:
                raise RuntimeError(f"Failed to generate config: {e}")

    def clear_screen(self):
        """Clear the terminal screen."""
        os.system("clear" if os.name != "nt" else "cls")

    def print_section_header(self, title: str):
        """Print a formatted section header."""
        print("\n" + "=" * 70)
        print(f"  {title}")
        print("=" * 70)

    def print_option(self, index: int, name: str, description: str, enabled: bool):
        """Print a formatted option."""
        status = "[X]" if enabled else "[ ]"
        if description:
            print(f"  {index:2d}. {status} {name:30s} - {description}")
        else:
            print(f"  {index:2d}. {status} {name}")

    def get_user_choice(
        self, prompt: str, valid_choices: List[str] = None, allow_back: bool = True
    ) -> Optional[str]:
        """Get user input with validation."""
        while True:
            try:
                nav_hint = " (b=back, q=quit)" if allow_back else " (q=quit)"
                choice = input(prompt + nav_hint + ": ").strip().lower()

                if choice == "q":
                    confirm = input("Really quit? (y/n): ").strip().lower()
                    if confirm == "y":
                        print("\nWizard cancelled by user.")
                        sys.exit(0)
                    continue

                if choice == "b" and allow_back:
                    return "BACK"

                if not choice:
                    return ""

                if valid_choices:
                    if choice in valid_choices:
                        return choice
                    print(f"Please enter one of: {', '.join(valid_choices)}")
                else:
                    return choice

            except KeyboardInterrupt:
                print("\n\nWizard interrupted by user.")
                sys.exit(0)

    def get_location_names(self) -> List[str]:
        """Extract all location names from the config, excluding 'configuration'."""
        locations = []
        for key in self.config.keys():
            if key != "configuration":
                locations.append(key)
        return locations

    def get_all_opportunity_types(self) -> List[Tuple[str, str]]:
        """Extract all unique opportunity types from the config."""
        opportunities = []
        seen = set()

        for location in self.get_location_names():
            for opp_name in self.config[location].keys():
                if opp_name not in seen:
                    seen.add(opp_name)
                    # Get description from first occurrence
                    opp_config = self.config[location][opp_name]
                    desc = ""

                    # Look for a "description" property (without dot)
                    if "description" in opp_config and opp_config["description"]:
                        desc = opp_config["description"]

                    opportunities.append((opp_name, desc))

        return sorted(opportunities)

    def select_opportunities(self) -> bool:
        """Let user select which instrumentation opportunities to enable."""
        while True:
            self.clear_screen()
            self.print_section_header("Step 1: Select Instrumentation Types")

            opportunities = self.get_all_opportunity_types()

            print("\nSelect which types of instrumentation you want to configure:")
            print(
                "(You can toggle individual arguments for each type in the next steps)\n"
            )

            for idx, (opp_name, opp_desc) in enumerate(opportunities, 1):
                enabled = opp_name in self.enabled_opportunities
                self.print_option(idx, opp_name, opp_desc, enabled)

            print("\nCommands:")
            print("  - Enter numbers (space-separated) to toggle opportunities")
            print("  - 'all' to enable all, 'none' to disable all")
            print("  - Press Enter when done to continue")

            choice = self.get_user_choice("\nYour choice", allow_back=False)

            if choice == "BACK":
                continue
            elif choice == "":
                if not self.enabled_opportunities:
                    print("\n⚠ Please enable at least one instrumentation type!")
                    input("Press Enter to continue...")
                    continue
                return True
            elif choice == "all":
                self.enabled_opportunities = {opp[0] for opp in opportunities}
            elif choice == "none":
                self.enabled_opportunities.clear()
            else:
                try:
                    indices = [int(x) for x in choice.split()]
                    for idx in indices:
                        if 1 <= idx <= len(opportunities):
                            opp_name = opportunities[idx - 1][0]
                            if opp_name in self.enabled_opportunities:
                                self.enabled_opportunities.remove(opp_name)
                            else:
                                self.enabled_opportunities.add(opp_name)
                except ValueError:
                    print(
                        "\n⚠ Invalid input. Please enter numbers separated by spaces."
                    )
                    input("Press Enter to continue...")

    def configure_pre_post_mode(self) -> bool:
        """Ask if PRE and POST should have the same configuration."""
        while True:
            self.clear_screen()
            self.print_section_header("Step 2: PRE vs POST Configuration")

            print("\nInstrumentation can happen at two points:")
            print("  - PRE:  Before the instrumented operation")
            print("  - POST: After the instrumented operation")
            print("\nFor example, for a load instruction:")
            print("  - PRE:  Can inspect/modify the pointer before the load")
            print("  - POST: Can inspect/modify the loaded value after the load")

            print(
                f"\nCurrent mode: {'SAME configuration for PRE and POST' if self.same_pre_post else 'DIFFERENT configurations'}"
            )

            choice = self.get_user_choice(
                "\nUse same configuration for PRE and POST? (y/n/Enter to keep)",
                valid_choices=["y", "yes", "n", "no", ""],
            )

            if choice == "BACK":
                return False
            elif choice in ["y", "yes"]:
                self.same_pre_post = True
                return True
            elif choice in ["n", "no"]:
                self.same_pre_post = False
                return True
            elif choice == "":
                return True

    def configure_base_options(self) -> bool:
        """Configure base/global options."""
        while True:
            self.clear_screen()
            self.print_section_header("Step 3: Base Configuration")

            if "configuration" not in self.config:
                self.config["configuration"] = {}

            base_config = self.config["configuration"]

            # Collect all configuration options (excluding descriptions and runtime_stubs_file)
            config_options = []
            for key in sorted(base_config.keys()):
                if not key.endswith(".description") and key != "runtime_stubs_file":
                    config_options.append(key)

            # Display current settings
            print("\nCurrent settings:")
            for idx, key in enumerate(config_options, 1):
                value = base_config.get(key)
                # Format value display
                if value == "":
                    display_value = "(none)"
                else:
                    display_value = value
                print(f"  {idx}. {key:30s} {display_value}")

            print("\nEnter option number to modify, or press Enter to continue")
            choice = self.get_user_choice("Option")

            if choice == "BACK":
                return False
            elif choice == "":
                return True
            else:
                # Try to parse as a number
                try:
                    idx = int(choice)
                    if 1 <= idx <= len(config_options):
                        key = config_options[idx - 1]
                        current_value = base_config[key]
                        description = base_config.get(f"{key}.description", "")

                        # Show description if available
                        if description:
                            print(f"\n{description}")

                        # Handle different types
                        if isinstance(current_value, bool):
                            new_val = self.get_user_choice(
                                f"Set {key}? (y/n)",
                                ["y", "n"],
                                allow_back=False,
                            )
                            if new_val:
                                base_config[key] = new_val == "y"
                        elif isinstance(current_value, str):
                            # Special handling for filter properties
                            if "filter" in key.lower():
                                print(
                                    f"Note: '{key}' requires a valid regular expression"
                                )
                            new_val = input(
                                f"Enter new value for {key} (current: '{current_value}'): "
                            ).strip()
                            # Allow setting to empty string
                            base_config[key] = new_val
                        else:
                            print(f"Unsupported type for {key}: {type(current_value)}")
                            input("Press Enter to continue...")
                    else:
                        print(f"\n⚠ Invalid option number: {idx}")
                        input("Press Enter to continue...")
                except ValueError:
                    print("\n⚠ Please enter a valid number")
                    input("Press Enter to continue...")

    def configure_opportunity_args(
        self, opp_name: str, location: str, step_prefix: str = "Step 4"
    ) -> bool:
        """Configure arguments for a specific opportunity at a location."""
        while True:
            self.clear_screen()
            location_desc = "PRE (before)" if "pre" in location else "POST (after)"
            self.print_section_header(
                f"{step_prefix}: Configure {opp_name} - {location_desc}"
            )

            if location not in self.config or opp_name not in self.config[location]:
                print(f"\n⚠ {opp_name} not found in {location}")
                input("Press Enter to continue...")
                return True

            opp_config = self.config[location][opp_name]

            # Show enable/disable status
            enabled = opp_config.get("enabled", False)
            print(f"\nInstrumentation: {'ENABLED ✓' if enabled else 'DISABLED ✗'}")

            # Collect arguments
            args = []
            for key, value in sorted(opp_config.items()):
                if (
                    key == "enabled"
                    or key.endswith(".description")
                    or key.endswith(".replace")
                ):
                    continue
                desc = opp_config.get(f"{key}.description", "No description")
                can_replace = f"{key}.replace" in opp_config
                replace_enabled = (
                    opp_config.get(f"{key}.replace", False) if can_replace else False
                )
                args.append((key, value, desc, can_replace, replace_enabled))

            if args:
                print("\nAvailable arguments:")
                for idx, (
                    arg_name,
                    arg_enabled,
                    arg_desc,
                    can_replace,
                    replace_enabled,
                ) in enumerate(args, 1):
                    status = "[X]" if arg_enabled else "[ ]"
                    if can_replace:
                        replace_status = "REPLACE" if replace_enabled else "observe"
                        replace_mark = f" [replaceable: {replace_status}]"
                    else:
                        replace_mark = ""
                    print(
                        f"  {idx:2d}. {status} {arg_name:25s} - {arg_desc}{replace_mark}"
                    )

            print("\nCommands:")
            print("  - 'e' to toggle enabled/disabled")
            print("  - Enter numbers (space-separated) to toggle arguments")
            print(
                "  - 'r <num>' to toggle replacement for replaceable argument (e.g., 'r 1')"
            )
            print("  - 'all' to enable all args, 'none' to disable all args")
            print("  - Press Enter when done")

            choice = self.get_user_choice("\nYour choice")

            if choice == "BACK":
                return False
            elif choice == "":
                return True
            elif choice == "e":
                opp_config["enabled"] = not opp_config["enabled"]
            elif choice == "all":
                for arg_name, _, _, _, _ in args:
                    opp_config[arg_name] = True
            elif choice == "none":
                for arg_name, _, _, _, _ in args:
                    opp_config[arg_name] = False
            elif choice.startswith("r "):
                # Toggle replacement
                try:
                    parts = choice.split()
                    if len(parts) == 2:
                        idx = int(parts[1])
                        if 1 <= idx <= len(args):
                            arg_name, _, _, can_replace, _ = args[idx - 1]
                            if can_replace:
                                replace_key = f"{arg_name}.replace"
                                opp_config[replace_key] = not opp_config.get(
                                    replace_key, False
                                )
                            else:
                                print(f"\n⚠ Argument '{arg_name}' is not replaceable.")
                                input("Press Enter to continue...")
                        else:
                            print(f"\n⚠ Invalid argument number: {idx}")
                            input("Press Enter to continue...")
                    else:
                        print("\n⚠ Usage: r <number>")
                        input("Press Enter to continue...")
                except ValueError:
                    print("\n⚠ Invalid input for replacement toggle.")
                    input("Press Enter to continue...")
            else:
                try:
                    indices = [int(x) for x in choice.split()]
                    for idx in indices:
                        if 1 <= idx <= len(args):
                            arg_name = args[idx - 1][0]
                            opp_config[arg_name] = not opp_config[arg_name]
                except ValueError:
                    print("\n⚠ Invalid input.")
                    input("Press Enter to continue...")

    def configure_locations(self) -> bool:
        """Configure all enabled opportunities for PRE and optionally POST."""
        # First, disable all opportunities that are not in enabled_opportunities
        for location in self.get_location_names():
            for opp_name, opp_config in self.config[location].items():
                if opp_name not in self.enabled_opportunities:
                    opp_config["enabled"] = False

        # Get all location names and separate into PRE and POST
        all_locations = self.get_location_names()
        pre_locations = [loc for loc in all_locations if loc.endswith("_pre")]
        post_locations = [loc for loc in all_locations if loc.endswith("_post")]

        # Configure PRE locations
        step_num = 4
        for idx, opp_name in enumerate(sorted(self.enabled_opportunities), 1):
            # Try to find the opportunity in any PRE location
            location = None
            for loc in pre_locations:
                if opp_name in self.config[loc]:
                    location = loc
                    break

            if location:
                if not self.configure_opportunity_args(opp_name, location):
                    return False

        # If same config, copy PRE to POST
        if self.same_pre_post:
            for opp_name in self.enabled_opportunities:
                # Copy from PRE to POST for each location pair
                for pre_loc in pre_locations:
                    if opp_name not in self.config[pre_loc]:
                        continue

                    # Find corresponding POST location (e.g., function_pre -> function_post)
                    base_name = pre_loc[:-4]  # Remove "_pre"
                    post_loc = base_name + "_post"

                    if post_loc in self.config and opp_name in self.config[post_loc]:
                        pre_config = self.config[pre_loc][opp_name]
                        post_config = self.config[post_loc][opp_name]
                        # Copy enabled and argument settings
                        post_config["enabled"] = pre_config.get("enabled", False)
                        for key in pre_config:
                            if not key.endswith(".description") and key != "enabled":
                                if key in post_config:
                                    post_config[key] = pre_config[key]
        else:
            # Configure POST locations separately
            for opp_name in sorted(self.enabled_opportunities):
                # Try to find the opportunity in any POST location
                location = None
                for loc in post_locations:
                    if opp_name in self.config[loc]:
                        location = loc
                        break

                if location:
                    if not self.configure_opportunity_args(opp_name, location):
                        return False

        return True

    def generate_runtime_stubs(self, config_path: str, stub_path: str) -> bool:
        """Generate runtime stub files (both .c/.cpp and .h) using the configuration."""
        print(f"\nGenerating runtime stubs using: {self.opt_path}")

        # Create a minimal LLVM IR module
        minimal_ir = """
define i32 @main() {
  %1 = alloca i32
  store i32 0, ptr %1
  %2 = load i32, ptr %1
  ret i32 %2
}
"""

        # Calculate header path
        stub_base = stub_path.rsplit(".", 1)[0] if "." in stub_path else stub_path
        stub_header_path = stub_base + ".h"

        with tempfile.TemporaryDirectory() as tmpdir:
            ir_file = os.path.join(tmpdir, "input.ll")
            temp_config = os.path.join(tmpdir, "temp_config.json")

            # Write minimal IR
            with open(ir_file, "w") as f:
                f.write(minimal_ir)

            # Create a temporary config with stub file set
            temp_cfg = self.config.copy()
            if "configuration" not in temp_cfg:
                temp_cfg["configuration"] = {}
            temp_cfg["configuration"]["runtime_stubs_file"] = stub_path

            with open(temp_config, "w") as f:
                json.dump(temp_cfg, f, indent=2)

            # Run opt with instrumentor to generate stubs
            try:
                cmd = [
                    self.opt_path,
                    "-passes=instrumentor",
                    f"-instrumentor-read-config-files={temp_config}",
                    "-disable-output",
                    ir_file,
                ]

                result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)

                if result.returncode != 0:
                    print(
                        f"Warning: opt returned non-zero exit code: {result.returncode}"
                    )
                    if result.stderr:
                        print(f"stderr: {result.stderr}")

                # Check if both stub files were generated
                stub_exists = os.path.exists(stub_path)
                header_exists = os.path.exists(stub_header_path)

                if stub_exists and header_exists:
                    print(f"✓ Runtime stub files generated:")
                    print(f"  {stub_path}")
                    print(f"  {stub_header_path}")
                    return True
                elif stub_exists:
                    print(f"✓ Stub file generated: {stub_path}")
                    print(f"⚠ Header file was not generated: {stub_header_path}")
                    return True
                else:
                    print(f"✗ Stub files were not generated")
                    return False

            except subprocess.TimeoutExpired:
                print("✗ opt command timed out")
                return False
            except Exception as e:
                print(f"✗ Failed to generate stubs: {e}")
                return False

    def review_and_save(self, output_path: str) -> bool:
        """Review configuration and save."""
        stub_path = None

        while True:
            self.clear_screen()
            self.print_section_header("Step 5: Review and Save")

            print("\nEnabled instrumentation types:")
            for opp in sorted(self.enabled_opportunities):
                print(f"  ✓ {opp}")

            print(
                f"\nPRE/POST mode: {'Same configuration' if self.same_pre_post else 'Different configurations'}"
            )
            print(
                f"Runtime prefix: {self.config.get('configuration', {}).get('runtime_prefix', '__instrumentor_')}"
            )

            # Extract directory and filenames
            output_dir = os.path.dirname(os.path.abspath(output_path))
            output_filename = os.path.basename(output_path)

            if stub_path:
                stub_dir = os.path.dirname(os.path.abspath(stub_path))
                stub_filename = os.path.basename(stub_path)
                # Generate header filename by replacing extension with .h
                stub_base = (
                    stub_path.rsplit(".", 1)[0] if "." in stub_path else stub_path
                )
                stub_header_path = stub_base + ".h"
                stub_header_filename = os.path.basename(stub_header_path)
            else:
                default_stub = output_path.rsplit(".", 1)[0] + "_stubs.c"
                stub_dir = os.path.dirname(os.path.abspath(default_stub))
                stub_filename = os.path.basename(default_stub)
                stub_base = default_stub.rsplit(".", 1)[0]
                stub_header_path = stub_base + ".h"
                stub_header_filename = os.path.basename(stub_header_path)

            print("\n" + "-" * 70)
            print("Output Location:")
            print("-" * 70)
            print(f"Directory:          {output_dir}")
            print(f"Config filename:    {output_filename}")
            print(
                f"Stub filenames:     {stub_filename}{' (not generated yet)' if not stub_path else ''}"
            )
            print(
                f"                    {stub_header_filename}{' (not generated yet)' if not stub_path else ''}"
            )
            print(f"\nFull paths:")
            print(f"  Config: {output_path}")
            if stub_path:
                print(f"  Stubs:  {stub_path}")
                print(f"          {stub_header_path}")
            else:
                default_stub = output_path.rsplit(".", 1)[0] + "_stubs.c"
                print(f"  Stubs:  {default_stub} (when generated)")
                print(f"          {stub_header_path} (when generated)")

            print("\nCommands:")
            print("  - 's' to save configuration and finish")
            print("  - 'g' to generate runtime stub file")
            print("  - 'd' to change output directory")
            print("  - 'c' to change configuration filename")
            print("  - 't' to change stub filename")
            print("  - 'b' to go back and modify settings")

            choice = self.get_user_choice(
                "\nYour choice", valid_choices=["s", "g", "d", "c", "t", "b", ""]
            )

            if choice == "BACK" or choice == "b":
                return False
            elif choice == "s":
                try:
                    # Create output directory if it doesn't exist
                    output_dir = os.path.dirname(os.path.abspath(output_path))
                    if output_dir and not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                        print(f"✓ Created directory: {output_dir}")

                    # Remove runtime_stubs_file from config before saving
                    config_to_save = json.loads(json.dumps(self.config))
                    if "configuration" in config_to_save:
                        config_to_save["configuration"].pop("runtime_stubs_file", None)
                        config_to_save["configuration"].pop(
                            "runtime_stubs_file.description", None
                        )

                    with open(output_path, "w") as f:
                        json.dump(config_to_save, f, indent=2)
                    print(f"\n✓ Configuration saved to: {output_path}")

                    return True
                except Exception as e:
                    print(f"\n✗ Failed to save configuration: {e}")
                    input("Press Enter to continue...")
            elif choice == "d":
                # Change output directory
                current_dir = os.path.dirname(os.path.abspath(output_path))
                print(f"\nCurrent directory: {current_dir}")
                new_dir = input("Enter new output directory: ").strip()
                if new_dir:
                    # Expand user home directory if needed
                    new_dir = os.path.expanduser(new_dir)
                    new_dir = os.path.abspath(new_dir)

                    # Update output_path with new directory
                    output_filename = os.path.basename(output_path)
                    output_path = os.path.join(new_dir, output_filename)

                    # Update stub_path with new directory if it was set
                    if stub_path:
                        stub_filename = os.path.basename(stub_path)
                        stub_path = os.path.join(new_dir, stub_filename)

                    print(f"✓ Output directory updated to: {new_dir}")
                    print(f"  Config will be saved to: {output_path}")
                    if stub_path:
                        print(f"  Stubs will be saved to: {stub_path}")
                else:
                    print("Directory unchanged")
                input("Press Enter to continue...")
            elif choice == "c":
                # Change configuration filename only
                current_filename = os.path.basename(output_path)
                current_dir = os.path.dirname(os.path.abspath(output_path))
                print(f"\nCurrent filename: {current_filename}")
                new_filename = input("Enter new configuration filename: ").strip()
                if new_filename:
                    output_path = os.path.join(current_dir, new_filename)
                    print(f"✓ Configuration filename updated to: {new_filename}")
                    print(f"  Full path: {output_path}")
                else:
                    print("Filename unchanged")
                input("Press Enter to continue...")
            elif choice == "t":
                # Change stub filename only
                current_dir = os.path.dirname(os.path.abspath(output_path))
                if stub_path:
                    current_stub_filename = os.path.basename(stub_path)
                else:
                    default_stub = output_path.rsplit(".", 1)[0] + "_stubs.c"
                    current_stub_filename = os.path.basename(default_stub)

                print(f"\nCurrent stub filename: {current_stub_filename}")
                new_stub_filename = input("Enter new stub filename: ").strip()
                if new_stub_filename:
                    stub_path = os.path.join(current_dir, new_stub_filename)
                    print(f"✓ Stub filename updated to: {new_stub_filename}")
                    print(f"  Full path: {stub_path}")
                else:
                    print("Filename unchanged")
                input("Press Enter to continue...")
            elif choice == "g":
                print("\nGenerate runtime stub files")
                print("This creates C/C++ files with stub implementations of the")
                print(
                    "instrumentation runtime functions that you can use as a template."
                )
                print(
                    "Note: Both a .c/.cpp file and a .h header file will be generated."
                )

                # Use current stub_path or ask for one
                if not stub_path:
                    default_stub = output_path.rsplit(".", 1)[0] + "_stubs.c"
                    stub_input = input(
                        f"\nStub file path (default: {default_stub}): "
                    ).strip()
                    stub_path = stub_input if stub_input else default_stub

                # Calculate header path
                stub_base = (
                    stub_path.rsplit(".", 1)[0] if "." in stub_path else stub_path
                )
                stub_header_path = stub_base + ".h"

                # Generate stubs immediately
                print(f"\nGenerating stub files:")
                print(f"  {stub_path}")
                print(f"  {stub_header_path}")

                # First save the config to a temporary location if not already saved
                with tempfile.TemporaryDirectory() as tmpdir:
                    temp_config_path = os.path.join(tmpdir, "temp_config.json")
                    config_to_save = json.loads(json.dumps(self.config))
                    if "configuration" in config_to_save:
                        config_to_save["configuration"].pop("runtime_stubs_file", None)
                        config_to_save["configuration"].pop(
                            "runtime_stubs_file.description", None
                        )

                    with open(temp_config_path, "w") as f:
                        json.dump(config_to_save, f, indent=2)

                    if self.generate_runtime_stubs(temp_config_path, stub_path):
                        print(f"✓ Stub files generated successfully")
                        print(f"  Created: {stub_path}")
                        print(f"  Created: {stub_header_path}")
                    else:
                        print(f"✗ Failed to generate stub files")

                input("Press Enter to continue...")
            elif choice == "":
                continue

    def run_interactive(self, output_path: str):
        """Run the interactive configuration wizard."""
        self.clear_screen()
        print("=" * 70)
        print("  LLVM Instrumentor Configuration Wizard")
        print("=" * 70)
        print(
            "\nThis wizard will help you create a custom instrumentation configuration."
        )
        print("You can enable/disable instrumentation opportunities and configure")
        print("what information is passed to the runtime functions.")
        print("\nNavigation: Use 'b' to go back, 'q' to quit at any prompt.")
        input("\nPress Enter to continue...")

        # Generate or load config
        try:
            self.config = self.generate_default_config()
        except Exception as e:
            print(f"Error: {e}")
            return False

        # State machine for navigation
        state = 0
        while True:
            if state == 0:  # Select opportunities
                if self.select_opportunities():
                    state = 1
            elif state == 1:  # PRE/POST mode
                if self.configure_pre_post_mode():
                    state = 2
                else:
                    state = 0
            elif state == 2:  # Base configuration
                if self.configure_base_options():
                    state = 3
                else:
                    state = 1
            elif state == 3:  # Configure locations
                if self.configure_locations():
                    state = 4
                else:
                    state = 2
            elif state == 4:  # Review and save
                if self.review_and_save(output_path):
                    return True
                else:
                    state = 3


def main():
    parser = argparse.ArgumentParser(
        description="Interactive wizard for configuring LLVM Instrumentor pass",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Interactive mode (recommended)
  %(prog)s

  # Specify custom output location
  %(prog)s -o my_config.json

  # Use specific opt binary
  %(prog)s --opt-path /path/to/opt

  # Load existing config and modify it
  %(prog)s --input existing_config.json -o modified_config.json
        """,
    )

    parser.add_argument(
        "-o",
        "--output",
        default="instrumentor_config.json",
        help="Output configuration file (default: instrumentor_config.json)",
    )

    parser.add_argument(
        "--opt-path", help="Path to the opt binary (default: auto-detect)"
    )

    parser.add_argument(
        "--input", help="Load and modify an existing configuration file"
    )

    args = parser.parse_args()

    try:
        wizard = InstrumentorConfigWizard(opt_path=args.opt_path)

        # Load existing config if provided
        if args.input:
            print(f"Loading existing configuration from: {args.input}")
            with open(args.input, "r") as f:
                wizard.config = json.load(f)
            print("✓ Configuration loaded\n")
            # Extract enabled opportunities from loaded config
            for location in wizard.get_location_names():
                for opp_name, opp_config in wizard.config[location].items():
                    if opp_config.get("enabled", False):
                        wizard.enabled_opportunities.add(opp_name)

        success = wizard.run_interactive(args.output)

        if success:
            print("\n" + "=" * 70)
            print("Configuration complete!")
            print("=" * 70)
            print(f"\nTo use this configuration with opt:")
            print(f"  opt -passes=instrumentor \\")
            print(f"      -instrumentor-read-config-files={args.output} \\")
            print(f"      input.ll -S -o output.ll")
            print(f"\nTo use with clang:")
            print(f"  clang -mllvm -enable-instrumentor \\")
            print(f"        -mllvm -instrumentor-read-config-files={args.output} \\")
            print(f"        input.c -o output")
            return 0
        else:
            return 1

    except Exception as e:
        print(f"\nFatal error: {e}", file=sys.stderr)
        import traceback

        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())