import requests
import base64
import json
from typing import Optional

class BrowserAutomationClient:
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip('/')
        self.session_id: Optional[str] = None
    
    def create_session(self) -> str:
        """Create a new browser session"""
        response = requests.post(f"{self.base_url}/session/create")
        response.raise_for_status()
        data = response.json()
        self.session_id = data["session_id"]
        return self.session_id
    
    def close_session(self):
        """Close the current session"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.delete(f"{self.base_url}/session/{self.session_id}")
        response.raise_for_status()
        self.session_id = None
    
    def navigate(self, url: str) -> dict:
        """Navigate to a URL"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.post(
            f"{self.base_url}/session/{self.session_id}/navigate",
            json={"url": url}
        )
        response.raise_for_status()
        return response.json()
    
    def click(self, selector: str, selector_type: str = "css") -> dict:
        """Click an element"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.post(
            f"{self.base_url}/session/{self.session_id}/click",
            json={"selector": selector, "selector_type": selector_type}
        )
        response.raise_for_status()
        return response.json()
    
    def type_text(self, selector: str, text: str, selector_type: str = "css", clear_first: bool = True) -> dict:
        """Type text into an element"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.post(
            f"{self.base_url}/session/{self.session_id}/type",
            json={
                "selector": selector,
                "text": text,
                "selector_type": selector_type,
                "clear_first": clear_first
            }
        )
        response.raise_for_status()
        return response.json()
    
    def screenshot(self, save_path: Optional[str] = None) -> str:
        """Take a screenshot and optionally save it"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.get(f"{self.base_url}/session/{self.session_id}/screenshot")
        response.raise_for_status()
        data = response.json()
        
        if data["success"]:
            screenshot_data = data["screenshot"]
            
            if save_path:
                # Decode and save the image
                image_data = base64.b64decode(screenshot_data)
                with open(save_path, "wb") as f:
                    f.write(image_data)
            
            return screenshot_data
        else:
            raise Exception(f"Screenshot failed: {data['message']}")
    
    def get_page_info(self) -> dict:
        """Get information about the current page"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.get(f"{self.base_url}/session/{self.session_id}/page-info")
        response.raise_for_status()
        return response.json()
    
    def execute_js(self, script: str) -> dict:
        """Execute JavaScript on the page"""
        if not self.session_id:
            raise ValueError("No active session")
        
        response = requests.get(
            f"{self.base_url}/session/{self.session_id}/execute-js",
            params={"script": script}
        )
        response.raise_for_status()
        return response.json()

# Example usage
if __name__ == "__main__":
    # Replace with your Hugging Face Space URL
    client = BrowserAutomationClient("https://your-username-browser-api.hf.space")
    
    try:
        # Create session
        session_id = client.create_session()
        print(f"Created session: {session_id}")
        
        # Navigate to a website
        result = client.navigate("https://httpbin.org/forms/post")
        print(f"Navigation result: {result}")
        
        # Get page info
        page_info = client.get_page_info()
        print(f"Page title: {page_info['title']}")
        print(f"Found {len(page_info['elements'])} interactive elements")
        
        # Fill out a form
        client.type_text("input[name='custname']", "John Doe")
        client.type_text("input[name='custtel']", "123-456-7890")
        client.type_text("textarea[name='comments']", "This is a test comment")
        
        # Take a screenshot
        screenshot_b64 = client.screenshot("screenshot.png")
        print("Screenshot saved as screenshot.png")
        
        # Submit the form
        client.click("input[type='submit']")
        
        # Take another screenshot after submission
        client.screenshot("after_submit.png")
        print("After submit screenshot saved")
        
    except Exception as e:
        print(f"Error: {e}")
    
    finally:
        # Always close the session
        if client.session_id:
            client.close_session()
            print("Session closed")