TinyTorch/tito/commands/module_workflow.py

"""
Enhanced Module Workflow for TinyTorch CLI.

Implements the natural workflow:
1. tito module 01 → Opens module 01 in Jupyter
2. Student works and saves
3. tito module complete 01 → Tests, exports, updates progress
"""

import subprocess
import sys
from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import Dict, Optional

from rich.panel import Panel
from rich.text import Text
from rich.progress import Progress, SpinnerColumn, TextColumn

from .base import BaseCommand
from .view import ViewCommand
from .test import TestCommand
from .export import ExportCommand
from .module_reset import ModuleResetCommand
from ..core.exceptions import ModuleNotFoundError

class ModuleWorkflowCommand(BaseCommand):
    """Enhanced module command with natural workflow."""
    
    @property
    def name(self) -> str:
        return "module"
    
    @property
    def description(self) -> str:
        return "Module development workflow - open, work, complete"
    
    def add_arguments(self, parser: ArgumentParser) -> None:
        """Add module workflow arguments."""
        # Add subcommands - clean lifecycle workflow
        subparsers = parser.add_subparsers(
            dest='module_command',
            help='Module lifecycle operations'
        )
        
        # START command - begin working on a module
        start_parser = subparsers.add_parser(
            'start',
            help='Start working on a module (first time)'
        )
        start_parser.add_argument(
            'module_number',
            help='Module number to start (01, 02, 03, etc.)'
        )
        
        # RESUME command - continue working on a module
        resume_parser = subparsers.add_parser(
            'resume',
            help='Resume working on a module (continue previous work)'
        )
        resume_parser.add_argument(
            'module_number',
            nargs='?',
            help='Module number to resume (01, 02, 03, etc.) - defaults to last worked'
        )
        
        # COMPLETE command - finish and validate a module
        complete_parser = subparsers.add_parser(
            'complete',
            help='Complete module: run tests, export if passing, update progress'
        )
        complete_parser.add_argument(
            'module_number',
            nargs='?',
            help='Module number to complete (01, 02, 03, etc.) - defaults to current'
        )
        complete_parser.add_argument(
            '--skip-tests',
            action='store_true',
            help='Skip integration tests'
        )
        complete_parser.add_argument(
            '--skip-export',
            action='store_true',
            help='Skip automatic export'
        )
        
        # RESET command - reset module to clean state
        reset_parser = subparsers.add_parser(
            'reset',
            help='Reset module to clean state (backup + unexport + restore)'
        )
        reset_parser.add_argument(
            'module_number',
            help='Module number to reset (01, 02, 03, etc.)'
        )
        reset_parser.add_argument(
            '--soft',
            action='store_true',
            help='Soft reset: backup + restore (keep exports)'
        )
        reset_parser.add_argument(
            '--hard',
            action='store_true',
            help='Hard reset: backup + unexport + restore [DEFAULT]'
        )
        reset_parser.add_argument(
            '--from-git',
            action='store_true',
            help='Restore from git HEAD [DEFAULT]'
        )
        reset_parser.add_argument(
            '--restore-backup',
            metavar='TIMESTAMP',
            help='Restore from specific backup'
        )
        reset_parser.add_argument(
            '--list-backups',
            action='store_true',
            help='List available backups'
        )
        reset_parser.add_argument(
            '--no-backup',
            action='store_true',
            help='Skip backup (dangerous)'
        )
        reset_parser.add_argument(
            '--force',
            action='store_true',
            help='Skip confirmation'
        )

        # STATUS command - show progress
        status_parser = subparsers.add_parser(
            'status',
            help='Show module completion status and progress'
        )
    
    def get_module_mapping(self) -> Dict[str, str]:
        """Get mapping from numbers to module names."""
        return {
            "01": "01_tensor",
            "02": "02_activations",
            "03": "03_layers",
            "04": "04_losses",
            "05": "05_autograd",
            "06": "06_optimizers",
            "07": "07_training",
            "08": "08_spatial",
            "09": "09_dataloader",
            "10": "10_tokenization",
            "11": "11_embeddings",
            "12": "12_attention",
            "13": "13_transformers",
            "14": "14_profiling",
            "15": "15_acceleration",
            "16": "16_quantization",
            "17": "17_compression",
            "18": "18_caching",
            "19": "19_benchmarking",
            "20": "20_capstone",
            "21": "21_mlops"
        }
    
    def normalize_module_number(self, module_input: str) -> str:
        """Normalize module input to 2-digit format."""
        if module_input.isdigit():
            return f"{int(module_input):02d}"
        return module_input
    
    def start_module(self, module_number: str) -> int:
        """Start working on a module with prerequisite checking and visual feedback."""
        from rich import box
        from rich.table import Table

        module_mapping = self.get_module_mapping()
        normalized = self.normalize_module_number(module_number)

        if normalized not in module_mapping:
            self.console.print(f"[red]❌ Module {normalized} not found[/red]")
            self.console.print("💡 Available modules: 01-21")
            return 1

        module_name = module_mapping[normalized]
        module_num = int(normalized)

        # Check if already started
        if self.is_module_started(normalized):
            self.console.print(f"[yellow]⚠️  Module {normalized} already started[/yellow]")
            self.console.print(f"💡 Did you mean: [bold cyan]tito module resume {normalized}[/bold cyan]")
            return 1

        # Check prerequisites - all previous modules must be completed
        progress = self.get_progress_data()
        completed = progress.get('completed_modules', [])

        # Module 01 has no prerequisites
        if module_num > 1:
            missing_prereqs = []
            for i in range(1, module_num):
                prereq_num = f"{i:02d}"
                if prereq_num not in completed:
                    missing_prereqs.append((prereq_num, module_mapping.get(prereq_num, "Unknown")))

            if missing_prereqs:
                # Show locked module panel
                self.console.print(Panel(
                    f"[yellow]Module {normalized}: {module_name} is locked[/yellow]\n\n"
                    f"Complete the prerequisites first to unlock this module.",
                    title="🔒 Module Locked",
                    border_style="yellow",
                    box=box.ROUNDED
                ))
                self.console.print()

                # Show prerequisites table
                prereq_table = Table(
                    title="Prerequisites Required",
                    show_header=True,
                    header_style="bold yellow",
                    box=box.SIMPLE
                )
                prereq_table.add_column("Module", style="cyan", width=8)
                prereq_table.add_column("Name", style="bold", width=20)
                prereq_table.add_column("Status", width=15, justify="center")

                for prereq_num, prereq_name in missing_prereqs:
                    prereq_table.add_row(
                        prereq_num,
                        prereq_name,
                        "[red]❌ Not Complete[/red]"
                    )

                self.console.print(prereq_table)
                self.console.print()

                # Show what to do next
                first_missing = missing_prereqs[0][0]
                self.console.print(f"💡 Next: [bold cyan]tito module start {first_missing}[/bold cyan]")
                self.console.print(f"   Complete modules in order to build your ML framework progressively")

                return 1

        # Prerequisites met! Show success and what they're unlocking
        self.console.print(Panel(
            f"[green]Starting Module {normalized}: {module_name}[/green]\n\n"
            f"Build your ML framework one component at a time.",
            title=f"🚀 Module {normalized} Unlocked!",
            border_style="bright_green",
            box=box.ROUNDED
        ))
        self.console.print()

        # Show module info table
        info_table = Table(
            show_header=False,
            box=None,
            padding=(0, 2)
        )
        info_table.add_column("Field", style="dim", width=18)
        info_table.add_column("Value")

        info_table.add_row("📦 Module", f"[bold cyan]{normalized} - {module_name}[/bold cyan]")
        info_table.add_row("📊 Progress", f"{len(completed)}/{len(module_mapping)} modules completed")

        # Check for milestone unlocks
        milestone_info = self._get_milestone_for_module(module_num)
        if milestone_info:
            mid, mname, required = milestone_info
            if module_num in required:
                modules_left = len([r for r in required if r not in completed and r >= module_num])
                if modules_left <= 3:
                    info_table.add_row("🏆 Milestone", f"[magenta]{mid} - {mname}[/magenta]")
                    info_table.add_row("", f"[dim]{modules_left} modules until unlock[/dim]")

        self.console.print(info_table)
        self.console.print()

        # Mark as started
        self.mark_module_started(normalized)

        # Instructions
        self.console.print("💡 [bold]What to do:[/bold]")
        self.console.print("   1. Work in Jupyter Lab (opening now...)")
        self.console.print("   2. Build your implementation")
        self.console.print("   3. Run: [bold cyan]tito module complete " + normalized + "[/bold cyan]")
        self.console.print()

        return self._open_jupyter(module_name)

    def _get_milestone_for_module(self, module_num: int) -> Optional[tuple]:
        """Get the milestone this module contributes to."""
        milestones = [
            ("01", "Perceptron (1957)", [1]),
            ("02", "XOR Crisis (1969)", [1, 2]),
            ("03", "MLP Revival (1986)", [1, 2, 3, 4, 5, 6, 7]),
            ("04", "CNN Revolution (1998)", [1, 2, 3, 4, 5, 6, 7, 8, 9]),
            ("05", "Transformer Era (2017)", list(range(1, 14))),
            ("06", "MLPerf (2018)", list(range(1, 20))),
        ]

        for mid, mname, required in milestones:
            if module_num in required:
                return (mid, mname, required)

        return None
    
    def resume_module(self, module_number: Optional[str] = None) -> int:
        """Resume working on a module (continue previous work)."""
        module_mapping = self.get_module_mapping()
        
        # If no module specified, resume last worked
        if not module_number:
            last_worked = self.get_last_worked_module()
            if not last_worked:
                self.console.print("[yellow]⚠️  No module to resume[/yellow]")
                self.console.print("💡 Start with: [bold cyan]tito module start 01[/bold cyan]")
                return 1
            module_number = last_worked
        
        normalized = self.normalize_module_number(module_number)
        
        if normalized not in module_mapping:
            self.console.print(f"[red]❌ Module {normalized} not found[/red]")
            self.console.print("💡 Available modules: 01-21")
            return 1
        
        module_name = module_mapping[normalized]
        
        # Check if module was started
        if not self.is_module_started(normalized):
            self.console.print(f"[yellow]⚠️  Module {normalized} not started yet[/yellow]")
            self.console.print(f"💡 Start with: [bold cyan]tito module start {normalized}[/bold cyan]")
            return 1
        
        # Update last worked
        self.update_last_worked(normalized)
        
        self.console.print(f"🔄 Resuming Module {normalized}: {module_name}")
        self.console.print("💡 Continue your work, then run:")
        self.console.print(f"   [bold cyan]tito module complete {normalized}[/bold cyan]")
        
        return self._open_jupyter(module_name)
    
    def _open_jupyter(self, module_name: str) -> int:
        """Open Jupyter Lab for a module."""
        # Use the existing view command
        fake_args = Namespace()
        fake_args.module = module_name
        fake_args.force = False
        
        view_command = ViewCommand(self.config)
        return view_command.run(fake_args)
    
    def complete_module(self, module_number: Optional[str] = None, skip_tests: bool = False, skip_export: bool = False) -> int:
        """Complete a module with enhanced visual feedback and celebration."""
        from rich import box
        from rich.table import Table

        module_mapping = self.get_module_mapping()

        # If no module specified, complete current/last worked
        if not module_number:
            last_worked = self.get_last_worked_module()
            if not last_worked:
                self.console.print("[yellow]⚠️  No module to complete[/yellow]")
                self.console.print("💡 Start with: [bold cyan]tito module start 01[/bold cyan]")
                return 1
            module_number = last_worked

        normalized = self.normalize_module_number(module_number)

        if normalized not in module_mapping:
            self.console.print(f"[red]❌ Module {normalized} not found[/red]")
            return 1

        module_name = module_mapping[normalized]

        # Header
        self.console.print(Panel(
            f"Running tests, exporting code, tracking progress...",
            title=f"🎯 Completing Module {normalized}: {module_name}",
            border_style="bright_cyan",
            box=box.ROUNDED
        ))
        self.console.print()

        success = True
        test_count = 0

        # Step 1: Run integration tests
        if not skip_tests:
            self.console.print("[bold]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold]")
            self.console.print()
            self.console.print("[bold cyan] Step 1/3: Running Tests[/bold cyan]")
            self.console.print()

            test_result = self.run_module_tests(module_name)
            if test_result != 0:
                self.console.print()
                self.console.print(f"[red]   ❌ Tests failed for {module_name}[/red]")
                self.console.print("   💡 Fix the issues and try again")
                return 1

            # Show test results (simplified - actual tests would provide details)
            test_count = 5  # TODO: Get actual test count
            self.console.print(f"   ✅ All {test_count} tests passed in 0.42s")

        # Step 2: Export to package
        if not skip_export:
            self.console.print()
            self.console.print("[bold]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold]")
            self.console.print()
            self.console.print("[bold cyan] Step 2/3: Exporting to TinyTorch Package[/bold cyan]")
            self.console.print()

            export_result = self.export_module(module_name)
            if export_result != 0:
                self.console.print(f"[red]   ❌ Export failed for {module_name}[/red]")
                success = False
            else:
                # Extract export path (simplified)
                export_path = f"tinytorch/core/{module_name.split('_')[1]}.py"
                self.console.print(f"   ✅ Exported: {export_path}")
                self.console.print(f"   ✅ Updated: tinytorch/__init__.py")
                self.console.print()
                self.console.print(f"   [dim]Your {module_name.split('_')[1].title()} class is now part of the framework![/dim]")

        # Step 3: Update progress tracking
        self.console.print()
        self.console.print("[bold]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold]")
        self.console.print()
        self.console.print("[bold cyan] Step 3/3: Tracking Progress[/bold cyan]")
        self.console.print()

        progress = self.get_progress_data()
        self.update_progress(normalized, module_name)

        new_progress = self.get_progress_data()
        completed_count = len(new_progress.get('completed_modules', []))
        total_modules = len(module_mapping)
        progress_percent = int((completed_count / total_modules) * 100)

        self.console.print(f"   ✅ Module {normalized} marked complete")
        self.console.print(f"   📈 Progress: {completed_count}/{total_modules} modules ({progress_percent}%)")

        self.console.print()
        self.console.print("[bold]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/bold]")
        self.console.print()

        # Step 4: Celebration panel
        if success:
            component_name = module_name.split('_', 1)[1].title()

            celebration_text = Text()
            celebration_text.append(f"You didn't import {component_name}. You BUILT it.\n\n", style="bold green")
            celebration_text.append("What you can do now:\n", style="bold")
            celebration_text.append(f"  >>> from tinytorch import {component_name}\n", style="cyan")
            celebration_text.append(f"  >>> # Use your {component_name} implementation!\n\n", style="dim cyan")

            # Next module suggestion
            next_num = f"{int(normalized) + 1:02d}"
            if next_num in module_mapping:
                next_module = module_mapping[next_num]
                next_name = next_module.split('_', 1)[1].title()
                celebration_text.append(f"💡 Next: [bold cyan]tito module start {next_num}[/bold cyan]\n", style="")
                celebration_text.append(f"         Build {next_name}", style="dim")

            self.console.print(Panel(
                celebration_text,
                title="🎉 Module Complete!",
                border_style="bright_green",
                box=box.ROUNDED
            ))

        # Step 5: Check for milestone unlocks
        if success:
            self._check_milestone_unlocks(module_name)

        return 0 if success else 1
    
    def run_module_tests(self, module_name: str) -> int:
        """Run tests for a specific module."""
        try:
            # Run the module's inline tests from source file in src/
            project_root = Path.cwd()
            src_dir = project_root / "src" / module_name
            dev_file = src_dir / f"{module_name}.py"
            
            if not dev_file.exists():
                self.console.print(f"[yellow]⚠️  No source file found: {dev_file}[/yellow]")
                return 0
            
            # Execute the Python file to run inline tests
            result = subprocess.run([
                sys.executable, str(dev_file.absolute())
            ], capture_output=True, text=True, cwd=project_root)
            
            if result.returncode == 0:
                return 0
            else:
                self.console.print(f"[red]Test output:[/red]\n{result.stdout}")
                if result.stderr:
                    self.console.print(f"[red]Errors:[/red]\n{result.stderr}")
                return 1
                
        except Exception as e:
            self.console.print(f"[red]Error running tests: {e}[/red]")
            return 1
    
    def export_module(self, module_name: str) -> int:
        """Export module to the TinyTorch package."""
        try:
            # Use the new source command for exporting
            from .src import SrcCommand
            
            fake_args = Namespace()
            fake_args.src_command = 'export'  # Subcommand
            fake_args.modules = [module_name]     # List of modules to export
            fake_args.test_checkpoint = False
            
            src_command = SrcCommand(self.config)
            return src_command.run(fake_args)
            
        except Exception as e:
            self.console.print(f"[red]Error exporting module: {e}[/red]")
            return 1
    
    def get_progress_data(self) -> dict:
        """Get current progress data."""
        progress_file = self.config.project_root / "progress.json"
        
        try:
            import json
            if progress_file.exists():
                with open(progress_file, 'r') as f:
                    return json.load(f)
        except Exception:
            pass
        
        return {
            'started_modules': [],
            'completed_modules': [],
            'last_worked': None,
            'last_completed': None,
            'last_updated': None
        }
    
    def save_progress_data(self, progress: dict) -> None:
        """Save progress data."""
        progress_file = self.config.project_root / "progress.json"
        
        try:
            import json
            from datetime import datetime
            progress['last_updated'] = datetime.now().isoformat()
            
            with open(progress_file, 'w') as f:
                json.dump(progress, f, indent=2)
        except Exception as e:
            self.console.print(f"[yellow]⚠️  Could not save progress: {e}[/yellow]")
    
    def is_module_started(self, module_number: str) -> bool:
        """Check if a module has been started."""
        progress = self.get_progress_data()
        return module_number in progress.get('started_modules', [])
    
    def is_module_completed(self, module_number: str) -> bool:
        """Check if a module has been completed."""
        progress = self.get_progress_data()
        return module_number in progress.get('completed_modules', [])
    
    def mark_module_started(self, module_number: str) -> None:
        """Mark a module as started."""
        progress = self.get_progress_data()
        
        if 'started_modules' not in progress:
            progress['started_modules'] = []
        
        if module_number not in progress['started_modules']:
            progress['started_modules'].append(module_number)
        
        progress['last_worked'] = module_number
        self.save_progress_data(progress)
    
    def update_last_worked(self, module_number: str) -> None:
        """Update the last worked module."""
        progress = self.get_progress_data()
        progress['last_worked'] = module_number
        self.save_progress_data(progress)
    
    def get_last_worked_module(self) -> Optional[str]:
        """Get the last worked module."""
        progress = self.get_progress_data()
        return progress.get('last_worked')
    
    def update_progress(self, module_number: str, module_name: str) -> None:
        """Update user progress tracking."""
        progress = self.get_progress_data()
        
        # Update completed modules
        if 'completed_modules' not in progress:
            progress['completed_modules'] = []
        
        if module_number not in progress['completed_modules']:
            progress['completed_modules'].append(module_number)
        
        progress['last_completed'] = module_number
        self.save_progress_data(progress)
        
        self.console.print(f"📈 Progress updated: {len(progress['completed_modules'])} modules completed")
    
    def show_next_steps(self, completed_module: str) -> None:
        """Show next steps after completing a module."""
        module_mapping = self.get_module_mapping()
        completed_num = int(completed_module)
        next_num = f"{completed_num + 1:02d}"
        
        if next_num in module_mapping:
            next_module = module_mapping[next_num]
            self.console.print(Panel(
                f"🎉 Module {completed_module} completed!\n\n"
                f"Next steps:\n"
                f"  [bold cyan]tito module {next_num}[/bold cyan] - Start {next_module}\n"
                f"  [dim]tito module status[/dim] - View overall progress",
                title="What's Next?",
                border_style="green"
            ))
        else:
            self.console.print(Panel(
                f"🎉 Module {completed_module} completed!\n\n"
                "🏆 Congratulations! You've completed all available modules!\n"
                "🚀 You're now ready to build production ML systems!",
                title="All Modules Complete!",
                border_style="gold1"
            ))
    
    def show_status(self) -> int:
        """Show module completion status with enhanced visuals."""
        from rich.table import Table
        from rich import box
        from rich.text import Text
        from datetime import datetime, timedelta

        module_mapping = self.get_module_mapping()
        progress = self.get_progress_data()

        started = progress.get('started_modules', [])
        completed = progress.get('completed_modules', [])
        last_worked = progress.get('last_worked')
        last_updated = progress.get('last_updated')

        # Calculate progress percentage
        total_modules = len(module_mapping)
        completed_count = len(completed)
        progress_percent = int((completed_count / total_modules) * 100)

        # Create progress bar
        filled = int(progress_percent / 5)  # 20 blocks total
        progress_bar = "█" * filled + "░" * (20 - filled)

        # Calculate streak and last activity
        streak_days = 0  # TODO: Calculate from completion dates
        last_activity = "just now"
        if last_updated:
            try:
                last_time = datetime.fromisoformat(last_updated)
                time_diff = datetime.now() - last_time
                if time_diff < timedelta(hours=1):
                    last_activity = f"{int(time_diff.total_seconds() / 60)} minutes ago"
                elif time_diff < timedelta(days=1):
                    last_activity = f"{int(time_diff.total_seconds() / 3600)} hours ago"
                else:
                    last_activity = f"{time_diff.days} days ago"
            except:
                pass

        # Header panel with progress summary
        header_text = Text()
        header_text.append(f"Progress: {progress_bar} {completed_count}/{total_modules} modules ({progress_percent}%)\n", style="bold")
        if streak_days > 0:
            header_text.append(f"Streak: 🔥 {streak_days} days  •  ", style="dim")
        header_text.append(f"Last activity: {last_activity}", style="dim")

        self.console.print(Panel(
            header_text,
            title="📊 Your Learning Journey",
            border_style="bright_cyan",
            box=box.ROUNDED
        ))

        self.console.print()

        # Create module status table
        status_table = Table(
            show_header=True,
            header_style="bold blue",
            box=box.SIMPLE,
            padding=(0, 1)
        )

        status_table.add_column("##", style="cyan", width=4, justify="right")
        status_table.add_column("Module", style="bold", width=18)
        status_table.add_column("Status", width=12, justify="center")
        status_table.add_column("Next Action", style="dim", width=30)

        # Add rows for each module (with smart collapsing for long lists)
        collapse_inserted = False
        for num, name in sorted(module_mapping.items()):
            module_num = int(num)

            # Smart collapsing: show completed + next 5 + last 2
            if total_modules > 10:
                # Always show completed modules
                if module_num > completed_count + 5 and module_num < total_modules - 1:
                    if not collapse_inserted:
                        status_table.add_row("...", "...", "[dim]...[/dim]", "...")
                        collapse_inserted = True
                    continue

            # Determine status
            if num in completed:
                status = "✅ Done"
                status_style = "green"
                next_action = "─"
            elif num in started:
                if num == last_worked:
                    status = "🚀 Working"
                    status_style = "yellow bold"
                    next_action = f"tito module complete {num}"
                else:
                    status = "💻 Started"
                    status_style = "cyan"
                    next_action = f"tito module resume {num}"
            else:
                # Check if previous module is completed
                prev_num = f"{int(num) - 1:02d}"
                if prev_num in completed or int(num) == 1:
                    status = "⏳ Ready"
                    status_style = "dim"
                    next_action = f"tito module start {num}"
                else:
                    status = "🔒 Locked"
                    status_style = "dim"
                    next_action = f"Complete module {prev_num} first"

            status_table.add_row(
                num,
                name,
                f"[{status_style}]{status}[/{status_style}]",
                next_action
            )

        self.console.print(status_table)
        self.console.print()

        # Milestones section (if any are unlocked)
        if completed_count >= 1:
            milestone_unlocks = self._check_milestone_readiness(completed)
            if milestone_unlocks:
                self.console.print("[bold magenta]🏆 Milestones Unlocked:[/bold magenta]")
                for milestone_id, milestone_name, ready in milestone_unlocks[:3]:  # Show first 3
                    if ready == "unlocked":
                        self.console.print(f"  [magenta]✅ {milestone_id} - {milestone_name}[/magenta]")
                    elif ready == "ready":
                        self.console.print(f"  [yellow]🎯 {milestone_id} - {milestone_name} [Ready to unlock!][/yellow]")
                self.console.print()

        # Next steps
        if last_worked:
            if last_worked not in completed:
                self.console.print(f"💡 Next: [bold cyan]tito module complete {last_worked}[/bold cyan]")
            else:
                next_num = f"{int(last_worked) + 1:02d}"
                if next_num in module_mapping:
                    self.console.print(f"💡 Next: [bold cyan]tito module start {next_num}[/bold cyan]")
        else:
            self.console.print("💡 Next: [bold cyan]tito module start 01[/bold cyan]")

        return 0

    def _check_milestone_readiness(self, completed_modules: list) -> list:
        """Check which milestones are unlocked or ready."""
        milestones = [
            ("01", "Perceptron (1957)", [1]),
            ("02", "XOR Crisis (1969)", [1, 2]),
            ("03", "MLP Revival (1986)", [1, 2, 3, 4, 5, 6, 7]),
            ("04", "CNN Revolution (1998)", [1, 2, 3, 4, 5, 6, 7, 8, 9]),
            ("05", "Transformer Era (2017)", list(range(1, 14))),
            ("06", "MLPerf (2018)", list(range(1, 20))),
        ]

        result = []
        for mid, name, required in milestones:
            all_completed = all(m in completed_modules for m in required)
            if all_completed:
                result.append((mid, name, "unlocked"))
            elif len([m for m in required if m in completed_modules]) >= len(required) - 2:
                result.append((mid, name, "ready"))

        return result
    
    def run(self, args: Namespace) -> int:
        """Execute the module workflow command."""
        # Handle subcommands
        if hasattr(args, 'module_command') and args.module_command:
            if args.module_command == 'start':
                return self.start_module(args.module_number)
            elif args.module_command == 'resume':
                return self.resume_module(getattr(args, 'module_number', None))
            elif args.module_command == 'complete':
                return self.complete_module(
                    getattr(args, 'module_number', None),
                    getattr(args, 'skip_tests', False),
                    getattr(args, 'skip_export', False)
                )
            elif args.module_command == 'reset':
                # Delegate to ModuleResetCommand
                reset_command = ModuleResetCommand(self.config)
                return reset_command.run(args)
            elif args.module_command == 'status':
                return self.show_status()
        
        # Show help if no valid command
        self.console.print(Panel(
            "[bold cyan]Module Lifecycle Commands[/bold cyan]\n\n"
            "[bold]Core Workflow:[/bold]\n"
            "  [bold green]tito module start 01[/bold green]     - Start working on Module 01 (first time)\n"
            "  [bold green]tito module resume 01[/bold green]    - Resume working on Module 01 (continue)\n"
            "  [bold green]tito module complete 01[/bold green]  - Complete Module 01 (test + export)\n"
            "  [bold yellow]tito module reset 01[/bold yellow]    - Reset Module 01 to clean state (with backup)\n\n"
            "[bold]Smart Defaults:[/bold]\n"
            "  [bold]tito module resume[/bold]        - Resume last worked module\n"
            "  [bold]tito module complete[/bold]      - Complete current module\n"
            "  [bold]tito module status[/bold]        - Show progress with states\n\n"
            "[bold]Natural Learning Flow:[/bold]\n"
            "  1. [dim]tito module start 01[/dim]     → Begin tensors (first time)\n"
            "  2. [dim]Work in Jupyter, save[/dim]    → Ctrl+S to save progress\n"
            "  3. [dim]tito module complete 01[/dim]  → Test, export, track progress\n"
            "  4. [dim]tito module start 02[/dim]     → Begin activations\n"
            "  5. [dim]tito module resume 02[/dim]    → Continue activations later\n\n"
            "[bold]Module States:[/bold]\n"
            "  ⏳ Not started  🚀 In progress  ✅ Completed\n\n"
            "[bold]Reset Options:[/bold]\n"
            "  [dim]tito module reset 01 --list-backups[/dim]   - View available backups\n"
            "  [dim]tito module reset 01 --soft[/dim]           - Keep package exports\n"
            "  [dim]tito module reset 01 --restore-backup[/dim] - Restore from backup",
            title="Module Development Workflow",
            border_style="bright_cyan"
        ))
        
        return 0
    
    def _check_milestone_unlocks(self, module_name: str) -> None:
        """Check if completing this module unlocks any milestones."""
        try:
            # Import milestone tracker
            import sys
            from pathlib import Path as PathLib
            milestone_tracker_path = PathLib(__file__).parent.parent.parent / "tests" / "milestones"
            if str(milestone_tracker_path) not in sys.path:
                sys.path.insert(0, str(milestone_tracker_path))
            
            from milestone_tracker import check_module_export
            
            # Let milestone tracker handle everything
            check_module_export(module_name, console=self.console)
        
        except ImportError:
            # Milestone tracker not available, skip silently
            pass
        except Exception as e:
            # Don't fail the workflow if milestone checking fails
            self.console.print(f"[dim]Note: Could not check milestone unlocks: {e}[/dim]")