"""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)
    """
    click.echo(f"\n🤖 Generating tests for {len(pure_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."
            )

        # Initialize components
        llm_client = LLMClient(api_key=api_key)
        spec_generator = SpecGenerator(llm_client)
        test_generator = TestGenerator(llm_client)
        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 = {}

        for func in pure_functions:
            func_display = (
                f"{func.parent_class}.{func.name}" if func.is_method else 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:
                    # Generate new test
                    test_path, test_code = generate_new_test(
                        func,
                        func_display,
                        file_path,
                        spec_generator,
                        test_generator,
                        test_dir,
                        config.root,
                    )
                    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,
) -> Tuple[Path, str]:
    """Generate a new test for a function."""
    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,
    )

    click.echo(f"    ✅ Spec generated (confidence: {spec_card.confidence.value})")
    click.echo(f"    🧪 Generating tests...")

    test_code = test_generator.generate_test(
        spec_card=spec_card,
        function_source=func.source,
        module_path=spec_card.module_path,
    )

    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


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],
) -> 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.
    
    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
        
        for func, (test_path, test_code) in function_test_map.items():
            func_display = (
                f"{func.parent_class}.{func.name}" if func.is_method else func.name
            )
            
            try:
                # Get the source file path for this function
                file_path = function_to_file.get(func)
                if not file_path:
                    continue
                
                # Regenerate test with failure context
                spec_card = spec_generator.generate_spec(
                    function_source=func.source,
                    file_path=file_path,
                    function_name=func.name,
                )
                
                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,
                )
                
                # 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)
        
        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)}")

