#!/usr/bin/env python3
"""
PlatformIO Helper - An all-in-one interface for common PlatformIO commands.

Features:
- Interactive menu mode (when run without args)
- CLI mode for quick access
- fzf integration for fuzzy selection
- Custom preset/workflow support

Usage:
    ./scripts/pio_helper.py              # Interactive menu
    ./scripts/pio_helper.py upload       # Quick upload
    ./scripts/pio_helper.py upload -e debug_memory
    ./scripts/pio_helper.py clean+upload+debug-monitor -l
    ./scripts/pio_helper.py --preset dev
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Optional

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────

SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent
DEBUGGING_MONITOR = SCRIPT_DIR / "debugging_monitor.py"
DEFAULT_LOG_PATH = PROJECT_ROOT / ".cursor" / "debug.log"
PRESETS_DIR = Path.home() / ".config" / "pio-helper"
PRESETS_FILE = PRESETS_DIR / "presets.json"

# Available environments from platformio.ini
ENVIRONMENTS = ["default", "debug_memory", "gh_release"]
DEFAULT_ENV = "default"

# Action definitions
ACTIONS = {
    "build": {
        "description": "Build without upload",
        "command": ["pio", "run"],
        "needs_env": True,
    },
    "upload": {
        "description": "Build and upload to device",
        "command": ["pio", "run", "-t", "upload"],
        "needs_env": True,
    },
    "clean": {
        "description": "Clean build files",
        "command": ["pio", "run", "-t", "clean"],
        "needs_env": True,
    },
    "monitor": {
        "description": "Basic serial monitor",
        "command": ["pio", "device", "monitor"],
        "needs_env": False,
    },
    "debug-monitor": {
        "description": "Enhanced monitor with memory graphs",
        "command": [sys.executable, str(DEBUGGING_MONITOR)],
        "needs_env": False,
    },
}

# Common workflow shortcuts for the menu
MENU_WORKFLOWS = [
    {"name": "Build only", "actions": ["build"]},
    {"name": "Upload (default env)", "actions": ["upload"], "env": "default"},
    {"name": "Upload (debug_memory env)", "actions": ["upload"], "env": "debug_memory"},
    {"name": "Clean + Upload", "actions": ["clean", "upload"]},
    {"name": "Upload + Monitor", "actions": ["upload", "monitor"]},
    {"name": "Upload + Debug Monitor", "actions": ["upload", "debug-monitor"]},
    {"name": "Clean + Upload + Debug Monitor", "actions": ["clean", "upload", "debug-monitor"]},
    {"name": "Monitor only", "actions": ["monitor"]},
    {"name": "Debug Monitor only", "actions": ["debug-monitor"]},
]

# ─────────────────────────────────────────────────────────────────────────────
# Terminal Colors
# ─────────────────────────────────────────────────────────────────────────────

class Colors:
    """ANSI color codes for terminal output."""
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    
    RED = "\033[31m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    BLUE = "\033[34m"
    MAGENTA = "\033[35m"
    CYAN = "\033[36m"
    WHITE = "\033[37m"
    
    @classmethod
    def disable(cls):
        """Disable colors (for non-TTY output)."""
        for attr in dir(cls):
            if not attr.startswith('_') and isinstance(getattr(cls, attr), str):
                setattr(cls, attr, "")


# Disable colors if not a TTY
if not sys.stdout.isatty():
    Colors.disable()


def print_header(text: str):
    """Print a styled header."""
    print(f"\n{Colors.CYAN}{Colors.BOLD}{text}{Colors.RESET}")
    print(f"{Colors.DIM}{'─' * len(text)}{Colors.RESET}")


def print_success(text: str):
    """Print success message."""
    print(f"{Colors.GREEN}✓ {text}{Colors.RESET}")


def print_error(text: str):
    """Print error message."""
    print(f"{Colors.RED}✗ {text}{Colors.RESET}")


def print_info(text: str):
    """Print info message."""
    print(f"{Colors.BLUE}→ {text}{Colors.RESET}")


def print_warning(text: str):
    """Print warning message."""
    print(f"{Colors.YELLOW}! {text}{Colors.RESET}")


# ─────────────────────────────────────────────────────────────────────────────
# Preset Management
# ─────────────────────────────────────────────────────────────────────────────

def load_presets() -> dict:
    """Load presets from the config file."""
    if not PRESETS_FILE.exists():
        return {}
    try:
        with open(PRESETS_FILE, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError) as e:
        print_warning(f"Could not load presets: {e}")
        return {}


def save_presets(presets: dict):
    """Save presets to the config file."""
    PRESETS_DIR.mkdir(parents=True, exist_ok=True)
    with open(PRESETS_FILE, "w") as f:
        json.dump(presets, f, indent=2)


def save_preset(name: str, actions: list[str], env: str, log: bool, log_path: Optional[str] = None):
    """Save a new preset."""
    presets = load_presets()
    presets[name] = {
        "actions": actions,
        "env": env,
        "log": log,
    }
    if log_path:
        presets[name]["log_path"] = log_path
    save_presets(presets)
    print_success(f"Saved preset '{name}'")


def delete_preset(name: str) -> bool:
    """Delete a preset."""
    presets = load_presets()
    if name in presets:
        del presets[name]
        save_presets(presets)
        print_success(f"Deleted preset '{name}'")
        return True
    print_error(f"Preset '{name}' not found")
    return False


def list_presets():
    """List all saved presets."""
    presets = load_presets()
    if not presets:
        print_info("No presets saved yet.")
        print_info(f"Presets are stored in: {PRESETS_FILE}")
        return
    
    print_header("Saved Presets")
    for name, config in presets.items():
        actions_str = " + ".join(config.get("actions", []))
        env = config.get("env", DEFAULT_ENV)
        log = "with logging" if config.get("log") else ""
        print(f"  {Colors.GREEN}{name}{Colors.RESET}: {actions_str} ({env}) {log}")


# ─────────────────────────────────────────────────────────────────────────────
# fzf Integration
# ─────────────────────────────────────────────────────────────────────────────

def has_fzf() -> bool:
    """Check if fzf is available."""
    return shutil.which("fzf") is not None


def fzf_select(options: list[str], prompt: str = "Select: ", multi: bool = False) -> Optional[list[str]]:
    """
    Use fzf for selection. Returns list of selected options or None if cancelled.
    Falls back to numbered menu if fzf is not available.
    """
    if not has_fzf():
        return numbered_select(options, prompt, multi)
    
    try:
        fzf_args = ["fzf", "--prompt", prompt, "--height", "40%", "--reverse"]
        if multi:
            fzf_args.append("--multi")
        
        result = subprocess.run(
            fzf_args,
            input="\n".join(options),
            capture_output=True,
            text=True,
        )
        
        if result.returncode == 0:
            selected = result.stdout.strip().split("\n")
            return [s for s in selected if s]
        return None
    except (subprocess.SubprocessError, OSError):
        return numbered_select(options, prompt, multi)


def numbered_select(options: list[str], prompt: str = "Select: ", multi: bool = False) -> Optional[list[str]]:
    """Fallback numbered selection menu."""
    print()
    for i, opt in enumerate(options, 1):
        print(f"  {Colors.CYAN}{i:2}{Colors.RESET}. {opt}")
    print(f"  {Colors.DIM} 0. Cancel{Colors.RESET}")
    print()
    
    try:
        if multi:
            print_info("Enter numbers separated by spaces (e.g., '1 3 5')")
        choice = input(f"{prompt}").strip()
        
        if not choice or choice == "0":
            return None
        
        if multi:
            indices = [int(c) - 1 for c in choice.split()]
            return [options[i] for i in indices if 0 <= i < len(options)]
        else:
            idx = int(choice) - 1
            if 0 <= idx < len(options):
                return [options[idx]]
            return None
    except (ValueError, IndexError, KeyboardInterrupt):
        return None


# ─────────────────────────────────────────────────────────────────────────────
# Action Execution
# ─────────────────────────────────────────────────────────────────────────────

def build_command(action: str, env: str, verbose: bool = False) -> list[str]:
    """Build the command for a given action."""
    if action not in ACTIONS:
        raise ValueError(f"Unknown action: {action}")
    
    action_def = ACTIONS[action]
    cmd = list(action_def["command"])
    
    # Add environment flag if needed
    if action_def["needs_env"]:
        cmd.extend(["-e", env])
    
    # Add verbose flag
    if verbose and action in ("build", "upload", "clean"):
        cmd.append("-v")
    
    return cmd


def run_action(action: str, env: str, log: bool = False, log_path: Optional[Path] = None, 
               verbose: bool = False) -> int:
    """
    Run a single action.
    
    Returns the exit code of the command.
    """
    cmd = build_command(action, env, verbose)
    action_def = ACTIONS[action]
    
    print_info(f"Running: {action} ({action_def['description']})")
    print(f"{Colors.DIM}  $ {' '.join(cmd)}{Colors.RESET}")
    
    if log:
        actual_log_path = log_path or DEFAULT_LOG_PATH
        actual_log_path.parent.mkdir(parents=True, exist_ok=True)
        print_info(f"Logging to: {actual_log_path}")
        
        # Use tee to both display and log output
        with open(actual_log_path, "a") as log_file:
            # Write a separator to the log
            log_file.write(f"\n{'='*60}\n")
            log_file.write(f"Action: {action} | Env: {env}\n")
            log_file.write(f"Command: {' '.join(cmd)}\n")
            log_file.write(f"{'='*60}\n")
            log_file.flush()
            
            # Run with tee-like behavior
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=PROJECT_ROOT,
            )
            
            for line in iter(process.stdout.readline, ""):
                sys.stdout.write(line)
                sys.stdout.flush()
                log_file.write(line)
                log_file.flush()
            
            process.wait()
            return process.returncode
    else:
        # Run directly without logging
        result = subprocess.run(cmd, cwd=PROJECT_ROOT)
        return result.returncode


def run_chain(actions: list[str], env: str, log: bool = False, 
              log_path: Optional[Path] = None, verbose: bool = False) -> int:
    """
    Run a chain of actions sequentially.
    
    Stops on first failure and returns the exit code.
    """
    print_header(f"Running: {' + '.join(actions)}")
    print_info(f"Environment: {env}")
    if log:
        print_info(f"Logging enabled")
    print()
    
    for i, action in enumerate(actions):
        if i > 0:
            print()  # Spacing between actions
        
        exit_code = run_action(action, env, log, log_path, verbose)
        
        if exit_code != 0:
            print_error(f"Action '{action}' failed with exit code {exit_code}")
            return exit_code
        
        print_success(f"Action '{action}' completed successfully")
    
    print()
    print_success(f"All actions completed successfully!")
    return 0


# ─────────────────────────────────────────────────────────────────────────────
# Interactive Menu
# ─────────────────────────────────────────────────────────────────────────────

def select_environment() -> Optional[str]:
    """Prompt user to select an environment."""
    result = fzf_select(ENVIRONMENTS, "Select environment: ")
    return result[0] if result else None


def interactive_menu():
    """Run the interactive menu interface."""
    while True:
        print_header("PlatformIO Helper")
        
        # Build menu options
        options = []
        for wf in MENU_WORKFLOWS:
            env_info = f" [{wf.get('env', 'select env')}]" if 'env' in wf else ""
            options.append(f"{wf['name']}{env_info}")
        
        options.append("─" * 30)  # Separator
        options.append("📋 Run preset...")
        options.append("💾 Save workflow as preset...")
        options.append("📝 List presets")
        options.append("🗑️  Delete preset...")
        options.append("─" * 30)  # Separator
        options.append("⚙️  Custom: Select actions...")
        options.append("❌ Quit")
        
        # Use fzf or numbered menu
        if has_fzf():
            result = fzf_select(options, "Select workflow: ")
            if not result:
                continue
            choice = result[0]
        else:
            print()
            for i, opt in enumerate(options, 1):
                if opt.startswith("─"):
                    print(f"  {Colors.DIM}{opt}{Colors.RESET}")
                else:
                    print(f"  {Colors.CYAN}{i:2}{Colors.RESET}. {opt}")
            print()
            
            try:
                choice_input = input(f"Select (1-{len(options)}, q to quit): ").strip().lower()
                if choice_input in ("q", "quit", "exit", ""):
                    break
                idx = int(choice_input) - 1
                if 0 <= idx < len(options):
                    choice = options[idx]
                else:
                    continue
            except (ValueError, KeyboardInterrupt):
                break
        
        # Handle choice
        if choice.startswith("─") or not choice:
            continue
        elif "Quit" in choice:
            break
        elif "Run preset" in choice:
            run_preset_menu()
        elif "Save workflow" in choice:
            save_preset_menu()
        elif "List presets" in choice:
            list_presets()
            input(f"\n{Colors.DIM}Press Enter to continue...{Colors.RESET}")
        elif "Delete preset" in choice:
            delete_preset_menu()
        elif "Custom" in choice:
            custom_workflow_menu()
        else:
            # Find matching workflow
            for wf in MENU_WORKFLOWS:
                env_info = f" [{wf.get('env', 'select env')}]" if 'env' in wf else ""
                if choice == f"{wf['name']}{env_info}":
                    env = wf.get("env")
                    if env is None:
                        env = select_environment()
                        if env is None:
                            continue
                    
                    # Ask about logging
                    log = ask_yes_no("Enable logging?", default=False)
                    
                    run_chain(wf["actions"], env, log)
                    input(f"\n{Colors.DIM}Press Enter to continue...{Colors.RESET}")
                    break


def ask_yes_no(prompt: str, default: bool = False) -> bool:
    """Ask a yes/no question."""
    default_str = "Y/n" if default else "y/N"
    try:
        answer = input(f"{prompt} [{default_str}]: ").strip().lower()
        if not answer:
            return default
        return answer in ("y", "yes", "1", "true")
    except (KeyboardInterrupt, EOFError):
        return default


def run_preset_menu():
    """Menu to run a saved preset."""
    presets = load_presets()
    if not presets:
        print_info("No presets saved yet.")
        return
    
    options = list(presets.keys())
    result = fzf_select(options, "Select preset: ")
    if not result:
        return
    
    name = result[0]
    config = presets[name]
    
    actions = config.get("actions", [])
    env = config.get("env", DEFAULT_ENV)
    log = config.get("log", False)
    log_path = Path(config["log_path"]) if config.get("log_path") else None
    
    run_chain(actions, env, log, log_path)
    input(f"\n{Colors.DIM}Press Enter to continue...{Colors.RESET}")


def save_preset_menu():
    """Menu to save a new preset."""
    # Select actions
    action_names = list(ACTIONS.keys())
    print_info("Select actions to include in the preset:")
    selected = fzf_select(action_names, "Select actions (multi): ", multi=True)
    if not selected:
        return
    
    # Select environment
    env = select_environment()
    if not env:
        return
    
    # Ask about logging
    log = ask_yes_no("Enable logging by default?", default=False)
    
    # Get preset name
    try:
        name = input("Preset name: ").strip()
        if not name:
            print_error("Preset name cannot be empty")
            return
    except (KeyboardInterrupt, EOFError):
        return
    
    save_preset(name, selected, env, log)


def delete_preset_menu():
    """Menu to delete a preset."""
    presets = load_presets()
    if not presets:
        print_info("No presets to delete.")
        return
    
    options = list(presets.keys())
    result = fzf_select(options, "Select preset to delete: ")
    if not result:
        return
    
    name = result[0]
    if ask_yes_no(f"Delete preset '{name}'?", default=False):
        delete_preset(name)


def custom_workflow_menu():
    """Menu to create a custom workflow."""
    # Select actions
    action_names = list(ACTIONS.keys())
    print_info("Select actions to run (in order):")
    selected = fzf_select(action_names, "Select actions (multi): ", multi=True)
    if not selected:
        return
    
    # Select environment
    env = select_environment()
    if not env:
        return
    
    # Ask about logging
    log = ask_yes_no("Enable logging?", default=False)
    
    run_chain(selected, env, log)
    
    # Offer to save as preset
    if ask_yes_no("\nSave this as a preset?", default=False):
        try:
            name = input("Preset name: ").strip()
            if name:
                save_preset(name, selected, env, log)
        except (KeyboardInterrupt, EOFError):
            pass
    
    input(f"\n{Colors.DIM}Press Enter to continue...{Colors.RESET}")


# ─────────────────────────────────────────────────────────────────────────────
# CLI Interface
# ─────────────────────────────────────────────────────────────────────────────

def parse_actions(action_str: str) -> list[str]:
    """Parse action string like 'clean+upload+monitor' into list of actions."""
    actions = [a.strip() for a in action_str.split("+")]
    
    # Validate actions
    for action in actions:
        if action not in ACTIONS:
            print_error(f"Unknown action: {action}")
            print_info(f"Available actions: {', '.join(ACTIONS.keys())}")
            sys.exit(1)
    
    return actions


def main():
    parser = argparse.ArgumentParser(
        description="PlatformIO Helper - An all-in-one interface for common PlatformIO commands.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                           Interactive menu
  %(prog)s upload                    Upload with default env
  %(prog)s upload -e debug_memory    Upload with debug_memory env
  %(prog)s clean+upload+monitor      Chain multiple actions
  %(prog)s upload+debug-monitor -l   Upload and monitor with logging
  %(prog)s --preset dev              Run saved preset 'dev'
  %(prog)s --list-presets            List all saved presets

Actions:
  build          Build without upload
  upload         Build and upload to device
  clean          Clean build files
  monitor        Basic serial monitor (pio device monitor)
  debug-monitor  Enhanced monitor with memory graphs

Environments:
  default        Standard development build
  debug_memory   Debug build with memory tracking
  gh_release     Release build
        """
    )
    
    parser.add_argument(
        "actions",
        nargs="?",
        help="Action(s) to run, can be chained with '+' (e.g., 'clean+upload+monitor')"
    )
    
    parser.add_argument(
        "-e", "--env",
        choices=ENVIRONMENTS,
        default=DEFAULT_ENV,
        help=f"PlatformIO environment (default: {DEFAULT_ENV})"
    )
    
    parser.add_argument(
        "-l", "--log",
        nargs="?",
        const=str(DEFAULT_LOG_PATH),
        metavar="PATH",
        help=f"Log output to file (default: {DEFAULT_LOG_PATH})"
    )
    
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Verbose output"
    )
    
    parser.add_argument(
        "--preset",
        metavar="NAME",
        help="Run a saved preset"
    )
    
    parser.add_argument(
        "--list-presets",
        action="store_true",
        help="List all saved presets"
    )
    
    parser.add_argument(
        "--save-preset",
        metavar="NAME",
        help="Save the current command as a preset"
    )
    
    parser.add_argument(
        "--delete-preset",
        metavar="NAME",
        help="Delete a saved preset"
    )
    
    args = parser.parse_args()
    
    # Handle preset listing
    if args.list_presets:
        list_presets()
        return 0
    
    # Handle preset deletion
    if args.delete_preset:
        return 0 if delete_preset(args.delete_preset) else 1
    
    # Handle preset execution
    if args.preset:
        presets = load_presets()
        if args.preset not in presets:
            print_error(f"Preset '{args.preset}' not found")
            list_presets()
            return 1
        
        config = presets[args.preset]
        actions = config.get("actions", [])
        env = config.get("env", DEFAULT_ENV)
        log = config.get("log", False)
        log_path = Path(config["log_path"]) if config.get("log_path") else None
        
        return run_chain(actions, env, log, log_path, args.verbose)
    
    # Interactive mode if no actions specified
    if not args.actions:
        try:
            interactive_menu()
            return 0
        except KeyboardInterrupt:
            print("\n")
            return 0
    
    # Parse and run actions
    actions = parse_actions(args.actions)
    log = args.log is not None
    log_path = Path(args.log) if args.log else None
    
    # Handle save preset
    if args.save_preset:
        save_preset(args.save_preset, actions, args.env, log, args.log)
        return 0
    
    return run_chain(actions, args.env, log, log_path, args.verbose)


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