import os
from typing import List, Union
from dotenv import load_dotenv
import anthropic
import re
import subprocess
import base64
import tempfile
import filetype

from loops.core.prompted_image import PromptedImage
from loops.core.system_prompt import SystemPrompt


load_dotenv()

# Get system_prompts folder path
from pathlib import Path
_current_file = Path(__file__)
_system_prompts_folder = _current_file.parent.parent / "system_prompts"


class RawAdVisual:
    MODEL = "claude-sonnet-4-5"
    BETAS = [
        "code-execution-2025-08-25",
        "interleaved-thinking-2025-05-14"
    ]
    TOOLS = [{
        "type": "code_execution_20250825",
        "name": "code_execution"
    }]
    THINKING = {
        "type": "enabled",
        "budget_tokens": 10000
    }
    MAX_TOKENS = 32000
    SYSTEM_PROMPT_PATH = str(_system_prompts_folder / "fix_text.txt")
    TEMPORARY_SCRIPT_PATH = "/tmp/correct_image.py"

    def __init__(self, visual: Union[str, bytes]):
        """
        Initialize with image path or bytes.
        
        Args:
            visual: Path to image file OR bytes of image data
        """
        self.visual = visual
        self.client = anthropic.Anthropic()
        self.system_prompt = SystemPrompt(self.SYSTEM_PROMPT_PATH).read()
        self._is_bytes = isinstance(visual, bytes)
        self._temp_path = None

    def execute_ocr_correction(self) -> bytes:
        """
        Executes OCR correction and returns fixed image as bytes.
        
        Returns:
            bytes: Fixed image data ready for Cloudinary upload
        """
        # If we have bytes, save temporarily for OCR processing
        if self._is_bytes:
            # Create temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
                tmp_file.write(self.visual)
                self._temp_path = tmp_file.name
            image_path = self._temp_path
            # Generate fixed image path from temp file
            fixed_image_path = str(Path(self._temp_path).with_name(f"{Path(self._temp_path).stem}_fixed.jpg"))
        else:
            image_path = self.visual
            # Generate fixed image path from original path
            fixed_image_path = image_path.replace("_raw", "_fixed")
            if not fixed_image_path.endswith(('.jpg', '.jpeg')):
                fixed_image_path = str(Path(fixed_image_path).with_suffix('.jpg'))
        
        # Generate and execute OCR script
        self.save_ocr_script()
        
        # Use python3 for better portability
        command = [
            "python3",
            self.TEMPORARY_SCRIPT_PATH,
            image_path,
            fixed_image_path,
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",  # Default system fonts
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
        ]
        
        try:
            result = subprocess.run(command, check=True, capture_output=True, text=True)
        except subprocess.CalledProcessError as e:
            stderr = e.stderr or ""
            stdout = e.stdout or ""
            
            # Check if it's a Tesseract not found error (expected in some environments)
            if 'tesseract' in stderr.lower() or 'TesseractNotFoundError' in stderr:
                raise RuntimeError(
                    "Tesseract OCR is not installed. "
                    "The OCR script was generated correctly, but Tesseract is required to execute it. "
                    "Install with: brew install tesseract (macOS) or apt-get install tesseract-ocr (Linux)"
                ) from e
            
            if 'ModuleNotFoundError' in stderr or 'ImportError' in stderr:
                print(f"⚠️  The generated script uses a library that's not available.")
                print(f"   This should not happen - the script should only use PIL and pytesseract.")
            
            # Print error for debugging
            print(f"⚠️  Error executing OCR script:")
            print(f"   Command: {' '.join(command)}")
            print(f"   Return code: {e.returncode}")
            if stdout:
                print(f"   Stdout: {stdout[:500]}")
            if stderr:
                print(f"   Stderr: {stderr[:500]}")
            
            # Try to read the generated script for debugging
            try:
                with open(self.TEMPORARY_SCRIPT_PATH, 'r') as f:
                    script_content = f.read()
                print(f"\n   Generated script (first 500 chars):\n{script_content[:500]}")
            except:
                pass
            
            # Fallback: try with system python
            command[0] = "python"
            try:
                result = subprocess.run(command, check=True, capture_output=True, text=True)
            except subprocess.CalledProcessError as e2:
                stderr2 = e2.stderr or ""
                if 'tesseract' in stderr2.lower() or 'TesseractNotFoundError' in stderr2:
                    raise RuntimeError(
                        "Tesseract OCR is not installed. "
                        "The OCR script was generated correctly, but Tesseract is required to execute it."
                    ) from e2
                print(f"⚠️  Fallback also failed:")
                print(f"   Return code: {e2.returncode}")
                if e2.stdout:
                    print(f"   Stdout: {e2.stdout[:500]}")
                if e2.stderr:
                    print(f"   Stderr: {e2.stderr[:500]}")
                raise
        
        # Read fixed image as bytes
        with open(fixed_image_path, "rb") as f:
            fixed_image_bytes = f.read()
        
        # Cleanup temp files if we created them
        if self._is_bytes and self._temp_path:
            try:
                os.unlink(self._temp_path)
            except:
                pass
            try:
                os.unlink(fixed_image_path)
            except:
                pass
        
        return fixed_image_bytes

    def save_ocr_script(self) -> str:
        """Generate and save OCR correction script to temporary file."""
        raw_model_content = self.fix_text_with()
        merged_text = self._extract_text_blocks(raw_model_content)
        python_code = self._extract_python_code(merged_text)
        
        # Validate the extracted code
        if not python_code or len(python_code) < 50:
            raise ValueError(f"Extracted Python code is too short or empty. Got {len(python_code) if python_code else 0} characters.")
        
        if 'sys.argv' not in python_code:
            raise ValueError("Extracted code does not contain sys.argv. The script format may be incorrect.")
        
        # Final aggressive cleanup: remove any non-Python lines
        lines = python_code.split('\n')
        cleaned_lines = []
        python_keywords = ['import', 'from', 'def', 'class', 'if', 'for', 'while', 'try', 'except', 'return', 'sys.argv', '=', '(', ')', '[', ']', '{', '}']
        markdown_chars = ['✅', '✓', '❌', '⚠️', '🎉', '**', '##', '###', '```']
        
        for line in lines:
            stripped = line.strip()
            
            # Skip completely empty lines at the very start
            if not stripped and not cleaned_lines:
                continue
            
            # Skip lines with any markdown/emoji characters
            if any(char in line for char in markdown_chars):
                continue
            
            # Skip lines that are clearly explanatory text (long, no Python syntax)
            if len(stripped) > 80:
                has_python_syntax = any(kw in stripped for kw in python_keywords)
                if not has_python_syntax:
                    continue
            
            # Skip markdown list items and headers
            if (stripped.startswith(('- ', '* ', '#', '##', '###', '1. ', '2. ', '3. ', '4. ', '5. ')) and
                not any(kw in stripped for kw in python_keywords)):
                continue
            
            # Keep the line
            cleaned_lines.append(line)
        
        python_code = '\n'.join(cleaned_lines).strip()
        
        # Remove any leading/trailing non-code
        while python_code and not python_code.split('\n')[0].strip().startswith(('import ', 'from ', '#!/', 'def ', 'class ')):
            python_code = '\n'.join(python_code.split('\n')[1:]).strip()
        
        # Final validation
        if not python_code or len(python_code) < 50:
            raise ValueError(f"Cleaned Python code is too short. Got {len(python_code) if python_code else 0} characters.")
        
        if 'sys.argv' not in python_code:
            raise ValueError("Cleaned code does not contain sys.argv.")
        
        # Write the script
        with open(self.TEMPORARY_SCRIPT_PATH, "w", encoding="utf-8") as f:
            f.write(python_code)
        
        # Verify the script is valid Python syntax
        try:
            compile(python_code, self.TEMPORARY_SCRIPT_PATH, 'exec')
        except SyntaxError as e:
            print(f"⚠️  Syntax error in generated script:")
            print(f"   {e}")
            print(f"   First 500 chars of code:\n{python_code[:500]}")
            raise
        
        return python_code


    def fix_text_with(self):
        """Generate Python code for OCR text correction using Claude API."""
        content = self._construct_full_prompt()
        with self.client.beta.messages.stream(
            model=self.MODEL,
            betas=self.BETAS,
            messages=content,
            max_tokens=self.MAX_TOKENS,
            thinking=self.THINKING,
            tools=self.TOOLS
        ) as stream:
            for event in stream:
                if event.type == "content_block_start":
                    print("start:", event.content_block, flush=True)
                elif event.type == "content_block_delta":
                    print("delta:", event.delta, flush=True)
                elif event.type == "content_block_stop":
                    print("stop block")
                elif event.type == "message_delta":
                    print("message partial:", event.delta, flush=True)
                elif event.type == "message_stop":
                    print("message complete", flush=True)

            output = stream.get_final_message()
        return output

    def _detect_media_type_from_bytes(self, image_bytes: bytes) -> str:
        """Detect media type from image bytes."""
        try:
            # Try using filetype library
            kind = filetype.guess(image_bytes)
            if kind:
                return f"image/{kind.extension}" if kind.extension != "jpg" else "image/jpeg"
        except:
            pass
        
        # Fallback: check magic bytes
        if image_bytes.startswith(b'\xff\xd8\xff'):
            return "image/jpeg"
        elif image_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
            return "image/png"
        elif image_bytes.startswith(b'GIF87a') or image_bytes.startswith(b'GIF89a'):
            return "image/gif"
        elif image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:12]:
            return "image/webp"
        
        # Default to JPEG
        return "image/jpeg"

    def _construct_full_prompt(self) -> List:
        # Get image data for prompt
        if self._is_bytes:
            # Convert bytes to base64 for Anthropic API
            image_base64 = base64.b64encode(self.visual).decode("utf-8")
            media_type = self._detect_media_type_from_bytes(self.visual)
            image_dict = {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": media_type,
                    "data": image_base64,
                }
            }
        else:
            image_dict = self._format_raw_visual()
        
        content = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text", 
                        "text": """Generate ONLY Python code. Output ONLY code inside ```python block.

EXACT REQUIREMENTS:
The script MUST accept exactly 4 arguments via sys.argv:
- sys.argv[1] = input image path (string)
- sys.argv[2] = output image path (string)  
- sys.argv[3] = normal font file path (string)
- sys.argv[4] = bold font file path (string)

Example structure:
```python
import sys
import pytesseract
from PIL import Image, ImageDraw, ImageFont

input_path = sys.argv[1]
output_path = sys.argv[2]
font_normal_path = sys.argv[3]
font_bold_path = sys.argv[4]

img = Image.open(input_path)
# ... OCR and text fixing code using pytesseract and PIL ...
img.save(output_path)
```

CRITICAL: 
- NO argparse
- NO optional arguments
- sys.argv[3] and sys.argv[4] are FONT FILE PATHS (strings), not font sizes
- Use ONLY pytesseract and PIL/Pillow (NO numpy, NO other libraries)
- Use pytesseract.image_to_data() for OCR
- Fix ALL text in the image
- Output ONLY the Python code, no explanations"""
                    },
                    image_dict,
                ]
            }
        ]
        return content
    

    def _format_raw_visual(self) -> dict:
        """Format image from path for Anthropic API."""
        image = PromptedImage(self.visual)
        data = image.get_data_base64()
        media_type = image.guess_media_type()
        return {
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": media_type,
                "data": data,
            }
        }
    
    @staticmethod
    def _extract_text_blocks(message) -> str:
        """Extracts text blocks from the message content."""
        text_blocks = []
        for block in message.content:
            if hasattr(block, "type") and block.type == "text":
                text_blocks.append(block.text)
        merged_text = "\n".join(text_blocks)
        decoded_text = merged_text.replace("\\n", "\n") #replace \n with actual new lines
        return decoded_text
    
    @staticmethod
    def _extract_python_code(text: str) -> str:
        """Extracts python code from a markdown-formatted text block.
        
        Returns only valid Python code, filtering out markdown and explanatory text.
        """
        # Strategy 1: Find code blocks (most reliable)
        code_blocks = []
        
        # Try multiple patterns for code blocks - be very specific
        patterns = [
            r"```python\n(.*?)```",  # ```python ... ```
            r"```py\n(.*?)```",      # ```py ... ```
            r"```\n(.*?)```",        # ``` ... ```
        ]
        
        for pattern in patterns:
            matches = re.finditer(pattern, text, flags=re.DOTALL)
            for match in matches:
                code = match.group(1).strip()
                # Validate it's actually Python code - must have sys.argv or start with import/from
                if code and len(code) > 50:  # Must be substantial
                    # Check if it contains Python keywords
                    has_python = any(kw in code for kw in ['sys.argv', 'import ', 'from ', 'def ', 'class ', 'if __name__'])
                    # Check it doesn't have too many markdown artifacts
                    has_markdown = code.count('✅') > 2 or code.count('**') > 5 or code.count('#') > 20
                    
                    if has_python and not has_markdown:
                        code_blocks.append((len(code), code))
        
        # If we found valid code blocks, return the largest one
        if code_blocks:
            code_blocks.sort(reverse=True)  # Sort by length, largest first
            best_code = code_blocks[0][1]
            
            # Final cleanup: remove any remaining markdown artifacts
            lines = best_code.split('\n')
            cleaned_lines = []
            for line in lines:
                stripped = line.strip()
                # Skip lines that are clearly markdown/explanatory
                if (stripped.startswith(('✅', '**', '##', '###', '- ', '* ')) or 
                    (len(stripped) > 100 and not any(c in stripped for c in ['=', '(', ')', '[', ']', 'import', 'def', 'class']))):
                    continue
                cleaned_lines.append(line)
            
            cleaned = '\n'.join(cleaned_lines).strip()
            if cleaned and 'sys.argv' in cleaned:
                return cleaned
        
        # Strategy 2: Extract lines that look like Python code
        lines = text.split('\n')
        code_lines = []
        in_code = False
        
        for line in lines:
            stripped = line.strip()
            
            # Detect start of code
            if stripped.startswith(('import ', 'from ', '#!/', 'def ', 'class ')) or 'sys.argv' in stripped:
                in_code = True
            
            if in_code:
                # Skip markdown markers
                if stripped.startswith('```'):
                    continue
                # Skip explanatory lines (markdown formatting, emojis, long text without code)
                if (stripped.startswith(('✅', '**', '##', '###', '- ', '* ')) or
                    (len(stripped) > 150 and not any(c in stripped for c in ['=', '(', ')', '[', ']', 'import', 'def', 'class', 'sys.argv']))):
                    continue
                code_lines.append(line)
            
            # Stop if we hit end marker and have collected code
            if in_code and stripped.startswith('```') and len(code_lines) > 10:
                break
        
        if code_lines:
            code = '\n'.join(code_lines)
            # Remove any remaining markdown
            code = re.sub(r'^```.*$', '', code, flags=re.MULTILINE)
            code = code.strip()
            
            if code and 'sys.argv' in code and len(code) > 100:
                return code
        
        return ""

