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

import json
from pathlib import Path
from typing import Optional

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}

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. Include proper imports
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

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
    ) -> 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.

        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)

        # Build prompt
        prompt = self.PROMPT_TEMPLATE.format(
            function_source=function_source,
            spec_card_json=spec_card_json,
            function_name=spec_card.function_name,
        )

        # Generate test code using LLM
        response = self.llm_client.generate(
            prompt=prompt,
            system_prompt=self.SYSTEM_PROMPT,
            temperature=0.3,
            max_tokens=3000,
        )

        # 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 regenerate_test_with_failure(
        self,
        spec_card: SpecCard,
        function_source: str,
        test_code: str,
        failure_output: str,
        module_path: str,
    ) -> 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.

        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)

        # 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}
```

SpecCard:
{spec_card_json}

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?
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. Do NOT repeat the same mistakes from the previous test

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

        # 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,
        )

        # 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.

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

        Returns:
            Path to saved test file.
        """
        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

