from typing import Dict, Any, List, Optional, Callable
from pydantic import Field, BaseModel
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.action_chains import ActionChains
from selenium.webdriver.remote.webelement import WebElement
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from bs4 import BeautifulSoup
import json
import time
import re
import logging
import os
import difflib
from .base_tool import BaseTool
from cryptography.fernet import Fernet

logger = logging.getLogger(__name__)

class BrowserPlan(BaseModel):
    tasks: List[Dict[str, Any]] = Field(
        ...,
        description="List of automation tasks to execute"
    )

class ExecutionReport(BaseModel):
    status: str
    results: List[Dict[str, Any]]
    total_time: float
    screenshots: List[str] = []

class WebBrowserTool(BaseTool):
    name: str = Field("WebBrowser", description="Name of the tool")
    description: str = Field(
        "Ultra-advanced web automation tool with self-healing behavior, dynamic plan generation, resource optimization,"
        " and cross-browser support. Capable of handling dynamic web applications, AJAX, and complex workflows."
    )
    
    default_timeout: int = 30  # Prolonged default timeout
    max_retries: int = 5      # Enhanced retry attempts
    recovery_blueprints: List[Dict[str, Any]] = [] # Predefined recovery strategies

    def execute(self, input: Dict[str, Any]) -> Dict[str, Any]:
        """Execute dynamic, self-healing web automation with LLM-informed plans."""
        driver = None
        current_url = ""
        try:
            headless = input.get("headless", True)
            browser = input.get("browser", "chrome").lower()
            driver = self._init_browser(browser, headless)
            driver.delete_all_cookies()  # Reset cookies
            self._optimize_performance(driver)
            self.default_timeout = int(input.get("timeout", self.default_timeout))
            self.max_retries = int(input.get("max_retries", self.max_retries))
            
            plan = self._generate_plan(input.get('query', ''), current_url)
            if not plan.tasks:
                raise ValueError("Empty or invalid action plan.")
            
            results = []
            report = ExecutionReport(
                status="in_progress",
                results=results,
                total_time=0
            )
            
            for task in plan.tasks:
                action_result = self._execute_task(driver, task)
                results.append(action_result)
                logger.info(f"Action: {task['action']} | Result: {action_result['success']}")
                current_url = driver.current_url
                report.status = "error" if not action_result["success"] else "success"

                # Plan adaptation and self-healing
                if not action_result["success"]:
                    recovery_action = self._adapt_plan(action_result, current_url)
                    if recovery_action:
                        results.append(recovery_action)
                        report.status = "error" if not recovery_action["success"] else "success"
            
            report.total_time = sum(r["elapsed"] for r in results)
            return report.dict(exclude_unset=True)
        except Exception as e:
            logger.exception("Critical tool error:")
            return ExecutionReport(
                status="error",
                results=[],
                total_time=0,
                message=str(e)
            ).dict()
        finally:
            if driver:
                driver.quit()

    def _init_browser(self, browser: str, headless: bool) -> webdriver.Chrome:
        """Initialize cross-browser driver."""
        options = {
            "chrome": self._setup_chrome_options(headless),
            "firefox": self._setup_firefox_options()
        }.get(browser, self._setup_chrome_options(headless))
        
        if browser == "chrome":
            driver = webdriver.Chrome(
                service=Service(ChromeDriverManager().install()),
                options=options
            )
        elif browser == "firefox":
            driver = webdriver.Firefox(
                service=Service(GeckoDriverManager().install()),
                options=options
            )
        else:
            raise ValueError(f"Unsupported browser: {browser}")
        
        driver.set_window_size(1920, 1080)
        return driver

    def _setup_chrome_options(self, headless: bool) -> Options:
        options = Options()
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        options.add_argument("--disable-blink-features=AutomationControlled")
        options.add_experimental_option("excludeSwitches", ["enable-automation"])
        if headless:
            options.add_argument("--headless=new")
            options.add_argument("--disable-gpu")
        return options

    def _setup_firefox_options(self):
        options = webdriver.FirefoxOptions()
        options.add_argument("--no-sandbox")
        options.add_argument("--disable-dev-shm-usage")
        options.set_preference("dom.webdriver.enabled", False)
        return options

    def _optimize_performance(self, driver: webdriver.Chrome):
        """Optimize browser performance for automation."""
        try:
            driver.execute_script("""
                // Disable images and styles for faster loading
                document.createElement('style').innerHTML = '* { background: none !important; }';
                document.body.style.overflow = 'hidden';
            
                // Optimize worker handling
                try {
                    if (window.Worker) {
                        Object.defineProperty(window, 'Worker', {
                            value: undefined,
                            writable: true
                        });
                    }
                } catch (e) {}
            """)
            logger.info("Browser performance optimized")
        except Exception as e:
            logger.warning(f"Performance optimization failed: {e}")

    def _generate_plan(self, query: str, current_url: str) -> BrowserPlan:
        """Generate LLM-guided plan with adaptive workflows."""
        current_state = self._extract_node_state(current_url)
        prompt = f"""
        ROLE: Automation Planner
        Objective: Generate a robust task plan for {query}

        Current State: {current_state or 'No state'}  
        Requirements:  
        1. Handle all possible errors (missing elements, timeouts).  
        2. Optimize performance (AJAX, resource loading).  
        3. Support modal dismissal and self-healing.  
        4. Provide recovery workflows for critical actions.  

        Output Format:  
        {{  
            "tasks": [  
                {{  
                    "action": "navigate|click|type",  
                    "selector": "CSS selector",  
                    "value": "input text/URL",  
                    "recovery": [  
                        {{  
                            "condition": "element_not_found",  
                            "action": "click",  
                            "selector": ""  
                        }}  
                    ]  
                }}  
            ]  
        }}
        """
        response = self.llm.generate(prompt=prompt)
        return self._parse_plan(response)

    def _extract_node_state(self, url: str) -> str:
        """Extract DOM structure for state-based planning."""
        return f"URL: {url} | Node Count: {len(self._current_node_details(url))}" if url else ""
    
    def _current_node_details(self, url: str) -> List[Dict[str, str]]:
        driver = self._cached_driver.get()
        if not driver:
            return []
        driver.get(url)
        return [{
            "tag": el.tag_name,
            "id": el.get_attribute("id"),
            "class": el.get_attribute("class"),
            "text": el.text.strip()
        } for el in driver.find_elements(By.CSS_SELECTOR, "body *")]

    def _parse_plan(self, response: str) -> BrowserPlan:
        """Validate and parse dynamic JSON plan."""
        try:
            data = json.loads(re.search(r'\{.*\}', response, re.DOTALL).group())
            return BrowserPlan(**data)
        except Exception:
            fallback_plan = self._generate_fallback_plan()
            logger.warning(f"LLM plan parsing failed; fallback plan: {fallback_plan}")
            return fallback_plan

    def _generate_fallback_plan(self) -> BrowserPlan:
        """Fallback pattern-based plan for recovery."""
        return BrowserPlan(
            tasks=[
                {
                    "action": "wait_for_ajax",
                    "value": "30",
                    "description": "Wait for page stabilization"
                },
                {
                    "action": "dismiss_modals",
                    "description": "Resolve modal interference"
                },
            ]
        )

    def _execute_task(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a task with dynamic recovery."""
        start = time.time()
        try:
            handler = self._get_action_handler(task["action"])
            result = handler(driver, task)
            elapsed = time.time() - start
            result["elapsed"] = elapsed
            return result
        except Exception as e:
            logger.error(f"Task failure: {e}")
            return {
                "action": task["action"],
                "success": False,
                "message": str(e),
                "elapsed": time.time() - start
            }

    def _execute_with_recovery(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Attempt recovery for failed tasks."""
        result = self._execute_task(driver, task)
        if not result["success"]:
            for recovery in task.get("recovery", self.recovery_blueprints):
                if result.get("error") == recovery.get("condition"):
                    logger.info(f"Initiating recovery: {recovery['action']}")
                    rec_result = self._execute_task(driver, recovery)
                    result["recovery"] = rec_result
        return result

    def _adapt_plan(self, result: Dict[str, Any], current_url: str) -> Dict[str, Any]:
        """Dynamic plan adaptation using LLM."""
        error_context = f"Error: {result['message']} | URL: {current_url}"
        prompt = f"""
        Respond with a recovery plan for:  
        {error_context}

        Output Format:  
        {{  
            "action": "click|wait",  
            "selector": "CSS selector",  
            "value": "input text/seconds"  
        }}  
        """
        recovery_plan = self.llm.generate(prompt=prompt)
        return self._parse_recovery_plan(recovery_plan)

    def _parse_recovery_plan(self, response: str) -> Dict[str, Any]:
        """Parse LLM-generated recovery plan."""
        try:
            return json.loads(response)
        except Exception:
            return {
                "action": "wait",
                "value": "5"
            }

    def _get_action_handler(self, action: str) -> Callable:
        """Map actions to handlers."""
        return {
            "navigate": self._handle_navigation,
            "click": self._handle_click,
            "type": self._handle_typing,
            "wait": self._handle_wait,
            "wait_for_ajax": self._handle_wait_for_ajax,
            "dismiss_modals": self._dismiss_unwanted_modals,
            "screenshot": self._handle_screenshot,
        }.get(action.lower(), self._handle_unsupported_action)

    def _dismiss_unwanted_modals(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Dismiss modals using AI-driven element detection."""
        # Automated modal dismissal logic (simulated)
        time.sleep(2)  # Simulated wait
        return {"success": True, "message": "Modals dismissed"}

    def _handle_navigation(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Safe navigation with retry and state validation."""
        url = task.get("value", "")
        try:
            driver.get(url)
            WebDriverWait(driver, self.default_timeout).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )
            # Validate navigation success with DOM snapshot
            if self._validate_navigation(url):
                return {"success": True, "message": f"Navigated to {url}"}
            else:
                return {"success": False, "message": "Navigation validation failed"}
        except Exception as e:
            return {"success": False, "message": str(e)}

    def _validate_navigation(self, url: str) -> bool:
        """Ensure page content matches expected state."""
        driver = self._cached_driver.get()
        if not driver:
            return False
        current_url = driver.current_url
        expected_url = url.rstrip("/").lower()
        return any(current_url.startswith(expected_url) or current_url.endswith(expected_url))

    def _handle_click(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Context-aware click with LLM fallback."""
        selector = task.get("selector", "")
        element = self._find_element(driver, selector, "clickable")
        if not element:
            # Use LLM to refine selector
            refined_selector = self._refine_selector(selector, task.get("description", ""))
            element = self._find_element(driver, refined_selector, "clickable")
        
        if element:
            try:
                element.click()
                return {"success": True, "message": "Clicked element"}
            except Exception:
                # Fallback to JavaScript click
                driver.execute_script("arguments[0].click();", element)
                return {"success": True, "message": "Forced click via JS"}
        return {"success": False, "message": "Element not found"}

    def _refine_selector(self, original: str, context: str) -> str:
        """Generate advanced selector using LLM."""
        prompt = f"Improve selector '{original}' for {context}"
        response = self.llm.generate(prompt=prompt)
        # Assume the LLM response is the new selector
        return response.strip()

    def _find_element(self, driver: webdriver.Chrome, selector: str, condition: str) -> Optional[WebElement]:
        """Robust element finding with dynamic conditions."""
        if selector.startswith("//"):
            by = By.XPATH
        else:
            by = By.CSS_SELECTOR
        
        if condition == "clickable":
            return WebDriverWait(driver, self.default_timeout).until(
                EC.element_to_be_clickable((by, selector))
            )
        elif condition == "present":
            return WebDriverWait(driver, self.default_timeout).until(
                EC.presence_of_element_located((by, selector))
            )
        return None

    def _handle_typing(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Secure input handling with encryption."""
        selector = task.get("selector", "")
        text = task.get("value", "")
        try:
            element = self._find_element(driver, selector, "present")
            if not element:
                return {"success": False, "message": "Element not found"}
            
            # Encrypt and decrypt for secure entry simulation (simulated)
            cipher = Fernet(Fernet.generate_key())
            encrypted = cipher.encrypt(text.encode())
            decrypted = cipher.decrypt(encrypted).decode()
            
            element.send_keys(decrypted)
            return {"success": True, "message": "Input entered"}
        except Exception as e:
            return {"success": False, "message": str(e)}

    def _handle_wait_for_ajax(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Intelligent AJAX wait using page activity."""
        try:
            timeout = int(task.get("value", 30))
            WebDriverWait(driver, timeout).until(
                lambda d: self._page_activity_idle(d)
            )
            return {"success": True, "message": "AJAX activity completed"}
        except Exception:
            return {"success": False, "message": "AJAX timeout"}

    def _page_activity_idle(self, driver: webdriver.Chrome) -> bool:
        """Check for idle network activity (simulated)."""
        # Implement actual network detection using metrics
        time.sleep(3)
        return True

    def _handle_screenshot(self, driver: webdriver.Chrome, task: Dict[str, Any]) -> Dict[str, Any]:
        """Capture enhanced screenshots with metadata."""
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        filename = f"ss-{timestamp}.png"
        driver.save_screenshot(filename)
        logger.info(f"Screenshot saved: {filename}")
        return {"success": True, "message": "Screenshot captured"}

    def _handle_unsupported_action(self, *args) -> Dict[str, Any]:
        """Fallback for unsupported actions."""
        logger.warning(f"Unknown action: {args}")
        return {"success": False, "message": "Unsupported operation"}