"""Self-Healing QA Loop for NIKAME.

Runs linting (ruff) and syntax checks on AI-generated Python code.
If errors are found, the tracebacks are fed back into the AISynthesizer
to automatically patch the broken files before they are written to disk.
"""

from __future__ import annotations

import subprocess
import sys
import tempfile
from pathlib import Path

from nikame.codegen.ai_synthesizer import AISynthesizer
from nikame.utils.logger import console, get_logger

_log = get_logger("qa_loop")


class QALoop:
    """Manages the verification and auto-correction of generated code."""

    def __init__(self, max_retries: int = 3) -> None:
        self.max_retries = max_retries
        self.ai = AISynthesizer()

    def run_checks(self, file_buffer: dict[str, str]) -> dict[str, str]:
        """Run QA checks and heal broken files. Returns the modified buffer."""
        py_files = {k: v for k, v in file_buffer.items() if k.endswith(".py")}
        if not py_files:
            return file_buffer

        attempts = 0
        while attempts < self.max_retries:
            failed_files = self._execute_check(py_files)
            if not failed_files:
                console.print("[success]✓ AI-generated code passed all syntax and lint checks.[/success]")
                break
                
            attempts += 1
            console.print(f"[warning]⚠ QA Loop detected errors in {len(failed_files)} files. Attempting auto-heal (Attempt {attempts}/{self.max_retries})...[/warning]")
            
            for file_path, error_msg in failed_files.items():
                original_code = py_files[file_path]
                fixed_code = self._heal_file(file_path, original_code, error_msg)
                py_files[file_path] = fixed_code
                file_buffer[file_path] = fixed_code

        if attempts == self.max_retries and failed_files:
            console.print(f"[bold red]❌ QA Loop exhausted. Generated code may contain errors.[/bold red]")

        return file_buffer

    def _execute_check(self, py_files: dict[str, str]) -> dict[str, str]:
        """Detect syntax errors and fatal linters. Returns {filepath: error_message}."""
        failed_files = {}
        
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = Path(tmpdir)
            
            # Write files to tempdir
            reverse_map = {}
            for rel_path, content in py_files.items():
                target = tmp_path / rel_path
                target.parent.mkdir(parents=True, exist_ok=True)
                target.write_text(content, encoding="utf-8")
                reverse_map[str(target)] = rel_path
                
            # Run basic syntax compile check
            for rel_path, content in py_files.items():
                target = tmp_path / rel_path
                result = subprocess.run(
                    [sys.executable, "-m", "py_compile", str(target)],
                    capture_output=True,
                    text=True,
                )
                if result.returncode != 0:
                    err = result.stderr.strip()
                    failed_files[rel_path] = f"SyntaxError: {err}"
                    
            # If no syntax errors, run ruff (focus on undefined names, imports, etc.)
            if not failed_files:
                try:
                    result = subprocess.run(
                        ["ruff", "check", "--select", "F821,F822,E999", "."],
                        cwd=tmpdir,
                        capture_output=True,
                        text=True
                    )
                    if result.returncode != 0:
                        # Parse ruff output
                        for line in result.stdout.splitlines():
                            if ".py:" in line:
                                parts = line.split(":", 2)
                                if len(parts) >= 3:
                                    # tmpdir/app/main.py -> app/main.py
                                    raw_path = parts[0]
                                    rel_path = raw_path.replace(str(tmp_path) + "/", "").replace(str(tmp_path) + "\\", "")
                                    if rel_path in py_files:
                                        msg = parts[2].strip()
                                        if rel_path in failed_files:
                                            failed_files[rel_path] += f"\n{msg}"
                                        else:
                                            failed_files[rel_path] = msg
                except FileNotFoundError:
                    _log.debug("Ruff not installed, skipping advanced linting.")

        return failed_files

    def _heal_file(self, file_path: str, original_code: str, error_msg: str) -> str:
        """Use the AI Synthesizer to fix the broken file."""
        prompt = (
            f"You are NIKAME Copilot's automated QA agent. The generated file '{file_path}' "
            f"failed compilation with the following errors:\n\n{error_msg}\n\n"
            "Your task is to fix the errors and return the ENTIRE corrected Python file. "
            "Output ONLY the syntactically correct Python code. Do not apologize or explain."
        )
        return self.ai.synthesize_code(prompt, original_code)
