"""LLM-based test generation."""

import ast
import json
import re
from pathlib import Path
from typing import Optional, Tuple

from guardrail.llm.client import LLMClient
from guardrail.models import SpecCard


class TestGenerator:
    """Generates pytest tests from SpecCard using LLM."""

    SYSTEM_PROMPT = """You are an expert Python test writer. Your task is to generate comprehensive pytest unit tests for Python functions.

Generate pytest test code that:
1. Tests all expected behaviors from the SpecCard
2. Covers all edge cases listed
3. Uses plain assert statements (no Hypothesis in V1)
4. Is deterministic (no randomness)
5. Does not require network/database access
6. Has correct imports
7. Tests the function thoroughly

Return ONLY the Python test code, no markdown, no explanations, just the test code."""

    PROMPT_TEMPLATE = """Generate pytest tests for this function. CRITICAL: Tests must match the function's ACTUAL behavior.

Function source code:
```python
{function_source}
```

SpecCard:
{spec_card_json}

{is_nested_section}

IMPORTANT - CORRECT IMPORT PATH:
The function is located at: {module_path}
You MUST use this EXACT import path: `from {module_path} import {function_name}`
If the module path includes package prefixes (e.g., "app.calculator"), use the FULL path.
Do NOT use relative imports or guess the import path.

REQUIREMENTS:
1. Read the function code FIRST to understand what it actually does - analyze the implementation
2. Generate tests that match the function's REAL behavior, not assumptions or generic patterns
3. Only test with valid input types - do NOT test with incompatible types (e.g., don't pass boolean to numeric functions, don't pass string to arithmetic operations unless the function handles strings)
4. If the SpecCard suggests invalid edge cases (wrong types, nonsensical inputs), IGNORE them and use your code analysis instead
5. Test the function as it IS implemented, not as it should be or as we assume it works
6. Use plain assert statements (no Hypothesis)
7. Ensure tests are deterministic
8. Use the EXACT import path provided above: `from {module_path} import {function_name}`
9. Test function name should be: test_{function_name}
10. Do not require network/database - if function has side effects, skip or mock appropriately
11. Validate input types match what the function actually expects based on its code
12. AVOID syntax errors: No leading zeros in numeric literals (e.g., use 1 not 01), no invalid date formats
13. Ensure all imports are valid and use the correct module path
{extraction_requirements}

VALIDATION CHECKLIST (before returning):
- [ ] Import statement uses the exact module path: `from {module_path} import {function_name}`
- [ ] No syntax errors (check for leading zeros, invalid literals)
- [ ] All test inputs are valid types for the function
- [ ] Test code is syntactically correct Python

Return ONLY the Python test code, no markdown code blocks, no explanations."""

    def __init__(self, llm_client: Optional[LLMClient] = None):
        """
        Initialize test generator.

        Args:
            llm_client: Optional LLM client. If None, creates one with default settings.
        """
        self.llm_client = llm_client or LLMClient()

    def generate_test(
        self, 
        spec_card: SpecCard, 
        function_source: str, 
        module_path: str,
        is_nested: bool = False,
        parent_function_source: Optional[str] = None,
        previous_validation_error: Optional[str] = None,
    ) -> str:
        """
        Generate pytest test code from SpecCard.

        Args:
            spec_card: SpecCard object with function specification.
            function_source: Source code of the function.
            module_path: Module path to import the function from.
            is_nested: Whether the function is nested inside another function.
            parent_function_source: Source code of parent function if nested.
            previous_validation_error: Optional error message from previous validation failure.

        Returns:
            Python test code as string.

        Raises:
            RuntimeError: If LLM call fails.
        """
        # Convert SpecCard to JSON for prompt
        spec_card_dict = spec_card.model_dump()
        spec_card_json = json.dumps(spec_card_dict, indent=2)

        # Handle nested functions
        if is_nested and parent_function_source:
            is_nested_section = f"""
IMPORTANT: This is a NESTED function (inside another function).

Parent function:
```python
{parent_function_source}
```

Nested function to test:
```python
{function_source}
```

You MUST extract the nested function to module level in the test file before testing it.
"""
            extraction_requirements = """
12. EXTRACT the nested function: Copy the nested function definition to module level in your test code
13. Handle closure variables: If the nested function uses variables from parent scope, extract those too
14. Test the extracted function directly (not through the parent function)
15. Ensure the extracted function has all necessary dependencies and imports
"""
        else:
            is_nested_section = ""
            extraction_requirements = ""

        # Add validation error context if this is a retry
        validation_error_section = ""
        if previous_validation_error:
            validation_error_section = f"""
⚠️ PREVIOUS ATTEMPT FAILED VALIDATION:
{previous_validation_error}

CRITICAL: The previous generated test had validation errors. You MUST fix these issues:
- Do NOT repeat the same mistakes
- Pay special attention to: {previous_validation_error}
- Review your code carefully before returning it
- If the error mentions "leading zero", ensure you use integers without leading zeros (e.g., use 1, 2, 3 instead of 01, 02, 03)
- If you need to represent dates/times with leading zeros, use STRINGS (e.g., "01/01/2024" not 01/01/2024)
- Double-check all numeric literals for leading zeros
- Make sure all syntax is valid Python

This is a RETRY attempt - you must generate DIFFERENT code that fixes the validation error.
"""

        # Build prompt
        prompt = self.PROMPT_TEMPLATE.format(
            function_source=function_source,
            spec_card_json=spec_card_json,
            function_name=spec_card.function_name,
            module_path=module_path,
            is_nested_section=is_nested_section,
            extraction_requirements=extraction_requirements,
        )
        
        # Prepend validation error section if present
        if validation_error_section:
            prompt = validation_error_section + "\n" + prompt

        # Get timeout and retry settings (defaults)
        api_timeout = 60
        api_max_retries = 3
        
        # Generate test code using LLM
        response = self.llm_client.generate(
            prompt=prompt,
            system_prompt=self.SYSTEM_PROMPT,
            temperature=0.3,
            max_tokens=3000,
            timeout=api_timeout,
            max_retries=api_max_retries,
        )

        # Clean up response (remove markdown code blocks if present)
        response = response.strip()
        if response.startswith("```python"):
            lines = response.split("\n")
            response = "\n".join(lines[1:-1]) if len(lines) > 2 else response
        elif response.startswith("```"):
            lines = response.split("\n")
            response = "\n".join(lines[1:-1]) if len(lines) > 2 else response

        return response

    def validate_test_code(self, test_code: str) -> Tuple[bool, Optional[str]]:
        """
        Validate generated test code for syntax errors and common issues.
        
        Args:
            test_code: Generated test code to validate.
            
        Returns:
            Tuple of (is_valid: bool, error_message: Optional[str])
        """
        # Check for syntax errors
        try:
            ast.parse(test_code)
        except SyntaxError as e:
            return False, f"Syntax error: {e.msg} at line {e.lineno}. Fix the syntax error in the generated code."
        
        # Check for common issues
        issues = []
        
        # Check for leading zeros in numeric literals (more precise check)
        # Parse the code to find actual numeric literals with leading zeros
        lines = test_code.split('\n')
        for i, line in enumerate(lines, 1):
            # Skip comments
            if '#' in line:
                line_content = line.split('#')[0]
            else:
                line_content = line
            
            # Skip if line is empty or only whitespace
            if not line_content.strip():
                continue
            
            # Check for leading zeros in numeric literals (not in strings)
            # Pattern: word boundary, 0, followed by 1+ digits, not part of a string
            # Look for patterns like: 01, 02, 0123, etc. but not in quotes
            leading_zero_matches = re.finditer(r'\b0\d+\b', line_content)
            for match in leading_zero_matches:
                # Check if this is inside a string
                start_pos = match.start()
                # Simple check: count quotes before this position
                before_match = line_content[:start_pos]
                # If even number of quotes before, we're not in a string
                quote_count = before_match.count('"') + before_match.count("'")
                if quote_count % 2 == 0:
                    # This is a numeric literal with leading zero
                    issues.append(f"Leading zero in numeric literal at line {i}: '{match.group()}' (use {int(match.group())} instead, or use string '{match.group()}' if you need the leading zero)")
                    break  # Only report one per line
        
        # Check for invalid import patterns (basic check - imports starting with numbers)
        invalid_imports = re.findall(r'from\s+(\d+)', test_code)
        if invalid_imports:
            issues.append(f"Invalid import statement starting with number: {invalid_imports}. Module names cannot start with numbers.")
        
        # Check for common date format issues (leading zeros in dates not in strings)
        # This is more of a warning - dates with leading zeros should be strings
        date_pattern = r'(?<!["\'])\b0\d+[/-]\d+[/-]\d+\b(?!["\'])'
        if re.search(date_pattern, test_code):
            issues.append("Date format with leading zeros detected outside of strings. Use string format like '01/01/2024' instead of 01/01/2024.")
        
        if issues:
            # Combine all issues into a clear error message
            error_msg = "Validation errors found:\n" + "\n".join(f"- {issue}" for issue in issues)
            return False, error_msg
        
        return True, None

    def regenerate_test_with_failure(
        self,
        spec_card: SpecCard,
        function_source: str,
        test_code: str,
        failure_output: str,
        module_path: str,
        is_nested: bool = False,
        parent_function_source: Optional[str] = None,
    ) -> str:
        """
        Regenerate tests with failure information to fix incorrect tests.

        Args:
            spec_card: SpecCard object with function specification.
            function_source: Source code of the function.
            test_code: Previous test code that failed.
            failure_output: Output from the failed test execution.
            module_path: Module path to import the function from.
            is_nested: Whether the function is nested inside another function.
            parent_function_source: Source code of parent function if nested.

        Returns:
            Regenerated Python test code as string.
        """
        # Convert SpecCard to JSON for prompt
        spec_card_dict = spec_card.model_dump()
        spec_card_json = json.dumps(spec_card_dict, indent=2)

        # Handle nested functions in regeneration
        nested_section = ""
        if is_nested and parent_function_source:
            nested_section = f"""
IMPORTANT: This is a NESTED function. You MUST extract it to module level.

Parent function:
```python
{parent_function_source}
```

Nested function to test:
```python
{function_source}
```
"""
        
        # Build regeneration prompt
        prompt = f"""The previous test failed. Regenerate tests based on the actual failure and function code:

Function code (READ THIS CAREFULLY):
```python
{function_source}
```

{nested_section}
SpecCard:
{spec_card_json}

IMPORTANT - CORRECT IMPORT PATH:
The function is located at: {module_path}
You MUST use this EXACT import path: `from {module_path} import {spec_card.function_name}`
If the module path includes package prefixes (e.g., "app.calculator"), use the FULL path.
Do NOT use relative imports or guess the import path.

Previous test (FAILED):
```python
{test_code}
```

Failure output:
{failure_output}

CRITICAL INSTRUCTIONS:
1. Analyze the failure - what went wrong? Was it a type error? Wrong expected value? Invalid input? Import error?
2. Read the function code to understand what it ACTUALLY does
3. Generate NEW tests that:
   - Match the function's ACTUAL behavior (from reading the code)
   - Use valid input types only (check what types the function actually accepts)
   - Test what the function really does, not what we assumed
   - Fix the issues from the previous test
   - Only include realistic edge cases that make sense for this function
4. If the failure was due to invalid input types, generate tests with correct types
5. If the failure was due to wrong expected values, analyze the function code to determine correct expected values
6. If the failure was due to import errors, use the EXACT import path: `from {module_path} import {spec_card.function_name}`
7. If this is a nested function, EXTRACT it to module level in the test file before testing
8. AVOID syntax errors: No leading zeros in numeric literals (e.g., use 1 not 01), no invalid date formats
9. Do NOT repeat the same mistakes from the previous test

VALIDATION CHECKLIST (before returning):
- [ ] Import statement uses the exact module path: `from {module_path} import {spec_card.function_name}`
- [ ] No syntax errors (check for leading zeros, invalid literals)
- [ ] All test inputs are valid types for the function
- [ ] Test code is syntactically correct Python
- [ ] All issues from the previous failure are addressed

Return ONLY the corrected Python test code, no markdown code blocks, no explanations."""

        # Get timeout and retry settings (defaults)
        api_timeout = 60
        api_max_retries = 3
        
        # Generate regenerated test code using LLM
        response = self.llm_client.generate(
            prompt=prompt,
            system_prompt=self.SYSTEM_PROMPT,
            temperature=0.2,  # Lower temperature for more focused regeneration
            max_tokens=3000,
            timeout=api_timeout,
            max_retries=api_max_retries,
        )

        # Clean up response (remove markdown code blocks if present)
        response = response.strip()
        if response.startswith("```python"):
            lines = response.split("\n")
            response = "\n".join(lines[1:-1]) if len(lines) > 2 else response
        elif response.startswith("```"):
            lines = response.split("\n")
            response = "\n".join(lines[1:-1]) if len(lines) > 2 else response

        return response

    def save_test(
        self,
        test_code: str,
        output_dir: Path,
        file_name: str,
        function_name: str,
    ) -> Path:
        """
        Save generated test to file with validation.

        Args:
            test_code: Generated test code.
            output_dir: Directory to save test file.
            file_name: Base name for the test file.
            function_name: Name of the function being tested.

        Returns:
            Path to saved test file.
            
        Raises:
            ValueError: If test code is invalid.
        """
        # Validate test code before saving
        is_valid, error_msg = self.validate_test_code(test_code)
        if not is_valid:
            raise ValueError(f"Generated test code is invalid: {error_msg}")
        
        output_dir.mkdir(parents=True, exist_ok=True)

        # Create test file name: test_<file>_<func>.py
        safe_file_name = file_name.replace("/", "_").replace("\\", "_").replace(".", "_")
        safe_function_name = function_name.replace(".", "_")
        test_file_name = f"test_{safe_file_name}_{safe_function_name}.py"
        test_path = output_dir / test_file_name

        # Write test code
        with open(test_path, "w", encoding="utf-8") as f:
            f.write(test_code)

        return test_path

