import gradio as gr
import os
import asyncio
import nest_asyncio
from datetime import datetime
from typing import Optional, Dict, Any

from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination
from autogen_agentchat.teams import SelectorGroupChat
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.agents.web_surfer import MultimodalWebSurfer

# Enable nested event loops for Jupyter compatibility
nest_asyncio.apply()

class AIShoppingAnalyzer:
    def __init__(self, api_key: str):
        self.api_key = api_key
        # Set the API key in environment
        os.environ["OPENAI_API_KEY"] = api_key
        self.model_client = OpenAIChatCompletionClient(model="gpt-4o")
        self.termination = MaxMessageTermination(max_messages=20) | TextMentionTermination("TERMINATE")
        
    def create_websurfer(self) -> MultimodalWebSurfer:
        """Initialize the web surfer agent for e-commerce research"""
        return MultimodalWebSurfer(
            name="websurfer_agent",
            description="""E-commerce research specialist that:
            1. Searches multiple retailers for product options
            2. Compares prices and reviews
            3. Checks product specifications and availability
            4. Analyzes website structure and findability
            5. Detects and analyzes structured data (Schema.org, JSON-LD, Microdata)
            6. Evaluates product markup and rich snippets
            7. Checks for proper semantic HTML and data organization""",
            model_client=self.model_client,
            headless=True
        )

    def create_assistant(self) -> AssistantAgent:
        """Initialize the shopping assistant agent"""
        return AssistantAgent(
            name="assistant_agent",
            description="E-commerce shopping advisor and website analyzer",
            system_message="""You are an expert shopping assistant and e-commerce analyst. Your role is to:
            1. Help find products based on user needs
            2. Compare prices and features across different sites
            3. Analyze website usability and product findability
            4. Evaluate product presentation and information quality
            5. Assess the overall e-commerce experience
            6. Analyze structured data implementation:
               - Check for Schema.org markup
               - Validate JSON-LD implementation
               - Evaluate microdata usage
               - Assess rich snippet potential
            7. Report on data structure quality:
               - Product markup completeness
               - Price and availability markup
               - Review and rating markup
               - Inventory status markup

            When working with the websurfer_agent:
            - Guide their research effectively
            - Verify the information they find
            - Analyze how easy it was to find products
            - Evaluate product page quality
            - Say 'keep going' if more research is needed
            - Say 'TERMINATE' only when you have a complete analysis""",
            model_client=self.model_client
        )

    def create_team(self, websurfer_agent: MultimodalWebSurfer, assistant_agent: AssistantAgent) -> SelectorGroupChat:
        """Set up the team of agents"""
        user_proxy = UserProxyAgent(
            name="user_proxy",
            description="An e-commerce site owner looking for AI shopping analysis"
        )

        return SelectorGroupChat(
            participants=[websurfer_agent, assistant_agent, user_proxy],
            selector_prompt="""You are coordinating an e-commerce analysis system. The following roles are available:
            {roles}

            Given the conversation history {history}, select the next role from {participants}.
            - The websurfer_agent searches products and analyzes website structure
            - The assistant_agent evaluates findings and makes recommendations
            - The user_proxy provides input when needed

            Return only the role name.""",
            model_client=self.model_client,
            termination_condition=self.termination
        )

    async def analyze_site(self, 
                         website_url: str, 
                         product_category: str, 
                         specific_product: Optional[str] = None) -> str:
        """Run the analysis with proper cleanup"""
        websurfer = None
        try:
            # Set up the analysis query
            query = f"""Analyze the e-commerce experience for {website_url} focusing on:
            1. Product findability in the {product_category} category
            2. Product information quality
            3. Navigation and search functionality
            4. Price visibility and comparison features"""
            
            if specific_product:
                query += f"\n5. Detailed analysis of this specific product: {specific_product}"

            # Initialize agents with automatic browser management
            websurfer = self.create_websurfer()
            assistant = self.create_assistant()

            # Create team
            team = self.create_team(websurfer, assistant)
            
            # Modified execution to handle EOF errors
            try:
                result = []
                async for message in team.run_stream(task=query):
                    if isinstance(message, str):
                        result.append(message)
                    else:
                        result.append(str(message))
                return "\n".join(result)
            except EOFError:
                return "Analysis completed with some limitations. Please try again if results are incomplete."
            except Exception as e:
                return f"Analysis error: {str(e)}"

        finally:
            if websurfer:
                try:
                    await websurfer.close()
                except Exception as e:
                    print(f"Cleanup error: {str(e)}")
                    # Continue even if cleanup fails

def create_gradio_interface() -> gr.Interface:
    """Create the Gradio interface for the AI Shopping Analyzer"""
    
    css = """
    @import url('https://fonts.googleapis.com/css2?family=Open+Sans:wght@300;400;600;700&display=swap');
    
    body { 
        font-family: 'Open Sans', sans-serif !important; 
    }
    
    .dashboard-container { 
        border: 1px solid #e0e5ff; 
        border-radius: 8px; 
        background-color: #ffffff;
    }
    
    .token-header { 
        font-size: 1.25rem; 
        font-weight: 600; 
        margin-top: 1rem; 
        margin-bottom: 0.5rem;
    }
    
    .feature-button { 
        display: inline-block; 
        margin: 0.25rem; 
        padding: 0.5rem 1rem; 
        background-color: #f3f4f6; 
        border: 1px solid #e5e7eb; 
        border-radius: 0.375rem; 
        font-size: 0.875rem;
    }
    
    .feature-button:hover { 
        background-color: #e5e7eb;
    }
    
    /* Custom styling for form elements */
    .gr-form {
        background: transparent !important;
        border: none !important;
        box-shadow: none !important;
    }
    
    .gr-input, .gr-textarea {
        border: 1px solid #e5e7eb !important;
        border-radius: 6px !important;
        padding: 8px 12px !important;
        font-size: 14px !important;
        transition: all 0.2s !important;
    }
    
    .gr-input:focus, .gr-textarea:focus {
        border-color: #4c4ce3 !important;
        outline: none !important;
        box-shadow: 0 0 0 2px rgba(76, 76, 227, 0.2) !important;
    }
    
    .gr-button {
        background-color: #4c4ce3 !important;
        color: white !important;
        border-radius: 6px !important;
        padding: 8px 16px !important;
        font-size: 14px !important;
        font-weight: 600 !important;
        transition: all 0.2s !important;
    }
    
    .gr-button:hover {
        background-color: #3a3ab8 !important;
    }
    """
    
    def validate_api_key(api_key: str) -> bool:
        """Validate the OpenAI API key format"""
        return api_key.startswith("sk-") and len(api_key) > 20

    async def run_analysis(api_key: str,
                         website_url: str,
                         product_category: str,
                         specific_product: str) -> str:
        """Handle the analysis submission"""
        if not validate_api_key(api_key):
            return "Please enter a valid OpenAI API key (should start with 'sk-')"
        
        if not website_url:
            return "Please enter a website URL"
            
        if not product_category:
            return "Please specify a product category"
            
        try:
            analyzer = AIShoppingAnalyzer(api_key)
            result = await analyzer.analyze_site(
                website_url=website_url,
                product_category=product_category,
                specific_product=specific_product if specific_product else None
            )
            return result
        except Exception as e:
            return f"Error during analysis: {str(e)}"

    # Create the interface
    return gr.Interface(
        fn=run_analysis,
        inputs=[
            gr.Textbox(label="OpenAI API Key", placeholder="sk-...", type="password"),
            gr.Textbox(label="Website URL", placeholder="https://your-store.com"),
            gr.Textbox(label="Product Category", placeholder="e.g., Electronics, Clothing, etc."),
            gr.Textbox(label="Specific Product (Optional)", placeholder="e.g., Blue Widget Model X")
        ],
        outputs=gr.Textbox(label="Analysis Results", lines=20),
        title="AI Shopping Agent Analyzer",
        description="""Analyze how your e-commerce site performs when the shopper is an AI agent.
        This tool helps you understand your site's effectiveness for AI-powered shopping assistants.""",
        theme="default",
        allow_flagging="never"
    )

if __name__ == "__main__":
    # Install Playwright browsers and dependencies
    import subprocess
    try:
        subprocess.run(["playwright", "install"], check=True)
        subprocess.run(["playwright", "install-deps"], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error installing Playwright dependencies: {e}")
    except Exception as e:
        print(f"Unexpected error during setup: {e}")

    # Create and launch the interface
    iface = create_gradio_interface()
    iface.launch()