#!/usr/bin/env python3
"""
Demo of BrowserAgent's advanced element detection system.
Features:
1. Improved modal visibility filtering (elements behind overlays properly excluded)
2. Intelligent combination of rule-based (DOM) and AI-based (vision) detection
3. Intersection ratio analysis for optimal element selection
"""

import asyncio
import os
from pathlib import Path
import dotenv

# Load environment variables
dotenv.load_dotenv()

# Import MARSYS framework components
import sys
sys.path.insert(0, str(Path(__file__).parent))
from marsys.models.models import ModelConfig
from marsys.agents.browser_agent import BrowserAgent


async def demo_advanced_element_detection():
    """
    Demo: Advanced element detection with modal visibility filtering and intelligent combination.
    """
    print("🚀 Demo: Advanced Element Detection System")
    print("=" * 60)
    
    # Step 1: Create model configurations
    print("🔧 Setting up models...")
    
    # Main model for the BrowserAgent
    main_model_config = ModelConfig(
        type="api",
        provider="google",
        name="gemini-2.5-pro",
        api_key=os.getenv("GOOGLE_API_KEY"),
        max_tokens=2048,
        temperature=0.1
    )
    
    # Vision model for element detection
    vision_model_config = ModelConfig(
        type="api",
        provider="google",
        name="gemini-2.5-pro",
        api_key=os.getenv("GOOGLE_API_KEY"),
        max_tokens=4096,
        temperature=0.1,
        thinking_budget=2000
    )
    
    browser_agent = None
    
    try:
        # Step 2: Create BrowserAgent with vision capabilities
        print("🤖 Creating BrowserAgent with vision capabilities...")
        
        browser_agent = await BrowserAgent.create_safe(
            model_config=main_model_config,
            description="Browser agent with AI-powered element detection",
            headless=False,  # Set to True for headless mode
            viewport_width=1980,
            viewport_height=1080,
            vision_model_config=vision_model_config,  # Enable AI element detection
            tmp_dir="./tmp",
        )
        
        print(f"   ✅ BrowserAgent created with vision capabilities")
        
        # Step 3: Navigate to a website
        print("📱 Navigating to website...")
        await browser_agent.browser_tool.goto("https://www.google.com")
        await asyncio.sleep(2)  # Wait for page to load
        print("   ✅ Navigated to Google")
        
        # Step 4: Detect interactive elements using AI
        print("🔍 Detecting interactive elements...")
        elements = await browser_agent.predict_interactive_elements()
        print(f"   ✅ Detected {len(elements)} interactive elements")
        
        # Step 5: Test improved visibility detection
        print("🔍 Testing improved modal visibility detection...")
        

        
        # Step 6: Test intelligent combination
        print("\n   🧠 Intelligent combination (AI + Rule-based with visibility fix):")
        combined_result = await browser_agent.highlight_interactive_elements(
            visible_only=True,
            use_prediction=True,
            use_rule_based=True,
            intersection_threshold=0.2  # Balanced threshold
        )
        
        elements = combined_result['elements']
        screenshot_path = combined_result['screenshot_path']
        
        # Analyze the intelligent combination results
        rule_based_count = sum(1 for elem in elements if elem.get('source') == 'rule_based')
        prediction_count = sum(1 for elem in elements if elem.get('source') == 'vision_prediction')
        matched_rule_based = sum(1 for elem in elements if elem.get('matched_with_prediction', False))
        unmatched_predictions = sum(1 for elem in elements if elem.get('no_rule_based_match', False))
        
        print(f"      - Total elements: {len(elements)}")
        print(f"      - Rule-based elements: {rule_based_count}")
        print(f"      - AI prediction elements: {prediction_count}")
        print(f"      - Rule-based matched with AI: {matched_rule_based}")
        print(f"      - Unique AI predictions: {unmatched_predictions}")
        print(f"      - Screenshot: {screenshot_path}")
        

        # Show sample elements with their sources
        if elements:
            print(f"\n   🔍 Sample combined elements:")
            for i, elem in enumerate(elements[:5]):
                source = elem.get('source', 'unknown')
                label = elem.get('label', 'unknown')
                bbox = elem.get('bbox', [])
                
                # Add source info
                if elem.get('matched_with_prediction'):
                    source_info = f"{source} (matched with AI, ratio: {elem.get('intersection_ratio', 0):.3f})"
                elif elem.get('no_rule_based_match'):
                    source_info = f"{source} (unique AI detection)"
                elif elem.get('no_prediction_match'):
                    source_info = f"{source} (unique rule-based)"
                else:
                    source_info = source
                
                print(f"   {i+1:2d}. [{source_info}]")
                print(f"       Label: {label}")
                print(f"       Bbox: {bbox}")
        

        
        print(f"\n✅ Demo completed successfully!")
        
    except Exception as e:
        print(f"❌ Demo failed: {e}")
        import traceback
        traceback.print_exc()
    
    finally:
        # Clean up
        if browser_agent:
            await browser_agent.close()
            print("🧹 Browser closed")


if __name__ == "__main__":
    print("MARSYS Framework - Advanced Element Detection System Demo")
    print("This demo showcases:")
    print("1. Modal visibility filtering (elements behind overlays excluded)")  
    print("2. Intelligent rule-based + AI-based detection combination")
    print("3. Intersection ratio analysis for optimal results")
    print("Make sure GOOGLE_API_KEY is set in your .env file\n")
    
    asyncio.run(demo_advanced_element_detection()) 