import os
import re
from google import genai
from google.genai import types as genai_types
import logging
from .gemini import SYSTEM_PROMPT, base_prompt_instructions

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def fix_manim_code(faulty_code: str, error_message: str, original_context: str):
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        logging.error("GEMINI_API_KEY not found in environment variables for fallback.")
        return None, None

    client = genai.Client(api_key=api_key)

    fix_prompt_text = (
        f"The following Manim code, intended to '{original_context}', failed with an error.\n\n"
        "### FAULTY CODE:\n"
        f"```python\n{faulty_code}\n```\n\n"
        "### ERROR MESSAGE:\n"
        f"```\n{error_message}\n```\n\n"
        "### INSTRUCTIONS:\n"
        "1. Analyze the error message and the faulty code.\n"
        "2. Correct the code to fix the specific error reported.\n"
        "3. Ensure the corrected code still fulfills the original request and adheres strictly to *all* the requirements listed below.\n"
        "4. Pay close attention to vector dimensions, matrix operations, allowed Manim methods, and total duration (30 seconds).\n"
        "5. If the code logic changes significantly, update the narration accordingly.\n"
        "6. Return *only* the corrected code and narration using the '### MANIM CODE:' and '### NARRATION:' delimiters, just like the original request.\n\n"
        "### REQUIREMENTS (Apply these to the corrected code):\n"
        f"{base_prompt_instructions}" 
    )

    contents = [fix_prompt_text] 

    logging.info("Attempting to fix Manim code via fallback...")
    try:
        generation_config = genai_types.GenerateContentConfig(
            system_instruction=SYSTEM_PROMPT
        )
        
        response = client.models.generate_content(
                model="gemini-2.0-flash",
                contents=contents,
                config=generation_config
                )
        if response:
            try:
                content = response.text
                logging.info("Received response from fallback attempt.")
                
                if "### NARRATION:" in content:
                    manim_code, narration = content.split("### NARRATION:", 1)
                    manim_code = re.sub(r"```python", "", manim_code).replace("```", "").strip()
                    narration = narration.strip()

                    if "from manim import *" not in manim_code:
                         logging.warning("Adding missing 'from manim import *' (fallback fix).")
                         manim_code = "from manim import *\nimport numpy as np\n" + manim_code
                    elif "import numpy as np" not in manim_code:
                         logging.warning("Adding missing 'import numpy as np' (fallback fix).")
                         lines = manim_code.splitlines()
                         for i, line in enumerate(lines):
                             if "from manim import *" in line:
                                 lines.insert(i + 1, "import numpy as np")
                                 manim_code = "\n".join(lines)
                                 break

                    logging.info("Successfully parsed fixed code and narration from fallback.")
                    return {"manim_code": manim_code, "output_file": "output.mp4"}, narration
                else:
                    logging.warning("Delimiter '### NARRATION:' not found in fallback response. Attempting fallback extraction.")
                    code_match = re.search(r'```python(.*?)```', content, re.DOTALL)
                    if code_match:
                        manim_code = code_match.group(1).strip()
                        narration_part = content.split('```', 2)[-1].strip()
                        narration = narration_part if len(narration_part) > 20 else ""
                        if not narration:
                            logging.warning("Fallback narration extraction resulted in empty or very short text (fallback fix).")
                        else:
                            logging.info("Successfully parsed code and narration using fallback regex (fallback fix).")

                        if "from manim import *" not in manim_code:
                             logging.warning("Adding missing 'from manim import *' (fallback fix, regex path).")
                             manim_code = "from manim import *\nimport numpy as np\n" + manim_code
                        elif "import numpy as np" not in manim_code:
                             logging.warning("Adding missing 'import numpy as np' (fallback fix, regex path).")
                             lines = manim_code.splitlines()
                             for i, line in enumerate(lines):
                                 if "from manim import *" in line:
                                     lines.insert(i + 1, "import numpy as np")
                                     manim_code = "\n".join(lines)
                                     break

                        logging.info("Successfully parsed fixed code using fallback extraction.")
                        return {"manim_code": manim_code, "output_file": "output.mp4"}, narration
                    else:
                         logging.error("Fallback extraction failed: No Python code block found in fallback response.")
                         logging.debug(f"Fallback content without code block:\n{content}")
                         return None, None

            except ValueError:
                logging.error("Could not extract text from the fallback response.")
                if response.prompt_feedback and response.prompt_feedback.block_reason:
                     logging.error(f"Fallback content generation blocked. Reason: {response.prompt_feedback.block_reason.name}")
                return None, None
            except Exception as e:
                 logging.exception(f"Error processing fallback response: {e}")
                 return None, None
        else:
            logging.error("No response received from Gemini during fallback attempt.")
            return None, None

    except Exception as e:
        logging.exception(f"Error calling Gemini API during fallback: {e}")
        return None, None