k8s-tool/k8s_tool/main.py

"""Main application entry point."""

import sys
import os
import subprocess
import tempfile
import yaml
from typing import Optional, Dict
import questionary
from questionary import Style
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax

from k8s_tool.k8s_client import K8sClient
from k8s_tool.config import ConfigManager

console = Console()

# Custom style for questionary prompts
custom_style = Style([
    ('qmark', 'fg:#673ab7 bold'),
    ('question', 'bold'),
    ('answer', 'fg:#f44336 bold'),
    ('pointer', 'fg:#673ab7 bold'),
    ('highlighted', 'fg:#673ab7 bold'),
    ('selected', 'fg:#cc5454'),
    ('separator', 'fg:#cc5454'),
    ('instruction', ''),
    ('text', ''),
])


class K8sTool:
    """Main application class."""

    def __init__(self):
        """Initialize the application."""
        self.k8s_client = K8sClient()
        self.config = ConfigManager()
        self.current_namespace: Optional[str] = None

    def run(self):
        """Run the main application loop."""
        console.print(Panel.fit(
            "[bold cyan]K8s Tool[/bold cyan]\n"
            "[dim]Interactive kubectl helper[/dim]",
            border_style="cyan"
        ))

        while True:
            try:
                self._main_menu()
            except KeyboardInterrupt:
                console.print("\n[yellow]Goodbye! 👋[/yellow]")
                sys.exit(0)
            except Exception as e:
                console.print(f"[red]Error:[/red] {e}")
                if questionary.confirm("Continue?", style=custom_style).ask():
                    continue
                else:
                    break

    def _main_menu(self):
        """Display main menu."""
        # Display current context
        current_context = self.k8s_client.current_context or "[dim]unknown[/dim]"
        console.print(f"[bold]Current context:[/bold] [cyan]{current_context}[/cyan]")

        # Display current namespace
        namespace_info = f"[cyan]{self.current_namespace}[/cyan]" if self.current_namespace else "[dim]not selected[/dim]"
        is_favorite = self.config.is_favorite(self.current_namespace) if self.current_namespace else False
        fav_indicator = " ⭐" if is_favorite else ""
        console.print(f"[bold]Current namespace:[/bold] {namespace_info}{fav_indicator}")

        choices = [
            "Switch Context",
            "Select Namespace",
            "Manage Favorites",
            "List Deployments",
            "Restart Deployment",
            "Scale Deployment",
            "View ConfigMaps",
            "Edit ConfigMap",
            "View Pod Logs",
            "Exit"
        ]

        action = questionary.select(
            "What would you like to do?",
            choices=choices,
            style=custom_style
        ).ask()

        if action == "Switch Context":
            self._switch_context()
        elif action == "Select Namespace":
            self._select_namespace()
        elif action == "Manage Favorites":
            self._manage_favorites()
        elif action == "List Deployments":
            self._list_deployments()
        elif action == "Restart Deployment":
            self._restart_deployment()
        elif action == "Scale Deployment":
            self._scale_deployment()
        elif action == "View ConfigMaps":
            self._view_configmaps()
        elif action == "Edit ConfigMap":
            self._edit_configmaps()
        elif action == "View Pod Logs":
            self._view_pod_logs()
        elif action == "Exit":
            console.print("[yellow]Goodbye! 👋[/yellow]")
            sys.exit(0)

    def _select_namespace(self):
        """Select namespace."""
        console.print("[dim]Fetching namespaces...[/dim]")
        namespaces = self.k8s_client.get_namespaces()

        if not namespaces:
            console.print("[red]No namespaces found[/red]")
            return

        # Get favorites
        favorites = self.config.get_favorites()

        # Separate favorites and non-favorites
        fav_namespaces = [ns for ns in namespaces if ns in favorites]
        other_namespaces = [ns for ns in namespaces if ns not in favorites]

        # Create choices with stars for favorites
        choices = []
        if fav_namespaces:
            choices.extend([f"⭐ {ns}" for ns in fav_namespaces])
            if other_namespaces:
                choices.append(questionary.Separator("─" * 40))
        choices.extend(other_namespaces)

        selected = questionary.select(
            "Select namespace:",
            choices=choices,
            style=custom_style
        ).ask()

        if selected:
            # Remove star prefix if present
            namespace = selected.replace("⭐ ", "")
            self.current_namespace = namespace
            is_fav = self.config.is_favorite(namespace)
            fav_text = " (favorite)" if is_fav else ""
            console.print(f"[green]✓[/green] Namespace set to: [cyan]{namespace}[/cyan]{fav_text}")

    def _switch_context(self):
        """Switch Kubernetes context."""
        console.print("[dim]Fetching contexts...[/dim]")
        contexts = self.k8s_client.get_contexts()

        if not contexts:
            console.print("[red]No contexts found[/red]")
            return

        # Show contexts table
        self.k8s_client.display_contexts_table(contexts)

        # Filter out current context for selection
        available_contexts = [ctx['name'] for ctx in contexts if not ctx['is_active']]

        if not available_contexts:
            console.print("[yellow]Only one context available[/yellow]")
            return

        # Add option to view current
        choices = available_contexts + [questionary.Separator(), "Cancel"]

        selected = questionary.select(
            "Select context to switch to:",
            choices=choices,
            style=custom_style
        ).ask()

        if selected and selected != "Cancel":
            if self.k8s_client.switch_context(selected):
                # Reset namespace selection after context switch
                self.current_namespace = None
                console.print("[yellow]Note:[/yellow] Namespace selection reset. Please select a namespace.")

    def _manage_favorites(self):
        """Manage favorite namespaces."""
        choices = []

        if self.current_namespace:
            if self.config.is_favorite(self.current_namespace):
                choices.append(f"Remove '{self.current_namespace}' from favorites")
            else:
                choices.append(f"Add '{self.current_namespace}' to favorites")

        choices.extend([
            "View all favorites",
            "Add namespace to favorites",
            "Remove namespace from favorites",
            "Clear all favorites",
            "Back to main menu"
        ])

        action = questionary.select(
            "Favorites management:",
            choices=choices,
            style=custom_style
        ).ask()

        if not action or action == "Back to main menu":
            return

        if action.startswith("Add '") and action.endswith("' to favorites"):
            # Add current namespace
            self.config.add_favorite(self.current_namespace)

        elif action.startswith("Remove '") and action.endswith("' from favorites"):
            # Remove current namespace
            self.config.remove_favorite(self.current_namespace)

        elif action == "View all favorites":
            self._view_favorites()

        elif action == "Add namespace to favorites":
            self._add_namespace_to_favorites()

        elif action == "Remove namespace from favorites":
            self._remove_namespace_from_favorites()

        elif action == "Clear all favorites":
            if questionary.confirm("Are you sure you want to clear all favorites?", style=custom_style, default=False).ask():
                self.config.clear_favorites()

    def _view_favorites(self):
        """View all favorite namespaces."""
        favorites = self.config.get_favorites()

        if not favorites:
            console.print("[yellow]No favorites yet[/yellow]")
            return

        console.print("\n[bold]Favorite namespaces:[/bold]")
        for ns in favorites:
            indicator = " [dim](current)[/dim]" if ns == self.current_namespace else ""
            console.print(f"  ⭐ [cyan]{ns}[/cyan]{indicator}")
        console.print()

    def _add_namespace_to_favorites(self):
        """Add a namespace to favorites."""
        console.print("[dim]Fetching namespaces...[/dim]")
        namespaces = self.k8s_client.get_namespaces()

        if not namespaces:
            console.print("[red]No namespaces found[/red]")
            return

        # Filter out already favorite namespaces
        favorites = self.config.get_favorites()
        available = [ns for ns in namespaces if ns not in favorites]

        if not available:
            console.print("[yellow]All namespaces are already in favorites[/yellow]")
            return

        namespace = questionary.select(
            "Select namespace to add to favorites:",
            choices=available,
            style=custom_style
        ).ask()

        if namespace:
            self.config.add_favorite(namespace)

    def _remove_namespace_from_favorites(self):
        """Remove a namespace from favorites."""
        favorites = self.config.get_favorites()

        if not favorites:
            console.print("[yellow]No favorites to remove[/yellow]")
            return

        namespace = questionary.select(
            "Select namespace to remove from favorites:",
            choices=favorites,
            style=custom_style
        ).ask()

        if namespace:
            self.config.remove_favorite(namespace)

    def _ensure_namespace_selected(self) -> bool:
        """Ensure namespace is selected."""
        if not self.current_namespace:
            console.print("[yellow]Please select a namespace first[/yellow]")
            self._select_namespace()
            return self.current_namespace is not None
        return True

    def _list_deployments(self):
        """List deployments in current namespace."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching deployments in {self.current_namespace}...[/dim]")
        deployments = self.k8s_client.get_deployments(self.current_namespace)

        if not deployments:
            console.print("[yellow]No deployments found[/yellow]")
            return

        self.k8s_client.display_deployments_table(deployments)

        if questionary.confirm("View pods for a deployment?", style=custom_style, default=False).ask():
            dep_names = [dep['name'] for dep in deployments]
            dep_name = questionary.select(
                "Select deployment:",
                choices=dep_names,
                style=custom_style
            ).ask()

            if dep_name:
                self._view_deployment_pods(dep_name)

    def _view_deployment_pods(self, deployment_name: str):
        """View pods for a deployment."""
        label_selector = f"app={deployment_name}"
        console.print(f"[dim]Fetching pods for deployment {deployment_name}...[/dim]")
        pods = self.k8s_client.get_pods(self.current_namespace, label_selector=label_selector)

        if not pods:
            # Try without label selector
            console.print("[yellow]No pods found with label selector, fetching all pods...[/yellow]")
            pods = self.k8s_client.get_pods(self.current_namespace)

        if pods:
            self.k8s_client.display_pods_table(pods)
        else:
            console.print("[yellow]No pods found[/yellow]")

    def _restart_deployment(self):
        """Restart a deployment."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching deployments in {self.current_namespace}...[/dim]")
        deployments = self.k8s_client.get_deployments(self.current_namespace)

        if not deployments:
            console.print("[yellow]No deployments found[/yellow]")
            return

        dep_names = [dep['name'] for dep in deployments]
        dep_name = questionary.select(
            "Select deployment to restart:",
            choices=dep_names,
            style=custom_style
        ).ask()

        if not dep_name:
            return

        confirm = questionary.confirm(
            f"Are you sure you want to restart {dep_name}?",
            style=custom_style,
            default=False
        ).ask()

        if confirm:
            self.k8s_client.restart_deployment(self.current_namespace, dep_name)

    def _scale_deployment(self):
        """Scale a deployment."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching deployments in {self.current_namespace}...[/dim]")
        deployments = self.k8s_client.get_deployments(self.current_namespace)

        if not deployments:
            console.print("[yellow]No deployments found[/yellow]")
            return

        dep_names = [f"{dep['name']} (current: {dep['replicas']})" for dep in deployments]
        selected = questionary.select(
            "Select deployment to scale:",
            choices=dep_names,
            style=custom_style
        ).ask()

        if not selected:
            return

        dep_name = selected.split(" (current:")[0]

        replicas = questionary.text(
            "Enter number of replicas:",
            validate=lambda text: text.isdigit() and int(text) >= 0,
            style=custom_style
        ).ask()

        if replicas:
            self.k8s_client.scale_deployment(self.current_namespace, dep_name, int(replicas))

    def _view_configmaps(self):
        """View ConfigMaps."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching ConfigMaps in {self.current_namespace}...[/dim]")
        configmaps = self.k8s_client.get_configmaps(self.current_namespace)

        if not configmaps:
            console.print("[yellow]No ConfigMaps found[/yellow]")
            return

        self.k8s_client.display_configmaps_table(configmaps)

        if questionary.confirm("View ConfigMap data?", style=custom_style, default=False).ask():
            cm_names = [cm['name'] for cm in configmaps]
            cm_name = questionary.select(
                "Select ConfigMap:",
                choices=cm_names,
                style=custom_style
            ).ask()

            if cm_name:
                data = self.k8s_client.get_configmap_data(self.current_namespace, cm_name)
                if data:
                    console.print(f"\n[bold]ConfigMap:[/bold] [cyan]{cm_name}[/cyan]\n")
                    for key, value in data.items():
                        console.print(f"[bold cyan]--- {key} ---[/bold cyan]")
                        console.print(Syntax(value, "yaml", line_numbers=False, background_color="default"))
                        console.print()

    def _edit_configmaps(self):
        """Edit a ConfigMap."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching ConfigMaps in {self.current_namespace}...[/dim]")
        configmaps = self.k8s_client.get_configmaps(self.current_namespace)

        if not configmaps:
            console.print("[yellow]No ConfigMaps found[/yellow]")
            return

        cm_names = [cm['name'] for cm in configmaps]
        cm_name = questionary.select(
            "Select ConfigMap to edit:",
            choices=cm_names,
            style=custom_style
        ).ask()

        if cm_name:
            data = self.k8s_client.get_configmap_data(self.current_namespace, cm_name)
            if data:
                self._edit_configmap(cm_name, data)

    def _edit_configmap(self, cm_name: str, current_data: Dict[str, str]):
        """Edit ConfigMap data in text editor."""
        # Get editor from environment or use default
        editor = os.environ.get('EDITOR', os.environ.get('VISUAL', 'vi'))

        # Create temporary file with ConfigMap data
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as tf:
            temp_file = tf.name
            # Write current data as YAML
            yaml.dump(current_data, tf, default_flow_style=False, allow_unicode=True)

        try:
            # Get modification time before editing
            mtime_before = os.path.getmtime(temp_file)

            # Open editor
            subprocess.call([editor, temp_file])

            # Check if file was modified
            mtime_after = os.path.getmtime(temp_file)

            if mtime_after == mtime_before:
                console.print("[yellow]No changes made[/yellow]")
                return

            # Read edited data
            with open(temp_file, 'r') as f:
                try:
                    new_data = yaml.safe_load(f)
                except yaml.YAMLError as e:
                    console.print(f"[red]Error parsing YAML:[/red] {e}")
                    if questionary.confirm("Retry editing?", style=custom_style, default=False).ask():
                        self._edit_configmap(cm_name, current_data)
                    return

            # Validate that new_data is a dict
            if not isinstance(new_data, dict):
                console.print("[red]Error:[/red] ConfigMap data must be a dictionary")
                if questionary.confirm("Retry editing?", style=custom_style, default=False).ask():
                    self._edit_configmap(cm_name, current_data)
                return

            # Convert all values to strings (ConfigMap requirement)
            new_data = {k: str(v) if not isinstance(v, str) else v for k, v in new_data.items()}

            # Show diff
            console.print("\n[bold]Changes:[/bold]")
            removed_keys = set(current_data.keys()) - set(new_data.keys())
            added_keys = set(new_data.keys()) - set(current_data.keys())
            modified_keys = {k for k in current_data.keys() & new_data.keys()
                           if current_data[k] != new_data[k]}

            if removed_keys:
                console.print(f"[red]Removed keys:[/red] {', '.join(removed_keys)}")
            if added_keys:
                console.print(f"[green]Added keys:[/green] {', '.join(added_keys)}")
            if modified_keys:
                console.print(f"[yellow]Modified keys:[/yellow] {', '.join(modified_keys)}")

            if not (removed_keys or added_keys or modified_keys):
                console.print("[yellow]No changes detected[/yellow]")
                return

            # Confirm update
            if questionary.confirm(
                f"Apply changes to ConfigMap '{cm_name}'?",
                style=custom_style,
                default=False
            ).ask():
                self.k8s_client.update_configmap(self.current_namespace, cm_name, new_data)

        finally:
            # Clean up temp file
            try:
                os.unlink(temp_file)
            except Exception:
                pass

    def _view_pod_logs(self):
        """View pod logs."""
        if not self._ensure_namespace_selected():
            return

        console.print(f"[dim]Fetching pods in {self.current_namespace}...[/dim]")
        pods = self.k8s_client.get_pods(self.current_namespace)

        if not pods:
            console.print("[yellow]No pods found[/yellow]")
            return

        pod_names = [f"{pod['name']} ({pod['status']})" for pod in pods]
        selected = questionary.select(
            "Select pod:",
            choices=pod_names,
            style=custom_style
        ).ask()

        if not selected:
            return

        pod_name = selected.split(" (")[0]

        # Check if pod has multiple containers
        containers = self.k8s_client.get_pod_containers(self.current_namespace, pod_name)
        container = None

        if len(containers) > 1:
            container = questionary.select(
                "Select container:",
                choices=containers,
                style=custom_style
            ).ask()
        elif len(containers) == 1:
            container = containers[0]

        tail_lines = questionary.text(
            "Number of lines to show (default: 100):",
            default="100",
            validate=lambda text: text.isdigit() and int(text) > 0,
            style=custom_style
        ).ask()

        if tail_lines:
            logs = self.k8s_client.get_pod_logs(
                self.current_namespace,
                pod_name,
                container=container,
                tail_lines=int(tail_lines)
            )

            if logs:
                title = f"[bold cyan]Logs: {pod_name}"
                if container:
                    title += f" [{container}]"
                title += "[/bold cyan]"
                console.print(f"\n{title}\n")
                console.print(Syntax(logs, "log", line_numbers=False, background_color="default"))
                console.print()


def main():
    """Main entry point."""
    try:
        app = K8sTool()
        app.run()
    except KeyboardInterrupt:
        console.print("\n[yellow]Goodbye! 👋[/yellow]")
        sys.exit(0)
    except Exception as e:
        console.print(f"[red]Fatal error:[/red] {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()