"""Iterative Refinement Loop for Fix Generation (Phase 5).

Refines fixes through multiple validation stages:
1. Syntax validation
2. Type checking
3. Test execution
4. Determinism verification (N runs)

Each stage provides feedback for LLM to refine the fix.
"""

from __future__ import annotations

import asyncio
import logging
import os
import re
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from autodebug.llm_first.models import (
    FailureContext,
    FixHypothesis,
    FixResponse,
    RefinementResult,
    RefinementStage,
    VerificationResult,
)
from autodebug.prompts.sanitizer import sanitize_code, sanitize_error, sanitize_plain

logger = logging.getLogger(__name__)

# Default timeout for test execution (seconds)
DEFAULT_TEST_TIMEOUT = 60

# Default number of verification runs for determinism
DEFAULT_VERIFICATION_RUNS = 10

# Maximum parallel verification runs
MAX_PARALLEL_RUNS = 4


class IterativeRefiner:
    """Refine fixes based on validation feedback.

    Iteratively improves fixes through stages:
    1. Syntax: Parse and validate Python syntax
    2. Types: Run mypy/pyright for type errors
    3. Test: Run the actual test
    4. Verify: Run N times for determinism
    """

    def __init__(
        self,
        instructor_client=None,
        max_iterations: int = 5,
        verification_runs: int = DEFAULT_VERIFICATION_RUNS,
        test_timeout: int = DEFAULT_TEST_TIMEOUT,
        parallel_runs: int = MAX_PARALLEL_RUNS,
        use_replay: bool = True,
        transcripts_dir: Optional[Path] = None,
        cleanup_on_success: bool = True,
    ):
        """Initialize iterative refiner.

        Args:
            instructor_client: LLM client for refinement.
            max_iterations: Maximum refinement attempts.
            verification_runs: Number of runs for determinism check.
            test_timeout: Timeout for test execution in seconds.
            parallel_runs: Maximum parallel verification runs.
            use_replay: Whether to use replay infrastructure.
            transcripts_dir: Directory containing transcripts for replay.
            cleanup_on_success: Whether to clean up temp directories on success.
        """
        self.instructor_client = instructor_client
        self.max_iterations = max_iterations
        self.verification_runs = verification_runs
        self.test_timeout = test_timeout
        self.parallel_runs = parallel_runs
        self.use_replay = use_replay
        self.transcripts_dir = transcripts_dir
        self.cleanup_on_success = cleanup_on_success
        self._work_dirs: List[Path] = []

    async def _llm_generate_with_retry(self, prompt: str, response_model, max_retries: int = 3):
        """C-43: LLM generate with exponential backoff retry for transient errors."""
        last_error: Exception | None = None
        for attempt in range(max_retries):
            try:
                return await self.instructor_client.generate(
                    prompt, response_model=response_model,
                )
            except Exception as e:
                last_error = e
                if attempt < max_retries - 1:
                    delay = 2 ** attempt
                    logger.warning(
                        "Refiner LLM call failed (attempt %d/%d), retrying in %ds: %s",
                        attempt + 1, max_retries, delay, e,
                    )
                    await asyncio.sleep(delay)
        raise last_error  # type: ignore[misc]

    async def refine_until_valid(
        self,
        initial_fix: FixHypothesis,
        context: FailureContext,
    ) -> Tuple[FixHypothesis, List[RefinementResult]]:
        """Refine fix until it passes all validation stages.

        Args:
            initial_fix: Initial fix to refine.
            context: Failure context.

        Returns:
            Tuple of (refined fix, refinement history).
        """
        current_fix = initial_fix
        history: List[RefinementResult] = []

        for iteration in range(self.max_iterations):
            # Stage 1: Syntax validation
            syntax_result = await self._validate_syntax(current_fix)
            if not syntax_result.valid:
                pre_refine = current_fix.model_copy() if hasattr(current_fix, 'model_copy') else current_fix
                current_fix = await self._refine_for_syntax(
                    current_fix, syntax_result.errors, history
                )
                # Convergence detection: if fix unchanged, break early
                if hasattr(current_fix, 'diff') and hasattr(pre_refine, 'diff') and current_fix.diff == pre_refine.diff:
                    history.append(RefinementResult(
                        stage=RefinementStage.SYNTAX,
                        input_fix=pre_refine,
                        output_fix=current_fix,
                        error_message=syntax_result.errors,
                        success=False,
                        reasoning="Syntax error (no improvement, stopping)",
                    ))
                    break
                history.append(RefinementResult(
                    stage=RefinementStage.SYNTAX,
                    input_fix=pre_refine,
                    output_fix=current_fix,
                    error_message=syntax_result.errors,
                    success=False,
                    reasoning="Syntax error",
                ))
                continue

            # Stage 2: Type checking
            type_result = await self._validate_types(current_fix, context)
            if not type_result.valid:
                pre_refine = current_fix.model_copy() if hasattr(current_fix, 'model_copy') else current_fix
                current_fix = await self._refine_for_types(
                    current_fix, type_result.errors, history
                )
                # H-134: Convergence detection -- stop if fix unchanged
                if hasattr(current_fix, 'diff') and hasattr(pre_refine, 'diff') and current_fix.diff == pre_refine.diff:
                    history.append(RefinementResult(
                        stage=RefinementStage.TYPES,
                        input_fix=pre_refine,
                        output_fix=current_fix,
                        error_message=type_result.errors,
                        success=False,
                        reasoning="Type error (no improvement, stopping)",
                    ))
                    break
                history.append(RefinementResult(
                    stage=RefinementStage.TYPES,
                    input_fix=pre_refine,
                    output_fix=current_fix,
                    error_message=type_result.errors,
                    success=False,
                    reasoning="Type error",
                ))
                continue

            # Stage 3: Test execution
            # H-133: Snapshot fix before refinement (same pattern as stages 1-2)
            test_result = await self._run_test(current_fix, context)
            if not test_result.passed:
                pre_refine = current_fix.model_copy() if hasattr(current_fix, 'model_copy') else current_fix
                current_fix = await self._refine_for_test_failure(
                    current_fix, test_result, context, history
                )
                # H-134: Convergence detection -- stop if fix unchanged
                if hasattr(current_fix, 'diff') and hasattr(pre_refine, 'diff') and current_fix.diff == pre_refine.diff:
                    history.append(RefinementResult(
                        stage=RefinementStage.TEST,
                        input_fix=pre_refine,
                        output_fix=current_fix,
                        failure_output=test_result.output,
                        success=False,
                        reasoning=f"Test failed (no improvement, stopping): {test_result.error_summary}",
                    ))
                    break
                history.append(RefinementResult(
                    stage=RefinementStage.TEST,
                    input_fix=pre_refine,
                    output_fix=current_fix,
                    failure_output=test_result.output,
                    success=False,
                    reasoning=f"Test failed: {test_result.error_summary}",
                ))
                continue

            # Stage 4: Verification (N runs)
            # H-133: Snapshot fix before refinement (same pattern as stages 1-2)
            verify_result = await self._verify_determinism(current_fix, context)
            if not verify_result.is_deterministic:
                pre_refine = current_fix.model_copy() if hasattr(current_fix, 'model_copy') else current_fix
                current_fix = await self._refine_for_flakiness(
                    current_fix, verify_result, context, history
                )
                # H-134: Convergence detection -- stop if fix unchanged
                if hasattr(current_fix, 'diff') and hasattr(pre_refine, 'diff') and current_fix.diff == pre_refine.diff:
                    history.append(RefinementResult(
                        stage=RefinementStage.VERIFY,
                        input_fix=pre_refine,
                        output_fix=current_fix,
                        failure_output=str(verify_result.failure_patterns),
                        success=False,
                        reasoning="Fix is flaky (no improvement, stopping)",
                    ))
                    break
                history.append(RefinementResult(
                    stage=RefinementStage.VERIFY,
                    input_fix=pre_refine,
                    output_fix=current_fix,
                    failure_output=str(verify_result.failure_patterns),
                    success=False,
                    reasoning="Fix is flaky",
                ))
                continue

            # All stages passed!
            current_fix.syntax_valid = True
            current_fix.types_valid = True
            current_fix.test_passed = True
            current_fix.verified_deterministic = True

            history.append(RefinementResult(
                stage=RefinementStage.VERIFY,
                input_fix=initial_fix,
                output_fix=current_fix,
                success=True,
                reasoning="All validation stages passed",
            ))

            return current_fix, history

        # Return best attempt even if not perfect
        return current_fix, history

    async def _validate_syntax(
        self,
        fix: FixHypothesis,
    ) -> ValidationResult:
        """Validate Python syntax of the fix.

        Args:
            fix: Fix to validate.

        Returns:
            Validation result with errors.
        """
        import ast

        # Extract code from diff
        code = self._extract_code_from_diff(fix.diff)

        try:
            ast.parse(code)
            return ValidationResult(valid=True, errors="")
        except SyntaxError as e:
            return ValidationResult(
                valid=False,
                errors=f"SyntaxError at line {e.lineno}: {e.msg}",
            )

    async def _validate_types(
        self,
        fix: FixHypothesis,
        context: FailureContext,
    ) -> ValidationResult:
        """Run type checker on the fix.

        L-03: Integrates the FixValidator from the validation module
        (which supports Ty and mypy) rather than relying on naive
        string pattern matching.

        Args:
            fix: Fix to validate.
            context: Failure context.

        Returns:
            Validation result with type errors.
        """
        diff = fix.diff
        issues = []

        # Try to use the real FixValidator with Ty/mypy type checking
        try:
            from autodebug.validation.fix_validator import FixValidator as _FixValidator

            validator = _FixValidator(use_ty=True)
            # Extract patched code from context + diff
            patched_code = getattr(fix, "patched_code", None)
            if patched_code:
                result = validator.validate_code(patched_code)
                if not result.is_valid:
                    type_errors = [
                        e for e in (result.errors or [])
                        if "type" in e.lower() or "incompatible" in e.lower()
                            or "attribute" in e.lower()
                    ]
                    if type_errors:
                        return ValidationResult(
                            valid=False,
                            errors="\n".join(type_errors[:5]),
                        )
                return ValidationResult(valid=True, errors="")
        except (ImportError, Exception):
            pass  # FixValidator unavailable; fall back to heuristic checks

        # Fallback: basic string-pattern checks for common type issues
        if "None." in diff:
            issues.append("Possible NoneType access")

        if "str + int" in diff or "int + str" in diff:
            issues.append("Type mismatch in concatenation")

        if issues:
            return ValidationResult(valid=False, errors="\n".join(issues))

        return ValidationResult(valid=True, errors="")

    async def _run_test(
        self,
        fix: FixHypothesis,
        context: FailureContext,
    ) -> TestRunResult:
        """Run the test with the fix applied.

        Creates a temporary directory, applies the patch, writes the test file,
        and runs pytest in a subprocess with proper isolation.

        Args:
            fix: Fix to test.
            context: Failure context.

        Returns:
            Test run result with actual execution output.
        """
        work_dir = None
        try:
            # Create temporary working directory
            work_dir = Path(tempfile.mkdtemp(prefix="hikaflow_test_"))
            self._work_dirs.append(work_dir)

            # Extract the target file path from the diff
            target_file = self._extract_target_file(fix.diff)
            if not target_file:
                # Try to get from context
                target_file = context.test_file

            # Apply the patch to create the patched source file
            patched_code = self._apply_patch_to_context(fix, context)
            if patched_code is None:
                return TestRunResult(
                    passed=False,
                    output="Failed to apply patch",
                    error_summary="Patch application failed",
                )

            # Write patched code to temp directory
            source_path = work_dir / "source.py"
            source_path.write_text(patched_code, encoding="utf-8")

            # Wrap and write the test file
            test_code = self._wrap_test(context)
            test_path = work_dir / "test_patched.py"
            test_path.write_text(test_code, encoding="utf-8")

            # Write conftest if needed for fixtures
            conftest_code = self._generate_conftest(context)
            if conftest_code:
                conftest_path = work_dir / "conftest.py"
                conftest_path.write_text(conftest_code, encoding="utf-8")

            # Run pytest with replay support
            result = await self._run_pytest_with_replay(
                work_dir=work_dir,
                test_file=test_path,
                timeout=self.test_timeout,
            )

            return result

        except Exception as e:
            logger.exception("Error running test: %s", e)
            return TestRunResult(
                passed=False,
                output=str(e),
                error_summary=f"Execution error: {type(e).__name__}: {e}",
            )
        finally:
            # Clean up on success if configured
            if self.cleanup_on_success and work_dir and work_dir.exists():
                try:
                    shutil.rmtree(work_dir)
                    self._work_dirs.remove(work_dir)
                except Exception as cleanup_error:
                    logger.warning("Failed to clean up work dir: %s", cleanup_error)

    async def _verify_determinism(
        self,
        fix: FixHypothesis,
        context: FailureContext,
    ) -> VerificationResult:
        """Run test multiple times in parallel to verify determinism.

        Runs the test N times (default 10) using asyncio.gather for parallel execution.
        Tracks pass/fail counts and captures failure patterns.

        Args:
            fix: Fix to verify.
            context: Failure context.

        Returns:
            VerificationResult with statistics and failure patterns.
        """
        work_dir = None
        run_details: List[Dict[str, Any]] = []
        failure_patterns: List[str] = []

        try:
            # Create a single working directory for all verification runs
            work_dir = Path(tempfile.mkdtemp(prefix="hikaflow_verify_"))
            self._work_dirs.append(work_dir)

            # Apply the patch once
            patched_code = self._apply_patch_to_context(fix, context)
            if patched_code is None:
                return VerificationResult(
                    total_runs=0,
                    passed_runs=0,
                    failed_runs=0,
                    is_deterministic=False,
                    failure_patterns=["Failed to apply patch"],
                    run_details=[],
                )

            # Write patched source
            source_path = work_dir / "source.py"
            source_path.write_text(patched_code, encoding="utf-8")

            # Write test file
            test_code = self._wrap_test(context)
            test_path = work_dir / "test_verify.py"
            test_path.write_text(test_code, encoding="utf-8")

            # Write conftest
            conftest_code = self._generate_conftest(context)
            if conftest_code:
                conftest_path = work_dir / "conftest.py"
                conftest_path.write_text(conftest_code, encoding="utf-8")

            # Run tests in parallel with semaphore for concurrency control
            semaphore = asyncio.Semaphore(self.parallel_runs)

            async def run_single(run_number: int) -> TestRunResult:
                async with semaphore:
                    return await self._run_pytest_with_replay(
                        work_dir=work_dir,
                        test_file=test_path,
                        timeout=self.test_timeout,
                        run_number=run_number,
                    )

            # Execute all runs in parallel
            tasks = [run_single(i + 1) for i in range(self.verification_runs)]
            results = await asyncio.gather(*tasks, return_exceptions=True)

            # Process results
            passed_count = 0
            failed_count = 0

            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    failed_count += 1
                    failure_patterns.append(f"Run {i+1}: {type(result).__name__}: {result}")
                    run_details.append({
                        "run_number": i + 1,
                        "passed": False,
                        "error": str(result),
                    })
                elif isinstance(result, TestRunResult):
                    if result.passed:
                        passed_count += 1
                    else:
                        failed_count += 1
                        if result.error_summary:
                            failure_patterns.append(f"Run {i+1}: {result.error_summary}")
                    run_details.append({
                        "run_number": i + 1,
                        "passed": result.passed,
                        "output": result.output[:500] if result.output else "",
                        "error_summary": result.error_summary,
                    })
                else:
                    failed_count += 1
                    run_details.append({
                        "run_number": i + 1,
                        "passed": False,
                        "error": "Unknown result type",
                    })

            # Determine if deterministic (all passes or all consistent failures)
            unique_patterns = set(failure_patterns)
            is_deterministic = (
                passed_count == self.verification_runs or
                (failed_count == self.verification_runs and len(unique_patterns) <= 1)
            )

            return VerificationResult(
                total_runs=self.verification_runs,
                passed_runs=passed_count,
                failed_runs=failed_count,
                is_deterministic=is_deterministic and passed_count == self.verification_runs,
                failure_patterns=list(unique_patterns)[:10],  # Limit patterns
                run_details=run_details,
            )

        except Exception as e:
            logger.exception("Error during verification: %s", e)
            return VerificationResult(
                total_runs=self.verification_runs,
                passed_runs=0,
                failed_runs=self.verification_runs,
                is_deterministic=False,
                failure_patterns=[f"Verification error: {e}"],
                run_details=run_details,
            )
        finally:
            if self.cleanup_on_success and work_dir and work_dir.exists():
                try:
                    shutil.rmtree(work_dir)
                    self._work_dirs.remove(work_dir)
                except Exception as cleanup_error:
                    logger.warning("Failed to clean up verify dir: %s", cleanup_error)

    async def _run_pytest_with_replay(
        self,
        work_dir: Path,
        test_file: Path,
        timeout: int,
        run_number: int = 1,
    ) -> TestRunResult:
        """Run pytest with replay patches and proper isolation.

        Sets up the replay environment if transcripts are available,
        then runs pytest in a subprocess with timeout handling.

        Args:
            work_dir: Working directory containing the test.
            test_file: Path to the test file.
            timeout: Timeout in seconds.
            run_number: Run number for logging/identification.

        Returns:
            TestRunResult with actual execution results.
        """
        start_time = time.time()

        try:
            # Build environment
            env = os.environ.copy()
            env["PYTHONPATH"] = str(work_dir)

            # Configure replay if available
            if self.use_replay and self.transcripts_dir and self.transcripts_dir.exists():
                env["HIKAFLOW_REPLAY"] = "1"
                env["HIKAFLOW_REPLAY_DIR"] = str(self.transcripts_dir)
                # Add replay bootstrap to conftest or pytest ini
                replay_conftest = work_dir / "replay_bootstrap.py"
                replay_bootstrap_code = self._generate_replay_bootstrap()
                replay_conftest.write_text(replay_bootstrap_code, encoding="utf-8")

            # Build pytest command
            pytest_args = [
                sys.executable, "-m", "pytest",
                str(test_file),
                "-v",
                "--tb=short",
                "-x",  # Stop on first failure
                "--no-header",
                "-q",  # Quiet output
            ]

            # Run pytest subprocess
            process = await asyncio.create_subprocess_exec(
                *pytest_args,
                cwd=work_dir,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env=env,
            )

            try:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(),
                    timeout=timeout,
                )
            except asyncio.TimeoutError:
                process.kill()
                await process.communicate()
                duration = time.time() - start_time
                return TestRunResult(
                    passed=False,
                    output=f"Test timed out after {timeout}s",
                    error_summary=f"Timeout after {timeout}s",
                    exit_code=-1,
                    duration=duration,
                    run_number=run_number,
                )

            duration = time.time() - start_time
            stdout = stdout_bytes.decode("utf-8", errors="replace")
            stderr = stderr_bytes.decode("utf-8", errors="replace")
            combined_output = stdout + "\n" + stderr

            # Parse test result
            passed = process.returncode == 0
            error_summary = ""

            if not passed:
                # Extract error summary from output
                error_summary = self._extract_error_summary(combined_output)

            return TestRunResult(
                passed=passed,
                output=combined_output[:5000],  # Limit output size
                error_summary=error_summary,
                exit_code=process.returncode,
                duration=duration,
                run_number=run_number,
            )

        except Exception as e:
            duration = time.time() - start_time
            return TestRunResult(
                passed=False,
                output=str(e),
                error_summary=f"Subprocess error: {type(e).__name__}: {e}",
                exit_code=-1,
                duration=duration,
                run_number=run_number,
            )

    def _wrap_test(self, context: FailureContext) -> str:
        """Wrap test code with proper imports and module path setup.

        Creates a standalone test file that can run in isolation.

        Args:
            context: Failure context containing test information.

        Returns:
            Complete test code with imports.
        """
        # Extract test code from context
        test_code = context.code_snippet or ""
        surrounding_code = context.surrounding_code or ""

        # Determine what to use - prefer surrounding_code if it has the test
        if context.function_name and context.function_name in surrounding_code:
            code_to_wrap = surrounding_code
        else:
            code_to_wrap = test_code if test_code else surrounding_code

        # Generate imports based on code analysis
        imports = self._infer_imports(code_to_wrap)

        # Build the test file
        lines = [
            '"""Auto-generated test wrapper for fix verification."""',
            "",
            "import sys",
            "from pathlib import Path",
            "",
            "# Add work directory to path for imports",
            "sys.path.insert(0, str(Path(__file__).parent))",
            "",
        ]

        # Add inferred imports
        for imp in imports:
            lines.append(imp)

        lines.append("")
        lines.append("# Import the patched source module")
        lines.append("import source")
        lines.append("")

        # Add the test code
        lines.append("# Test code")
        lines.append(code_to_wrap)
        lines.append("")

        return "\n".join(lines)

    def _generate_conftest(self, context: FailureContext) -> Optional[str]:
        """Generate conftest.py content for pytest fixtures if needed.

        Args:
            context: Failure context.

        Returns:
            Conftest content or None if not needed.
        """
        # Basic conftest with common fixtures
        conftest_lines = [
            '"""Auto-generated conftest for fix verification."""',
            "",
            "import pytest",
            "import sys",
            "from pathlib import Path",
            "",
            "# Ensure source module is importable",
            "sys.path.insert(0, str(Path(__file__).parent))",
            "",
        ]

        # Add replay setup if enabled
        if self.use_replay:
            conftest_lines.extend([
                "@pytest.fixture(autouse=True)",
                "def setup_replay(request):",
                '    """Set up replay environment for test isolation."""',
                "    import os",
                '    if os.environ.get("HIKAFLOW_REPLAY"):',
                "        try:",
                "            from autodebug_replay import runtime",
                "            from pathlib import Path",
                '            transcripts_dir = Path(os.environ.get("HIKAFLOW_REPLAY_DIR", ""))',
                "            if transcripts_dir.exists():",
                "                runtime.initialize(transcripts_dir)",
                "        except ImportError:",
                "            pass  # Replay not available",
                "    yield",
                "",
            ])

        return "\n".join(conftest_lines)

    def _generate_replay_bootstrap(self) -> str:
        """Generate replay bootstrap code.

        Returns:
            Python code for initializing replay.
        """
        return '''"""Replay bootstrap for test isolation."""
import os
import sys

def init_replay():
    """Initialize replay infrastructure."""
    replay_dir = os.environ.get("HIKAFLOW_REPLAY_DIR")
    if replay_dir and os.path.exists(replay_dir):
        try:
            from autodebug_replay import runtime
            from pathlib import Path
            runtime.initialize(Path(replay_dir))
        except ImportError as e:
            print(f"Warning: Could not initialize replay: {e}", file=sys.stderr)
        except Exception as e:
            print(f"Warning: Replay initialization failed: {e}", file=sys.stderr)

# Auto-initialize on import
if os.environ.get("HIKAFLOW_REPLAY") == "1":
    init_replay()
'''

    def _extract_target_file(self, diff: str) -> Optional[str]:
        """Extract the target file path from a unified diff.

        Args:
            diff: Unified diff string.

        Returns:
            Target file path or None if not found.
        """
        # Look for +++ line in unified diff format
        match = re.search(r'^\+\+\+ [ab]/(.+)$', diff, re.MULTILINE)
        if match:
            return match.group(1)

        # Try --- line as fallback
        match = re.search(r'^--- [ab]/(.+)$', diff, re.MULTILINE)
        if match:
            return match.group(1)

        return None

    def _apply_patch_to_context(
        self,
        fix: FixHypothesis,
        context: FailureContext,
    ) -> Optional[str]:
        """Apply a fix's diff to the context's original code.

        Args:
            fix: Fix containing the diff.
            context: Failure context with original code.

        Returns:
            Patched code or None on failure.
        """
        # Use the patched_code if already available
        if fix.patched_code:
            return fix.patched_code

        # Try to apply the diff to surrounding code
        original = context.surrounding_code or context.code_snippet or ""
        if not original:
            # If no original code, extract from diff
            return self._extract_code_from_diff(fix.diff)

        try:
            return PatchApplicator._apply_simple_patch(original, fix.diff)
        except Exception as e:
            logger.warning("Failed to apply patch: %s", e)
            # Fallback to extracting new code from diff
            return self._extract_code_from_diff(fix.diff)

    def _infer_imports(self, code: str) -> List[str]:
        """Infer required imports from code.

        Args:
            code: Python code to analyze.

        Returns:
            List of import statements.
        """
        imports = []

        # Common test imports
        imports.append("import pytest")

        # Check for common patterns
        if "asyncio" in code or "async def" in code:
            imports.append("import asyncio")

        if "datetime" in code:
            imports.append("from datetime import datetime, timedelta")

        if "json" in code:
            imports.append("import json")

        if "time.sleep" in code or "time." in code:
            imports.append("import time")

        if "mock" in code.lower() or "Mock" in code:
            imports.append("from unittest.mock import Mock, patch, MagicMock")

        if "re." in code or "re.search" in code:
            imports.append("import re")

        if "Path" in code:
            imports.append("from pathlib import Path")

        if "os." in code:
            imports.append("import os")

        return imports

    def _extract_error_summary(self, output: str) -> str:
        """Extract a concise error summary from test output.

        Args:
            output: Full test output.

        Returns:
            Concise error summary.
        """
        lines = output.split("\n")
        error_lines = []

        # Look for assertion errors, exceptions, and failure summaries
        in_error = False
        for line in lines:
            line_lower = line.lower()
            if any(pattern in line_lower for pattern in [
                "error:", "exception:", "failed", "assert",
                "traceback", "raise", "assertionerror"
            ]):
                in_error = True

            if in_error:
                error_lines.append(line)
                if len(error_lines) >= 10:  # Limit error lines
                    break

            # Reset after blank line
            if in_error and not line.strip():
                in_error = False

        if error_lines:
            return "\n".join(error_lines[:10])

        # Fallback: return last few lines
        return "\n".join(lines[-5:]) if lines else "Unknown error"

    async def _refine_for_syntax(
        self,
        fix: FixHypothesis,
        errors: str,
        history: List[RefinementResult],
    ) -> FixHypothesis:
        """Refine fix to address syntax errors.

        Args:
            fix: Current fix.
            errors: Syntax errors.
            history: Previous refinement attempts.

        Returns:
            Refined fix.
        """
        if not self.instructor_client:
            return fix

        prompt = f"""
The fix has a syntax error. Please correct it.

Current fix:
```diff
{sanitize_code(fix.diff)}
```

Syntax error:
{sanitize_error(errors)}

Previous attempts:
{self._format_history(history)}

Correct the syntax error while preserving the fix intent.
"""

        try:
            response = await self._llm_generate_with_retry(
                prompt,
                response_model=FixResponse,
            )

            return FixHypothesis(
                id=fix.id,
                strategy=fix.strategy,
                diff=response.diff,
                description=response.description,
                confidence=response.confidence,
                iteration=fix.iteration + 1,
                refinement_reason=f"Syntax fix: {errors[:100]}",
            )

        except Exception as e:
            logger.error(f"Syntax refinement failed: {e}")
            return fix

    async def _refine_for_types(
        self,
        fix: FixHypothesis,
        errors: str,
        history: List[RefinementResult],
    ) -> FixHypothesis:
        """Refine fix to address type errors.

        Args:
            fix: Current fix.
            errors: Type errors.
            history: Previous refinement attempts.

        Returns:
            Refined fix.
        """
        if not self.instructor_client:
            return fix

        prompt = f"""
The fix has type errors. Please correct them.

Current fix:
```diff
{sanitize_code(fix.diff)}
```

Type errors:
{sanitize_error(errors)}

Previous attempts:
{self._format_history(history)}

Fix the type issues:
1. Add proper type handling/conversion
2. Check for None before accessing attributes
3. Ensure consistent types in operations
"""

        try:
            response = await self._llm_generate_with_retry(
                prompt,
                response_model=FixResponse,
            )

            return FixHypothesis(
                id=fix.id,
                strategy=fix.strategy,
                diff=response.diff,
                description=response.description,
                confidence=response.confidence,
                iteration=fix.iteration + 1,
                refinement_reason=f"Type fix: {errors[:100]}",
            )

        except Exception as e:
            logger.error(f"Type refinement failed: {e}")
            return fix

    async def _refine_for_test_failure(
        self,
        fix: FixHypothesis,
        test_result: TestRunResult,
        context: FailureContext,
        history: List[RefinementResult],
    ) -> FixHypothesis:
        """Refine fix based on test failure output.

        Args:
            fix: Current fix.
            test_result: Test execution result.
            context: Failure context.
            history: Previous refinement attempts.

        Returns:
            Refined fix.
        """
        if not self.instructor_client:
            return fix

        prompt = f"""
The fix attempt failed during testing. Refine it.

Original error: {sanitize_plain(context.error_type)}: {sanitize_plain(context.error_message)}
Root cause: {sanitize_plain(context.root_cause.summary) if context.root_cause else 'Unknown'}

Current fix attempt:
```diff
{sanitize_code(fix.diff)}
```

Test failure output:
{sanitize_error(test_result.output[:2000])}

Previous refinement history:
{self._format_history(history)}

Analyze why the fix didn't work and generate an improved version.
DO NOT repeat previous failed approaches.
"""

        try:
            response = await self._llm_generate_with_retry(
                prompt,
                response_model=FixResponse,
            )

            return FixHypothesis(
                id=fix.id,
                strategy=fix.strategy,
                diff=response.diff,
                description=response.description,
                confidence=response.confidence,
                iteration=fix.iteration + 1,
                refinement_reason=f"Test failed: {test_result.error_summary[:100]}",
            )

        except Exception as e:
            logger.error(f"Test refinement failed: {e}")
            return fix

    async def _refine_for_flakiness(
        self,
        fix: FixHypothesis,
        verify_result: VerificationResult,
        context: FailureContext,
        history: List[RefinementResult],
    ) -> FixHypothesis:
        """Refine fix that is still flaky.

        Args:
            fix: Current fix.
            verify_result: Verification result showing flakiness.
            context: Failure context.
            history: Previous refinement attempts.

        Returns:
            Refined fix.
        """
        if not self.instructor_client:
            return fix

        prompt = f"""
The fix is still flaky (test passed {verify_result.passed_runs}/{verify_result.total_runs} times).

Original error: {sanitize_plain(context.error_type)}
Root cause category: {context.root_cause.category.value if context.root_cause else 'Unknown'}

Current fix:
```diff
{sanitize_code(fix.diff)}
```

Failure patterns observed:
{sanitize_error(chr(10).join(verify_result.failure_patterns[:5]))}

The fix addresses the issue sometimes but not consistently. This suggests:
1. The fix may be incomplete
2. There may be additional race conditions
3. The fix may introduce new flakiness

Strengthen the fix to be deterministically correct.
"""

        try:
            response = await self._llm_generate_with_retry(
                prompt,
                response_model=FixResponse,
            )

            return FixHypothesis(
                id=fix.id,
                strategy=fix.strategy,
                diff=response.diff,
                description=response.description,
                confidence=response.confidence,
                iteration=fix.iteration + 1,
                refinement_reason=f"Flaky: {verify_result.passed_runs}/{verify_result.total_runs}",
            )

        except Exception as e:
            logger.error(f"Flakiness refinement failed: {e}")
            return fix

    def _extract_code_from_diff(self, diff: str) -> str:
        """Extract new code from a diff.

        Args:
            diff: Diff string.

        Returns:
            Extracted code.
        """
        lines = []
        for line in diff.split("\n"):
            if line.startswith("+") and not line.startswith("+++"):
                lines.append(line[1:])  # Remove +
            elif not line.startswith("-"):
                lines.append(line)

        return "\n".join(lines)

    def _format_history(self, history: List[RefinementResult]) -> str:
        """Format refinement history for prompt.

        Args:
            history: Refinement history.

        Returns:
            Formatted history.
        """
        if not history:
            return "No previous attempts."

        parts = []
        for i, result in enumerate(history[-3:]):  # Last 3 attempts
            parts.append(
                f"Attempt {i+1} ({result.stage.value}): {result.reasoning}"
            )

        return "\n".join(parts)


class ValidationResult:
    """Result of a validation check."""

    def __init__(self, valid: bool, errors: str):
        self.valid = valid
        self.errors = errors


class TestRunResult:
    """Result of running a test."""

    def __init__(
        self,
        passed: bool,
        output: str,
        error_summary: str,
        exit_code: int = 0,
        duration: float = 0.0,
        run_number: int = 1,
        divergence_code: Optional[str] = None,
    ):
        self.passed = passed
        self.output = output
        self.error_summary = error_summary
        self.exit_code = exit_code
        self.duration = duration
        self.run_number = run_number
        self.divergence_code = divergence_code

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "passed": self.passed,
            "output": self.output,
            "error_summary": self.error_summary,
            "exit_code": self.exit_code,
            "duration": self.duration,
            "run_number": self.run_number,
            "divergence_code": self.divergence_code,
        }


class PatchApplicator:
    """Utility to apply patches to code."""

    @staticmethod
    def apply_patch(
        original_file: Path,
        diff: str,
        output_dir: Optional[Path] = None,
    ) -> Path:
        """Apply a diff patch to a file.

        Args:
            original_file: Original file to patch.
            diff: Diff to apply.
            output_dir: Directory for patched file.

        Returns:
            Path to patched file.
        """
        if output_dir is None:
            output_dir = Path(tempfile.mkdtemp())

        output_file = output_dir / original_file.name

        # Read original
        content = original_file.read_text()

        # Apply patch (simplified - in practice would use proper patch tool)
        patched = PatchApplicator._apply_simple_patch(content, diff)

        output_file.write_text(patched)
        return output_file

    @staticmethod
    def _apply_simple_patch(content: str, diff: str) -> str:
        """Apply unified diff format to content.

        Properly handles unified diff format with:
        - Header lines (---, +++)
        - Hunk headers (@@ -start,count +start,count @@)
        - Context lines (no prefix or space prefix)
        - Removed lines (- prefix)
        - Added lines (+ prefix)

        Args:
            content: Original content.
            diff: Unified diff to apply.

        Returns:
            Patched content.

        Raises:
            ValueError: If diff cannot be applied.
        """
        lines = content.split("\n")
        diff_lines = diff.split("\n")
        result_lines = lines.copy()

        # Parse hunks from diff
        hunks = PatchApplicator._parse_unified_diff(diff_lines)

        # Apply hunks in reverse order to maintain line number accuracy
        for hunk in reversed(hunks):
            start_line = hunk["start"]
            old_lines = hunk["old_lines"]
            new_lines = hunk["new_lines"]

            # Validate context (optional, helps catch mismatches)
            if start_line < 0:
                start_line = 0
            if start_line > len(result_lines):
                # Append mode
                result_lines.extend(new_lines)
                continue

            # Find the actual start position by matching context
            actual_start = PatchApplicator._find_hunk_position(
                result_lines, old_lines, start_line
            )

            if actual_start is None:
                # Fallback: try fuzzy matching
                actual_start = PatchApplicator._fuzzy_find_position(
                    result_lines, old_lines
                )

            if actual_start is None:
                # Last resort: use the specified position
                actual_start = max(0, start_line - 1)  # Convert to 0-indexed

            # Replace the old lines with new lines
            end_idx = actual_start + len(old_lines)
            result_lines = result_lines[:actual_start] + new_lines + result_lines[end_idx:]

        return "\n".join(result_lines)

    @staticmethod
    def _parse_unified_diff(diff_lines: List[str]) -> List[Dict[str, Any]]:
        """Parse unified diff into hunks.

        Args:
            diff_lines: Lines of the diff.

        Returns:
            List of hunk dictionaries with start, old_lines, new_lines.
        """
        hunks = []
        current_hunk = None

        for line in diff_lines:
            # Skip header lines
            if line.startswith("---") or line.startswith("+++"):
                continue

            # Hunk header
            if line.startswith("@@"):
                # Save previous hunk
                if current_hunk:
                    hunks.append(current_hunk)

                # Parse hunk header: @@ -start,count +start,count @@
                match = re.match(r"@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@", line)
                if match:
                    old_start = int(match.group(1))
                    new_start = int(match.group(2))
                else:
                    old_start = 1
                    new_start = 1

                current_hunk = {
                    "start": old_start,
                    "old_lines": [],
                    "new_lines": [],
                }
                continue

            if current_hunk is None:
                # No hunk started yet, skip
                continue

            # Process diff lines
            if line.startswith("-"):
                # Line to remove (include in old_lines for matching)
                current_hunk["old_lines"].append(line[1:])
            elif line.startswith("+"):
                # Line to add
                current_hunk["new_lines"].append(line[1:])
            elif line.startswith(" ") or line == "":
                # Context line (present in both old and new)
                context_line = line[1:] if line.startswith(" ") else line
                current_hunk["old_lines"].append(context_line)
                current_hunk["new_lines"].append(context_line)
            else:
                # No prefix - treat as context
                current_hunk["old_lines"].append(line)
                current_hunk["new_lines"].append(line)

        # Save final hunk
        if current_hunk:
            hunks.append(current_hunk)

        return hunks

    @staticmethod
    def _find_hunk_position(
        lines: List[str],
        old_lines: List[str],
        suggested_start: int,
    ) -> Optional[int]:
        """Find the position in lines where old_lines match.

        Args:
            lines: Content lines to search.
            old_lines: Lines to find.
            suggested_start: Suggested starting position (1-indexed).

        Returns:
            0-indexed position or None if not found.
        """
        if not old_lines:
            return max(0, suggested_start - 1)

        # Convert to 0-indexed and search in a window around suggested position
        suggested_idx = max(0, suggested_start - 1)
        search_window = 20  # Search up to 20 lines in each direction

        # Try exact match first at suggested position
        if PatchApplicator._lines_match(lines, suggested_idx, old_lines):
            return suggested_idx

        # Search in expanding window
        for offset in range(1, search_window + 1):
            # Try before
            pos = suggested_idx - offset
            if pos >= 0 and PatchApplicator._lines_match(lines, pos, old_lines):
                return pos

            # Try after
            pos = suggested_idx + offset
            if pos < len(lines) and PatchApplicator._lines_match(lines, pos, old_lines):
                return pos

        return None

    @staticmethod
    def _fuzzy_find_position(
        lines: List[str],
        old_lines: List[str],
    ) -> Optional[int]:
        """Fuzzy search for hunk position using first line matching.

        Args:
            lines: Content lines to search.
            old_lines: Lines to find.

        Returns:
            0-indexed position or None if not found.
        """
        if not old_lines:
            return 0

        first_line = old_lines[0].strip()
        if not first_line:
            return None

        # Search for first line
        for i, line in enumerate(lines):
            if line.strip() == first_line:
                # Verify remaining lines match
                if PatchApplicator._lines_match(lines, i, old_lines):
                    return i

        return None

    @staticmethod
    def _lines_match(
        content_lines: List[str],
        start_idx: int,
        expected_lines: List[str],
    ) -> bool:
        """Check if lines match at a given position.

        Args:
            content_lines: Content to check.
            start_idx: Starting index.
            expected_lines: Lines to match.

        Returns:
            True if all lines match.
        """
        if start_idx + len(expected_lines) > len(content_lines):
            return False

        for i, expected in enumerate(expected_lines):
            actual = content_lines[start_idx + i]
            # Strip for comparison to handle whitespace differences
            if actual.strip() != expected.strip():
                return False

        return True
