"""
Monkey patches to enhance browser-use with vision sidecar capabilities.
These patches are applied automatically when the package is imported.
"""

import asyncio
import os
from typing import Optional


def apply_all_patches():
    """Apply all local vision enhancement patches to browser-use."""
    try:
        patch_config()
        patch_prompts()
        patch_message_manager()
        patch_agent_service()
        print("✅ Local vision patches applied successfully")
    except Exception as e:
        print(f"⚠️ Could not apply some patches: {e}")


def patch_config():
    """Add local vision config fields to browser-use config."""
    try:
        import browser_use.config as config

        # Add local vision environment variables to CONFIG
        class EnhancedConfig:
            def __init__(self, original_config):
                self._original = original_config

            def __getattr__(self, name):
                if name == 'LOCAL_VISION_CONFIDENCE_THRESHOLD':
                    return float(os.getenv('LOCAL_VISION_CONFIDENCE_THRESHOLD', '0.85'))
                elif name == 'LOCAL_VISION_ENABLED':
                    return os.getenv('LOCAL_VISION_ENABLED', 'true').lower() in ('true', '1', 'yes')
                else:
                    return getattr(self._original, name)

        # Enhance the existing CONFIG object
        if hasattr(config, 'CONFIG'):
            config.CONFIG = EnhancedConfig(config.CONFIG)

    except ImportError:
        pass


def patch_prompts():
    """Patch browser-use prompts to use async WebP resize."""
    try:
        import browser_use.agent.prompts as prompts

        async def _resize_screenshot_async(screenshot_b64: str) -> str:
            """Async version of screenshot resize with WebP format."""
            return await asyncio.to_thread(_sync_resize_webp, screenshot_b64)

        def _sync_resize_webp(screenshot_b64: str) -> str:
            """Sync resize with WebP format for better compression."""
            import base64
            from io import BytesIO
            from PIL import Image

            try:
                # Decode base64 to image
                image_data = base64.b64decode(screenshot_b64)
                img = Image.open(BytesIO(image_data))

                # Resize if needed
                max_size = (1024, 768)
                if img.size[0] > max_size[0] or img.size[1] > max_size[1]:
                    img.thumbnail(max_size, Image.Resampling.LANCZOS)

                # Convert to WebP for better compression
                buffer = BytesIO()
                img.save(buffer, format='WebP', quality=85, optimize=True)

                # Re-encode to base64
                return base64.b64encode(buffer.getvalue()).decode('utf-8')
            except Exception:
                # Fallback to original
                return screenshot_b64

        # Replace the resize function if it exists
        if hasattr(prompts, '_resize_screenshot'):
            prompts._resize_screenshot = _resize_screenshot_async

    except ImportError:
        pass


def patch_message_manager():
    """Patch message manager to integrate local vision processing."""
    try:
        import browser_use.agent.message_manager.service as mm_service
        from .local_vision import analyze_screenshot_locally

        original_get_user_message = mm_service.MessageManager.get_user_message

        async def enhanced_get_user_message(self, browser_state, prompt_text: str):
            """Enhanced get_user_message with local vision integration."""

            # Try local vision analysis before creating LLM message
            if hasattr(browser_state, 'screenshot') and browser_state.screenshot:
                try:
                    # Check if local vision is enabled
                    import os
                    if not os.getenv('LOCAL_VISION_ENABLED', 'true').lower() in ('true', '1', 'yes'):
                        return await original_get_user_message(self, browser_state, prompt_text)

                    last_action_type = getattr(self, '_last_action_type', 'none')
                    session_id = getattr(self, '_session_id', 'default')
                    confidence_threshold = float(os.getenv('LOCAL_VISION_CONFIDENCE_THRESHOLD', '0.85'))

                    local_result = await analyze_screenshot_locally(
                        screenshot_b64=browser_state.screenshot,
                        last_action_type=last_action_type,
                        session_id=session_id,
                        confidence_threshold=confidence_threshold
                    )

                    if local_result and not local_result.should_escalate:
                        # Use local vision description - modify browser_state
                        from copy import deepcopy
                        enhanced_state = deepcopy(browser_state)

                        # Add local vision insight to extracted content
                        if hasattr(enhanced_state, 'extracted_content'):
                            enhanced_state.extracted_content = (
                                f"[Local Vision - Confidence: {local_result.confidence:.2f}] "
                                f"{local_result.description}\n\n"
                                f"{enhanced_state.extracted_content}"
                            )

                        return await original_get_user_message(self, enhanced_state, prompt_text)

                except Exception:
                    # Fail silently, use original path
                    pass

            # Use original implementation
            return await original_get_user_message(self, browser_state, prompt_text)

        # Replace the method
        mm_service.MessageManager.get_user_message = enhanced_get_user_message

    except ImportError:
        pass


def patch_agent_service():
    """Patch agent service to track action types."""
    try:
        import browser_use.agent.service as agent_service

        original_init = agent_service.Agent.__init__

        def enhanced_init(self, *args, **kwargs):
            """Enhanced Agent init to track actions."""
            # Call original init
            original_init(self, *args, **kwargs)

            # Add action tracking
            self._last_action_type = 'none'
            self._session_id = kwargs.get('session_id', 'default')

        def _classify_action_type(self, action) -> str:
            """Classify action as mutating or non-mutating."""
            if not action:
                return 'none'

            action_str = str(action).lower()
            MUTATING_KEYWORDS = {"click", "type", "fill", "submit", "select", "upload", "press", "input", "check", "drag"}

            for keyword in MUTATING_KEYWORDS:
                if keyword in action_str:
                    return keyword

            return 'none'

        # Patch the Agent class
        agent_service.Agent.__init__ = enhanced_init
        agent_service.Agent._classify_action_type = _classify_action_type

        # Patch action execution to track types
        if hasattr(agent_service.Agent, '_execute_action'):
            original_execute = agent_service.Agent._execute_action

            async def enhanced_execute(self, action):
                """Enhanced action execution with tracking."""
                # Track action type
                self._last_action_type = self._classify_action_type(action)

                # Execute original
                return await original_execute(self, action)

            agent_service.Agent._execute_action = enhanced_execute

    except ImportError:
        pass