"""Test generation logic."""

from pathlib import Path
from typing import Dict, List, Optional, Tuple

import click

from guardrail.cache import Cache
from guardrail.config import GuardRailConfig
from guardrail.locator import FunctionInfo
from guardrail.llm.client import LLMClient
from guardrail.llm.spec import SpecGenerator
from guardrail.llm.testgen import TestGenerator
from guardrail.cli.detection import detect_and_analyze_changes


def generate_tests_for_functions(
    config: GuardRailConfig,
    pure_functions: List[FunctionInfo],
    function_to_file: Dict[FunctionInfo, Path],
) -> Tuple[List[Path], List[Path], Dict[FunctionInfo, Tuple[Path, str]]]:
    """
    Generate tests for all pure functions.
    
    Returns:
        Tuple of (generated_tests, cached_tests, function_test_map)
    """
    config_dict = config.load()
    max_functions = config_dict.get("max_functions_per_run", 20)
    max_function_size = config_dict.get("max_function_size", 5000)
    max_failures = config_dict.get("max_failures", 5)
    
    # Filter out functions that are too large
    processable_functions = []
    skipped_large = []
    
    for func in pure_functions:
        if len(func.source) > max_function_size:
            skipped_large.append(func)
        else:
            processable_functions.append(func)
    
    if skipped_large:
        click.echo(f"\n⚠️  Skipped {len(skipped_large)} function(s) that are too large (> {max_function_size} chars):")
        for func in skipped_large[:5]:  # Show first 5
            click.echo(f"    - {func.name} ({len(func.source)} chars)")
        if len(skipped_large) > 5:
            click.echo(f"    ... and {len(skipped_large) - 5} more")
    
    # Limit number of functions
    original_count = len(processable_functions)
    if len(processable_functions) > max_functions:
        click.echo(f"\n⚠️  Limiting to {max_functions} functions (found {original_count} changed functions)")
        click.echo(f"    Adjust 'max_functions_per_run' in config to process more")
        processable_functions = processable_functions[:max_functions]
    
    if not processable_functions:
        click.echo("\n⚠️  No processable functions found after filtering.")
        return [], [], {}
    
    click.echo(f"\n🤖 Generating tests for {len(processable_functions)} pure function(s)...")

    try:
        api_key = config.get_openai_api_key()
        if not api_key:
            raise ValueError(
                "OpenAI API key not found. Run 'guardrail init' to configure it."
            )

        # Get API settings from config
        api_timeout = config_dict.get("api_timeout", 60)
        api_max_retries = config_dict.get("api_max_retries", 3)
        
        # Initialize components
        llm_client = LLMClient(api_key=api_key)
        spec_generator = SpecGenerator(llm_client)
        test_generator = TestGenerator(llm_client)
        
        # Store API settings for use in generation
        # (We'll pass these through the generate methods)
        test_dir = config.get_test_dir()
        test_dir.mkdir(parents=True, exist_ok=True)

        # Initialize cache
        cache = None
        config_dict = config.load()
        if config_dict.get("cache_enabled", True):
            cache = Cache(config.cache_file)

        # Generate tests
        generated_tests = []
        cached_tests = []
        function_test_map = {}
        
        success_count = 0
        failure_count = 0
        total = len(processable_functions)

        for i, func in enumerate(processable_functions, 1):
            click.echo(f"    [{i}/{total}] Processing {func.name}...")
            if func.is_method:
                func_display = f"{func.parent_class}.{func.name}"
            elif func.is_nested and func.parent_function:
                func_display = f"{func.parent_function}.{func.name}"
            else:
                func_display = func.name

            try:
                file_path = function_to_file.get(func)
                if not file_path:
                    click.echo(f"    ⚠️  Could not find file for {func_display}")
                    continue

                # Check cache
                cached_test_path = None
                if cache:
                    cached_test_path = cache.get_cached_test(
                        function_source=func.source,
                        file_path=file_path,
                        function_name=func.name,
                    )

                if cached_test_path:
                    # Use cached test
                    click.echo(f"    💾 Using cached test for {func_display}")
                    generated_tests.append(cached_test_path)
                    with open(cached_test_path, "r") as f:
                        test_code = f.read()
                    function_test_map[func] = (cached_test_path, test_code)
                    cached_tests.append(cached_test_path)
                else:
                    # Get parent function source if nested
                    parent_function_source = None
                    if func.is_nested and func.parent_function:
                        # Find parent function in the same file
                        from guardrail.locator import FunctionLocator
                        locator = FunctionLocator(file_path)
                        all_functions = locator.get_all_functions()
                        for parent_func in all_functions:
                            if parent_func.name == func.parent_function and not parent_func.is_nested:
                                parent_function_source = parent_func.source
                                break
                    
                    # Generate new test
                    test_path, test_code = generate_new_test(
                        func,
                        func_display,
                        file_path,
                        spec_generator,
                        test_generator,
                        test_dir,
                        config.root,
                        parent_function_source=parent_function_source,
                    )
                    generated_tests.append(test_path)
                    function_test_map[func] = (test_path, test_code)

                    # Store in cache
                    if cache:
                        cache.store_test(
                            function_source=func.source,
                            file_path=file_path,
                            function_name=func.name,
                            test_file_path=test_path,
                        )

            except (ValueError, RuntimeError) as e:
                click.echo(f"    ❌ Failed to generate test for {func_display}: {e}", err=True)
            except Exception as e:
                click.echo(f"    ❌ Unexpected error for {func_display}: {e}", err=True)

        return generated_tests, cached_tests, function_test_map

    except ValueError as e:
        click.echo(f"⚠️  LLM not configured: {e}", err=True)
        click.echo("Set OPENAI_API_KEY environment variable to enable test generation.", err=True)
        return [], [], {}
    except Exception as e:
        click.echo(f"⚠️  Error during test generation: {e}", err=True)
        return [], [], {}


def generate_new_test(
    func: FunctionInfo,
    func_display: str,
    file_path: Path,
    spec_generator: SpecGenerator,
    test_generator: TestGenerator,
    test_dir: Path,
    root: Path,
    parent_function_source: Optional[str] = None,
) -> Tuple[Optional[Path], Optional[str]]:
    """Generate a new test for a function.
    
    Returns:
        Tuple of (test_path, test_code) or (None, None) if generation fails.
    """
    try:
        if func.is_nested:
            click.echo(f"    📝 Generating spec for nested function {func_display}...")
        else:
            click.echo(f"    📝 Generating spec for {func_display}...")

        spec_card = spec_generator.generate_spec(
            function_source=func.source,
            file_path=file_path,
            function_name=func.name,
            project_root=root,  # Pass project root for correct module path resolution
        )
    except Exception as e:
        click.echo(f"    ❌ Failed to generate spec: {e}", err=True)
        return None, None

        click.echo(f"    ✅ Spec generated (confidence: {spec_card.confidence.value}, module: {spec_card.module_path})")
        if func.is_nested:
            click.echo(f"    🧪 Generating tests (will extract nested function)...")
        else:
            click.echo(f"    🧪 Generating tests...")

        max_attempts = 3
        test_code = None
        previous_error = None  # Track previous validation error
        
        for attempt in range(max_attempts):
            try:
                test_code = test_generator.generate_test(
                    spec_card=spec_card,
                    function_source=func.source,
                    module_path=spec_card.module_path,
                    is_nested=func.is_nested,
                    parent_function_source=parent_function_source,
                    previous_validation_error=previous_error,  # Pass previous error to LLM
                )
                
                # Validate test code
                is_valid, error_msg = test_generator.validate_test_code(test_code)
                if is_valid:
                    break
                else:
                    previous_error = error_msg  # Store error for next attempt
                    if attempt < max_attempts - 1:
                        click.echo(f"    ⚠️  Validation failed (attempt {attempt + 1}/{max_attempts})")
                        # Show first line of error for brevity
                        error_preview = error_msg.split('\n')[0] if error_msg else "Unknown error"
                        click.echo(f"    📋 Error: {error_preview}")
                        click.echo(f"    🔄 Regenerating with error context...")
                        continue
                    else:
                        raise ValueError(f"Generated test failed validation after {max_attempts} attempts: {error_msg}")
            except ValueError as e:
                if attempt < max_attempts - 1:
                    previous_error = str(e)  # Store error for next attempt
                    click.echo(f"    ⚠️  Error (attempt {attempt + 1}/{max_attempts}): {e}")
                    click.echo(f"    🔄 Regenerating with error context...")
                    continue
                else:
                    raise

        if not test_code:
            raise ValueError("Failed to generate valid test code")

        file_stem = file_path.stem
        test_path = test_generator.save_test(
            test_code=test_code,
            output_dir=test_dir,
            file_name=file_stem,
            function_name=func_display,
        )

        click.echo(f"    ✅ Test saved: {test_path.relative_to(root)}")
        return test_path, test_code
    except Exception as e:
        click.echo(f"    ❌ Error generating test: {e}", err=True)
        return None, None


def regenerate_failing_tests(
    function_test_map: Dict[FunctionInfo, Tuple[Path, str]],
    failure_output: str,
    test_dir: Path,
    root: Path,
    function_to_file: Dict[FunctionInfo, Path],
    failing_tests: Optional[List[str]] = None,
) -> bool:
    """
    Regenerate tests for failing functions.
    
    Args:
        function_test_map: Map of functions to their test paths and code.
        failure_output: Test failure output.
        test_dir: Directory containing test files.
        root: Project root directory.
        function_to_file: Map of functions to their source file paths.
        failing_tests: Optional list of failing test names (e.g., ["test_file.py::test_function"]).
                      If provided, only regenerates tests for these specific failures.
    
    Returns:
        True if any tests were regenerated, False otherwise.
    """
    from guardrail.config import GuardRailConfig
    from guardrail.llm.client import LLMClient
    from guardrail.llm.spec import SpecGenerator
    from guardrail.llm.testgen import TestGenerator
    
    config = GuardRailConfig(root=root)
    
    try:
        api_key = config.get_openai_api_key()
        if not api_key:
            return False
        
        llm_client = LLMClient(api_key=api_key)
        spec_generator = SpecGenerator(llm_client)
        test_generator = TestGenerator(llm_client)
        
        regenerated_count = 0
        
        # If failing_tests is provided, only regenerate those specific tests
        functions_to_regenerate = []
        
        if failing_tests:
            # Map failing test names to functions
            for func, (test_path, test_code) in function_test_map.items():
                test_file_name = test_path.name  # e.g., "test_calculator_add.py"
                test_function_name = f"test_{func.name}"  # e.g., "test_add"
                
                # Check if this test file/function is in the failing tests
                # Pytest format: "test_file.py::test_function" or just "test_file.py"
                for failing_test in failing_tests:
                    # Extract test file and function from pytest output
                    if "::" in failing_test:
                        failing_file, failing_func = failing_test.split("::", 1)
                    else:
                        failing_file = failing_test
                        failing_func = None
                    
                    # Match test file name (with or without .py extension)
                    if test_file_name in failing_file or failing_file in test_file_name:
                        # If function name specified, check it matches
                        if failing_func is None or test_function_name == failing_func:
                            functions_to_regenerate.append((func, test_path, test_code))
                            break
        else:
            # Fallback: regenerate all if no specific failing tests provided
            functions_to_regenerate = [(func, test_path, test_code) 
                                      for func, (test_path, test_code) in function_test_map.items()]
        
        if not functions_to_regenerate:
            click.echo("    ⚠️  No matching functions found for failing tests.")
            return False
        
        click.echo(f"    🔄 Regenerating {len(functions_to_regenerate)} failing test(s)...")
        
        for func, test_path, test_code in functions_to_regenerate:
            if func.is_method:
                func_display = f"{func.parent_class}.{func.name}"
            elif func.is_nested and func.parent_function:
                func_display = f"{func.parent_function}.{func.name}"
            else:
                func_display = func.name
            
            try:
                # Get the source file path for this function
                file_path = function_to_file.get(func)
                if not file_path:
                    click.echo(f"    ⚠️  Could not find source file for {func_display}")
                    continue
                
                # Get parent function source if nested
                parent_function_source = None
                if func.is_nested and func.parent_function:
                    from guardrail.locator import FunctionLocator
                    locator = FunctionLocator(file_path)
                    all_functions = locator.get_all_functions()
                    for parent_func in all_functions:
                        if parent_func.name == func.parent_function and not parent_func.is_nested:
                            parent_function_source = parent_func.source
                            break
                
                # Regenerate test with failure context
                spec_card = spec_generator.generate_spec(
                    function_source=func.source,
                    file_path=file_path,
                    function_name=func.name,
                    project_root=root,  # Pass project root for correct module path resolution
                )
                
                regenerated_test_code = test_generator.regenerate_test_with_failure(
                    spec_card=spec_card,
                    function_source=func.source,
                    test_code=test_code,
                    failure_output=failure_output[:2000],  # Limit failure output
                    module_path=spec_card.module_path,
                    is_nested=func.is_nested,
                    parent_function_source=parent_function_source,
                )
                
                # Save regenerated test (overwrite existing)
                test_generator.save_test(
                    test_code=regenerated_test_code,
                    output_dir=test_dir,
                    file_name=file_path.stem,
                    function_name=func_display,
                )
                
                # Update function_test_map with new test code
                function_test_map[func] = (test_path, regenerated_test_code)
                regenerated_count += 1
                click.echo(f"    ✅ Regenerated test for {func_display}")
                
            except Exception as e:
                click.echo(f"    ❌ Failed to regenerate test for {func_display}: {e}", err=True)
        
        if regenerated_count > 0:
            click.echo(f"    ✅ Successfully regenerated {regenerated_count} test(s)")
        
        return regenerated_count > 0
        
    except Exception as e:
        click.echo(f"    ⚠️  Error during test regeneration: {e}", err=True)
        return False


def print_test_generation_summary(
    generated_tests: List[Path], cached_tests: List[Path], test_dir: Path, root: Path
) -> None:
    """Print summary of test generation."""
    new_count = len(generated_tests) - len(cached_tests)
    if cached_tests:
        click.echo(
            f"\n✅ {len(generated_tests)} test file(s) ready ({new_count} new, {len(cached_tests)} cached)"
        )
    else:
        click.echo(f"\n✅ Generated {len(generated_tests)} test file(s) in {test_dir.relative_to(root)}")

