"""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


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 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)}")

