"""
Test command for TinyTorch CLI: runs module tests using pytest.
"""

import subprocess
import sys
from argparse import ArgumentParser, Namespace
from pathlib import Path
from rich.panel import Panel
from rich.text import Text
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn

from .base import BaseCommand

class TestCommand(BaseCommand):
    @property
    def name(self) -> str:
        return "test"

    @property
    def description(self) -> str:
        return "Run module tests"

    def add_arguments(self, parser: ArgumentParser) -> None:
        parser.add_argument("module", nargs="?", help="Module to test (optional)")
        parser.add_argument("--all", action="store_true", help="Run all module tests")

    def validate_args(self, args: Namespace) -> None:
        """Validate test command arguments."""
        # Allow running without arguments to show helpful error message
        pass

    def run(self, args: Namespace) -> int:
        console = self.console
        
        if args.all:
            # Run all tests with progress bar
            failed_modules = []
            
            # Find all modules with tests
            source_dir = Path("modules/source")
            if not source_dir.exists():
                console.print(Panel("[red]❌ modules/source/ directory not found[/red]", 
                                  title="Error", border_style="red"))
                return 1
            
            # Find existing test files
            existing_tests = []
            exclude_dirs = {'.quarto', '__pycache__', '.git', '.pytest_cache'}
            for module_dir in source_dir.iterdir():
                if module_dir.is_dir() and module_dir.name not in exclude_dirs:
                    # Extract the short name from the module directory name
                    module_name = module_dir.name
                    if module_name.startswith(tuple(f"{i:02d}_" for i in range(100))):
                        short_name = module_name[3:]  # Remove "00_" prefix
                    else:
                        short_name = module_name
                    
                    # Check for centralized test file
                    centralized_test = Path("tests") / f"test_{short_name}.py"
                    if centralized_test.exists():
                        existing_tests.append(module_dir.name)
            
            if not existing_tests:
                console.print(Panel("[yellow]⚠️  No test files found in tests/ directory[/yellow]", 
                                  title="No Tests", border_style="yellow"))
                return 0
            
            console.print(Panel(f"🧪 Running tests for {len(existing_tests)} modules", 
                              title="Test Suite", border_style="bright_cyan"))
            
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
                console=console
            ) as progress:
                
                task = progress.add_task("Running tests...", total=len(existing_tests))
                
                for module in existing_tests:
                    progress.update(task, description=f"Testing {module}...")
                    
                    # Extract the short name from the module directory name
                    if module.startswith(tuple(f"{i:02d}_" for i in range(100))):
                        short_name = module[3:]  # Remove "00_" prefix
                    else:
                        short_name = module
                    
                    test_file = f"tests/test_{short_name}.py"
                    try:
                        result = subprocess.run([sys.executable, "-m", "pytest", test_file, "-v"], 
                                              capture_output=True, text=True, timeout=300)
                        
                        if result.returncode != 0:
                            failed_modules.append(module)
                            console.print(f"[red]❌ {module} tests failed[/red]")
                        else:
                            console.print(f"[green]✅ {module} tests passed[/green]")
                    except subprocess.TimeoutExpired:
                        failed_modules.append(module)
                        console.print(f"[red]❌ {module} tests timed out (5 minutes)[/red]")
                    except Exception as e:
                        failed_modules.append(module)
                        console.print(f"[red]❌ {module} tests failed with error: {e}[/red]")
                    
                    progress.advance(task)
            
            # Results summary
            if failed_modules:
                console.print(Panel(f"[red]❌ Failed modules: {', '.join(failed_modules)}[/red]", 
                                  title="Test Results", border_style="red"))
                return 1
            else:
                console.print(Panel("[green]✅ All tests passed![/green]", 
                                  title="Test Results", border_style="green"))
                return 0
        
        elif args.module:
            # Run specific module tests with detailed output
            # Extract the short name from the module directory name
            module_name = args.module
            if module_name.startswith(tuple(f"{i:02d}_" for i in range(100))):
                short_name = module_name[3:]  # Remove "00_" prefix
            else:
                short_name = module_name
            
            test_file = f"tests/test_{short_name}.py"
            
            console.print(Panel(f"🧪 Running tests for module: [bold cyan]{args.module}[/bold cyan]", 
                              title="Single Module Test", border_style="bright_cyan"))
            
            if not Path(test_file).exists():
                console.print(Panel(f"[yellow]⏳ Test file not found: {test_file}\n"
                                  f"Module '{args.module}' may not be implemented yet.[/yellow]", 
                                  title="Test Not Found", border_style="yellow"))
                return 1
            
            console.print(f"Running: pytest {test_file} -v")
            console.print()
            
            try:
                result = subprocess.run([sys.executable, "-m", "pytest", test_file, "-v"], 
                                      capture_output=True, text=True, timeout=300)
                
                # Print test output with syntax highlighting
                if result.stdout:
                    console.print("[dim]--- Test Output ---[/dim]")
                    console.print(result.stdout)
                if result.stderr:
                    console.print("[dim]--- Error Output ---[/dim]")
                    console.print(result.stderr)
                
                if result.returncode == 0:
                    console.print(Panel("[green]✅ All tests passed for {}![/green]".format(args.module), 
                                      title="Test Results", border_style="green"))
                else:
                    console.print(Panel("[red]❌ Some tests failed for {}[/red]".format(args.module), 
                                      title="Test Results", border_style="red"))
                
                return result.returncode
            except subprocess.TimeoutExpired:
                console.print(Panel("[red]❌ Tests timed out after 5 minutes for {}[/red]".format(args.module), 
                                  title="Test Results", border_style="red"))
                return 1
            except Exception as e:
                console.print(Panel("[red]❌ Test execution failed: {}[/red]".format(str(e)), 
                                  title="Test Results", border_style="red"))
                return 1
        
        else:
            # List available modules
            source_dir = Path("modules/source")
            if source_dir.exists():
                available_modules = []
                for module_dir in source_dir.iterdir():
                    if module_dir.is_dir():
                        # Extract the short name from the module directory name
                        module_name = module_dir.name
                        if module_name.startswith(tuple(f"{i:02d}_" for i in range(100))):
                            short_name = module_name[3:]  # Remove "00_" prefix
                        else:
                            short_name = module_name
                        
                        dev_file = module_dir / f"{short_name}_dev.py"
                        if dev_file.exists():
                            available_modules.append(module_dir.name)
                
                console.print(Panel(f"[red]❌ Please specify a module to test[/red]\n\n"
                                  f"Available modules: {', '.join(sorted(available_modules))}\n\n"
                                  f"[dim]Example: tito module test tensor[/dim]\n"
                                  f"[dim]For all modules: tito module test --all[/dim]", 
                                  title="Module Required", border_style="red"))
            else:
                console.print(Panel("[red]❌ No modules/source directory found[/red]", 
                                  title="Error", border_style="red"))
            return 1 