"""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 unit tests for the following function based on its SpecCard:

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

SpecCard:
{spec_card_json}

Requirements:
1. Generate comprehensive pytest tests covering all expected behaviors and edge cases
2. Use plain assert statements (no Hypothesis)
3. Ensure tests are deterministic
4. Include proper imports
5. Test function name should be: test_{function_name}
6. Do not require network/database - if function has side effects, skip or mock appropriately
7. Test all edge cases from the SpecCard

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

