import groq, json, time, os, re
from utils.prompts import get_fixer_prompt, get_simplified_prompt

class FixerAgent:
    def __init__(self):
        self.client = groq.Groq(api_key=os.environ.get("GROQ_API_KEY"))
        self.model = "llama-3.1-8b-instant"
    
    def fix_code(self, buggy_code: str, error_type: str,
                 description: str, test_cases: list,
                 test_results: dict = None,
                 previous_explanation: str = None,
                 iteration: int = 1) -> dict:
        """
        3-tier fallback:
        Tier 1: Full LLM with JSON output
        Tier 2: Simplified LLM prompt  
        Tier 3: Heuristic rule-based fix
        
        Returns: {
          "fixed_code": str,
          "explanation": str,
          "confidence": float,
          "method": "llm"/"simplified"/"heuristic",
          "elapsed_seconds": float
        }
        """
        start = time.time()
        
        # Tier 1: Full LLM with JSON output
        try:
            prompt = get_fixer_prompt(
                buggy_code, error_type, description, test_cases, 
                test_results, previous_explanation, iteration
            )
            response_text = self._call_groq(prompt)
            parsed = self._parse_json_response(response_text)
            if parsed and "fixed_code" in parsed:
                parsed["method"] = "llm"
                parsed["confidence"] = 0.9
                parsed["elapsed_seconds"] = time.time() - start
                return parsed
        except Exception as e:
            print(f"[FixerAgent] Tier 1 failed: {e}")

        # Tier 2: Simplified LLM prompt
        try:
            simple_prompt = get_simplified_prompt(buggy_code, error_type)
            response_text = self._call_groq(simple_prompt)
            
            # extract python code block from response
            match = re.search(r"```(?:python)?\n?(.*?)\n?```", response_text, re.DOTALL)
            fixed_code = match.group(1).strip() if match else response_text.strip()
            
            return {
                "fixed_code": fixed_code,
                "explanation": "Simplified prompt fallback",
                "confidence": 0.6,
                "method": "simplified",
                "elapsed_seconds": time.time() - start
            }
        except Exception as e:
            print(f"[FixerAgent] Tier 2 failed: {e}")

        # Tier 3: Heuristic rule-based fix
        fixed = self._heuristic_fix(buggy_code, error_type)
        return {
            "fixed_code": fixed,
            "explanation": "Heuristic fallback rule",
            "confidence": 0.3,
            "method": "heuristic",
            "elapsed_seconds": time.time() - start
        }
    
    def _call_groq(self, prompt: str) -> str:
        """Call Groq API, return response text"""
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=1024,
            temperature=0.2
        )
        return response.choices[0].message.content
    
    def _parse_json_response(self, response: str) -> dict:
        """
        Parse JSON from LLM response.
        Try: json.loads(response)
        If fails: find { } block with regex and try again
        If fails: return None
        """
        def try_parse(s):
            try:
                res = json.loads(s)
                if isinstance(res, dict) and "fixed_code" in res:
                    return res
            except Exception:
                pass
            return None

        # 1. Direct parse
        res = try_parse(response.strip())
        if res: return res

        # 2. Extract JSON block inside markdown
        match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", response, re.DOTALL)
        if match:
            res = try_parse(match.group(1))
            if res: return res

        # 3. Raw { } block fallback
        match = re.search(r"(\{.*\})", response, re.DOTALL)
        if match:
            res = try_parse(match.group(1))
            if res: return res

        return None
    
    def _heuristic_fix(self, buggy_code: str, error_type: str) -> str:
        """
        Rule-based fixes for common patterns:
        IndexError: find lst[len(lst)] → change to lst[len(lst)-1]
        Missing return: add "return result" before last line
        Wrong operator: common >= vs > substitutions
        Returns modified code or original if no rule matches
        """
        if error_type == "IndexError":
            # Heuristic for off-by-one errors
            if "len(lst)]" in buggy_code:
                return buggy_code.replace("len(lst)]", "len(lst)-1]")
        
        # Missing return check
        if "return " not in buggy_code:
            lines = buggy_code.splitlines()
            for i in range(len(lines)-1, -1, -1):
                if lines[i].strip():
                    indent = len(lines[i]) - len(lines[i].lstrip())
                    lines[i] = " " * indent + "return " + lines[i].lstrip()
                    return "\n".join(lines)
                    
        # Common operator mistakes
        if ">" in buggy_code and ">=" not in buggy_code:
            return buggy_code.replace(">", ">=")
            
        return buggy_code