# web_browser.py
from typing import Dict, Any, List, Optional, Callable
from pydantic import Field, BaseModel
from playwright.sync_api import sync_playwright, Page, TimeoutError as PlaywrightTimeoutError
import json, time, re, logging, os, difflib
from .base_tool import BaseTool
import difflib

logger = logging.getLogger(__name__)

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

class WebBrowserTool(BaseTool):
    name: str = Field("WebBrowser", description="Name of the tool")
    description: str = Field(
        "Universal web automation tool with advanced element identification (HTML and image fallback), AJAX waiting, modal dismissal, multi-tab support, and custom JS injection.",
        description="Tool description"
    )
    
    # Playwright timeouts are in milliseconds (default is 15 seconds).
    default_timeout: int = 15000  
    max_retries: int = 3       

    # Assumes self.llm is set (and supports both text and vision via generate_from_image)
    # You can inject this instance during tool initialization.
    
    def execute(self, input: Dict[str, Any]) -> Dict[str, Any]:
        """Execute an advanced dynamic web automation workflow using Playwright."""
        overall_start = time.time()
        results = []
        current_url = ""
        try:
            headless = input.get("headless", False)
            # Convert seconds to milliseconds
            self.default_timeout = int(input.get("timeout", 15)) * 1000
            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("No valid tasks in the generated plan.")

            with sync_playwright() as p:
                browser = p.chromium.launch(headless=headless)
                context = browser.new_context()
                page = context.new_page()

                # Map actions to handlers.
                action_map: Dict[str, Callable[[Page, Dict[str, Any]], Dict[str, Any]]] = {
                    "navigate": lambda p, task: self._handle_navigation(p, task.get("value", "")),
                    "click": lambda p, task: self._handle_click(p, task.get("selector", "")),
                    "type": lambda p, task: self._handle_typing(p, task.get("selector", ""), task.get("value", ""), task),
                    "wait": lambda p, task: self._handle_wait(task.get("value", "")),
                    "wait_for_ajax": lambda p, task: self._handle_wait_for_ajax(p, task.get("value", "30")),
                    "scroll": lambda p, task: self._handle_scroll(p, task.get("selector", "")),
                    "hover": lambda p, task: self._handle_hover(p, task.get("selector", "")),
                    "screenshot": lambda p, task: self._handle_screenshot(p, task.get("value", "screenshot.png")),
                    "switch_tab": lambda p, task: self._handle_switch_tab(context, task.get("value", "0")),
                    "execute_script": lambda p, task: self._handle_execute_script(p, task.get("value", "")),
                    "drag_and_drop": lambda p, task: self._handle_drag_and_drop(p, task.get("selector", ""), task.get("value", "")),
                }
                
                for task in plan.tasks:
                    # Dismiss any modals before each action.
                    self._dismiss_unwanted_modals(page)
                    action = task.get("action", "").lower()
                    logger.info(f"Executing task: {task.get('description', action)}")
                    start_time = time.time()
                    handler = action_map.get(action)
                    if not handler:
                        results.append({
                            "action": action,
                            "success": False,
                            "message": f"Unsupported action: {action}"
                        })
                        continue

                    result = self._execute_with_retries(page, task, handler)
                    elapsed = time.time() - start_time
                    result["elapsed"] = elapsed
                    logger.info(f"Action '{action}' completed in {elapsed:.2f} seconds.")
                    results.append(result)

                    if not result.get('success', False):
                        logger.error(f"Task failed: {result.get('message')}")
                        self._capture_failure_screenshot(page, action)
                        break

                    current_url = page.url

                overall_elapsed = time.time() - overall_start
                logger.info(f"Total execution time: {overall_elapsed:.2f} seconds.")
                browser.close()
                return {"status": "success", "results": results, "total_time": overall_elapsed}
        except Exception as e:
            logger.exception("Execution error:")
            return {"status": "error", "message": str(e)}

    def _generate_plan(self, query: str, current_url: str) -> BrowserPlan:
        prompt = f"""Generate browser automation plan for: {query}

Current URL: {current_url or 'No page loaded yet'}

Required JSON format:
{{
    "tasks": [
        {{
            "action": "navigate|click|type|wait|wait_for_ajax|scroll|hover|screenshot|switch_tab|execute_script|drag_and_drop",
            "selector": "CSS selector (optional)",
            "value": "input text/URL/seconds/filename/target-selector",
            "description": "action purpose"
        }}
    ]
}}

Guidelines:
1. Prefer IDs in selectors (#element-id) and semantic attributes.
2. Include wait steps after navigation and wait for AJAX where applicable.
3. Dismiss any modals/pop-ups that are not part of the task.
4. For drag_and_drop, use source selector in 'selector' and target selector in 'value'.
5. For execute_script, 'value' should contain valid JavaScript.
6. For switch_tab, 'value' should be an index or keyword 'new'.
"""
        response = self.llm.generate(prompt=prompt)
        return self._parse_plan(response)

    def _parse_plan(self, response: str) -> BrowserPlan:
        try:
            json_match = re.search(r'```json\n?(.+?)\n?```', response, re.DOTALL)
            if json_match:
                plan_data = json.loads(json_match.group(1).strip())
            else:
                json_str_match = re.search(r'\{.*\}', response, re.DOTALL)
                if not json_str_match:
                    raise ValueError("No JSON object found in the response.")
                plan_data = json.loads(json_str_match.group())
            validated_tasks = []
            for task in plan_data.get("tasks", []):
                if not all(key in task for key in ["action", "description"]):
                    logger.warning(f"Skipping task due to missing keys: {task}")
                    continue
                validated_tasks.append({
                    "action": task["action"],
                    "selector": task.get("selector", ""),
                    "value": task.get("value", ""),
                    "description": task["description"]
                })
            return BrowserPlan(tasks=validated_tasks)
        except (json.JSONDecodeError, AttributeError, ValueError) as e:
            logger.error(f"Plan parsing failed: {e}")
            return BrowserPlan(tasks=[])

    def _execute_with_retries(self, page: Page, task: Dict[str, Any],
                                handler: Callable[[Page, Dict[str, Any]], Dict[str, Any]]) -> Dict[str, Any]:
        """Execute a task with retry logic and fall back to image-based identification if needed."""
        attempts = 0
        result = {}
        while attempts < self.max_retries:
            result = self._execute_safe_task(page, task, handler)
            if result.get("success", False):
                return result
            attempts += 1
            logger.info(f"Retrying task '{task.get('action')}' (attempt {attempts + 1}/{self.max_retries})")
            time.sleep(1 * attempts)
        # If HTML-based methods fail, use the fallback with the image LLM function.
        if task.get("action") in ["click"]:
            logger.info("HTML-based automation failed. Using fallback with image LLM.")
            result = self._fallback_with_image_llm(page, task)
        return result

    def _execute_safe_task(self, page: Page, task: Dict[str, Any],
                             handler: Callable[[Page, Dict[str, Any]], Dict[str, Any]]) -> Dict[str, Any]:
        try:
            return handler(page, task)
        except Exception as e:
            action = task.get("action", "unknown")
            logger.exception(f"Error executing task '{action}':")
            return {"action": action, "success": False, "message": f"Critical error: {str(e)}"}

    def _dismiss_unwanted_modals(self, page: Page):
        try:
            modal_selectors = [".modal", ".popup", '[role="dialog"]', ".overlay", ".lightbox"]
            for selector in modal_selectors:
                elements = page.query_selector_all(selector)
                for modal in elements:
                    if modal.is_visible():
                        close_selectors = [".close", ".btn-close", "[aria-label='Close']", "[data-dismiss='modal']"]
                        dismissed = False
                        for close_sel in close_selectors:
                            try:
                                close_button = modal.query_selector(close_sel)
                                if close_button and close_button.is_visible():
                                    close_button.click()
                                    dismissed = True
                                    logger.info(f"Dismissed modal using selector {close_sel}")
                                    time.sleep(1)
                                    break
                            except Exception:
                                continue
                        if not dismissed:
                            page.evaluate("(modal) => modal.remove()", modal)
                            logger.info(f"Removed overlay/modal with selector {selector}")
        except Exception as e:
            logger.debug(f"Modal dismissal error: {e}")

    # -------------------------
    # New: Fallback via Image-Based LLM
    # -------------------------
    def _annotate_page_with_numbers(self, page: Page):
        """Inject a JS snippet that draws numbered red boxes around interactive elements."""
        script = """
        (() => {
            // Remove previous overlays.
            document.querySelectorAll('.automation-annotation-overlay').forEach(el => el.remove());
            const elements = document.querySelectorAll('button, a, input, [onclick]');
            let counter = 1;
            elements.forEach(el => {
                const rect = el.getBoundingClientRect();
                if (rect.width === 0 || rect.height === 0) return;
                const overlay = document.createElement('div');
                overlay.classList.add('automation-annotation-overlay');
                overlay.style.position = 'absolute';
                overlay.style.left = (rect.left + window.scrollX) + 'px';
                overlay.style.top = (rect.top + window.scrollY) + 'px';
                overlay.style.width = rect.width + 'px';
                overlay.style.height = rect.height + 'px';
                overlay.style.border = '2px solid red';
                overlay.style.zIndex = 9999;
                overlay.style.pointerEvents = 'none';
                overlay.textContent = counter;
                overlay.style.fontSize = '16px';
                overlay.style.fontWeight = 'bold';
                overlay.style.color = 'red';
                overlay.style.backgroundColor = 'rgba(255, 255, 255, 0.7)';
                document.body.appendChild(overlay);
                counter += 1;
            });
        })();
        """
        page.evaluate(script)

    def _click_element_by_number(self, page: Page, number: int) -> Dict[str, Any]:
        """Locate interactive elements in order and click the one matching the given number."""
        elements = page.query_selector_all("button, a, input, [onclick]")
        index = number - 1  # numbering starts at 1
        if index < len(elements):
            try:
                elements[index].click()
                return {"action": "click", "success": True, "message": f"Clicked element number {number}"}
            except Exception as e:
                return {"action": "click", "success": False, "message": f"Click failed: {str(e)}"}
        else:
            return {"action": "click", "success": False, "message": f"Element number {number} not found."}

    def _fallback_with_image_llm(self, page: Page, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Fallback method: Annotate the page, capture a screenshot, and send it to the LLM's vision function.
        Uses generate_from_image to extract the numbered element corresponding to the task.
        """
        try:
            self._annotate_page_with_numbers(page)
            time.sleep(1)  # Allow time for the overlays to render.
            screenshot_bytes = page.screenshot(type="png")
            prompt = (f"Based on the annotated screenshot, identify the number of the element "
                      f"that best corresponds to the following task: {task.get('description', '')}. "
                      "Return only the number.")
            # Use the LLM's image function (same LLM instance) to process the screenshot.
            response_text = self.llm.generate_from_image(prompt, image_bytes=screenshot_bytes)
            match = re.search(r'\d+', response_text)
            if match:
                number = int(match.group(0))
                logger.info(f"LLM indicated element number {number} for task '{task.get('action')}'.")
                return self._click_element_by_number(page, number)
            else:
                return {"action": task.get("action"), "success": False, "message": "No element number found in LLM response."}
        except Exception as e:
            return {"action": task.get("action"), "success": False, "message": f"Fallback failed: {str(e)}"}

    # -------------------------
    # Existing Handlers
    # -------------------------
    def _handle_navigation(self, page: Page, url: str) -> Dict[str, Any]:
        if not url.startswith(("http://", "https://")):
            url = f"https://{url}"
        try:
            page.goto(url, timeout=self.default_timeout)
            page.wait_for_selector("body", timeout=self.default_timeout)
            return {"action": "navigate", "success": True, "message": f"Navigated to {url}"}
        except PlaywrightTimeoutError as e:
            logger.error(f"Navigation to {url} timed out: {e}")
            return {"action": "navigate", "success": False, "message": f"Navigation timed out: {str(e)}"}
        except Exception as e:
            logger.error(f"Navigation to {url} failed: {e}")
            return {"action": "navigate", "success": False, "message": f"Navigation failed: {str(e)}"}

    def _handle_click(self, page: Page, selector: str) -> Dict[str, Any]:
        try:
            page.wait_for_selector(selector, state="visible", timeout=self.default_timeout)
            page.click(selector, timeout=self.default_timeout)
            return {"action": "click", "success": True, "message": f"Clicked element: {selector}"}
        except PlaywrightTimeoutError as e:
            logger.error(f"Click action timed out on selector {selector}: {e}")
            return {"action": "click", "success": False, "message": f"Click timed out: {str(e)}"}
        except Exception as e:
            logger.error(f"Click action failed on selector {selector}: {e}")
            return {"action": "click", "success": False, "message": f"Click failed: {str(e)}"}

    def _handle_typing(self, page: Page, selector: str, text: str, task: Dict[str, Any]) -> Dict[str, Any]:
        try:
            page.wait_for_selector(selector, state="attached", timeout=self.default_timeout)
            page.fill(selector, text, timeout=self.default_timeout)
            return {"action": "type", "success": True, "message": f"Typed '{text}' into element."}
        except PlaywrightTimeoutError as e:
            if "search" in task.get("description", "").lower() or "search" in selector.lower():
                logger.info("Primary selector failed; using advanced fallback for element detection.")
                element = self._advanced_find_element(page, "search")
                if not element:
                    return {"action": "type", "success": False, "message": f"Typing failed: No search-like element found; error: {str(e)}"}
                try:
                    element.fill(text, timeout=self.default_timeout)
                    return {"action": "type", "success": True, "message": f"Typed '{text}' into fallback element."}
                except Exception as ex:
                    return {"action": "type", "success": False, "message": f"Typing failed on fallback element: {str(ex)}"}
            else:
                return {"action": "type", "success": False, "message": f"Typing failed: {str(e)}"}
        except Exception as e:
            logger.error(f"Typing action failed: {e}")
            return {"action": "type", "success": False, "message": f"Typing failed: {str(e)}"}

    def _handle_wait(self, seconds: str) -> Dict[str, Any]:
        try:
            wait_time = float(seconds)
            logger.info(f"Waiting for {wait_time} seconds")
            time.sleep(wait_time)
            return {"action": "wait", "success": True, "message": f"Waited {wait_time} seconds"}
        except ValueError as e:
            logger.error(f"Invalid wait time provided: {seconds}")
            return {"action": "wait", "success": False, "message": "Invalid wait time"}

    def _handle_wait_for_ajax(self, page: Page, seconds: str) -> Dict[str, Any]:
        try:
            timeout_seconds = int(seconds)
            logger.info(f"Waiting for AJAX/network activity for up to {timeout_seconds} seconds.")
            end_time = time.time() + timeout_seconds
            while time.time() < end_time:
                ajax_complete = page.evaluate("""
                    () => {
                        return (window.jQuery ? jQuery.active === 0 : true) &&
                               (typeof window.fetch === 'function' ? true : true);
                    }
                """)
                if ajax_complete:
                    break
                time.sleep(0.5)
            return {"action": "wait_for_ajax", "success": True, "message": "AJAX/network activity subsided."}
        except Exception as e:
            logger.error(f"Wait for AJAX failed: {e}")
            return {"action": "wait_for_ajax", "success": False, "message": f"Wait for AJAX failed: {str(e)}"}

    def _handle_scroll(self, page: Page, selector: str) -> Dict[str, Any]:
        try:
            if selector:
                page.wait_for_selector(selector, timeout=self.default_timeout)
                page.eval_on_selector(selector, "el => el.scrollIntoView({behavior: 'smooth', block: 'center'})")
                scroll_target = selector
            else:
                page.evaluate("window.scrollTo(0, document.body.scrollHeight);")
                scroll_target = "page bottom"
            return {"action": "scroll", "success": True, "message": f"Scrolled to {scroll_target}"}
        except Exception as e:
            logger.error(f"Scroll action failed on selector {selector}: {e}")
            return {"action": "scroll", "success": False, "message": f"Scroll failed: {str(e)}"}

    def _handle_hover(self, page: Page, selector: str) -> Dict[str, Any]:
        try:
            page.wait_for_selector(selector, state="visible", timeout=self.default_timeout)
            page.hover(selector, timeout=self.default_timeout)
            return {"action": "hover", "success": True, "message": f"Hovered over {selector}"}
        except Exception as e:
            logger.error(f"Hover action failed on selector {selector}: {e}")
            return {"action": "hover", "success": False, "message": f"Hover failed: {str(e)}"}

    def _handle_screenshot(self, page: Page, filename: str) -> Dict[str, Any]:
        try:
            page.screenshot(path=filename)
            return {"action": "screenshot", "success": True, "message": f"Screenshot saved as {filename}"}
        except Exception as e:
            logger.error(f"Screenshot capture failed: {e}")
            return {"action": "screenshot", "success": False, "message": f"Screenshot failed: {str(e)}"}

    def _handle_switch_tab(self, context, value: str) -> Dict[str, Any]:
        try:
            pages = context.pages
            if value.lower() == "new":
                target_page = pages[-1]
            else:
                idx = int(value)
                if idx < len(pages):
                    target_page = pages[idx]
                else:
                    return {"action": "switch_tab", "success": False, "message": f"Tab index {value} out of range"}
            return {"action": "switch_tab", "success": True, "message": f"Switched to tab {value}"}
        except Exception as e:
            logger.error(f"Switch tab failed: {e}")
            return {"action": "switch_tab", "success": False, "message": f"Switch tab failed: {str(e)}"}

    def _handle_execute_script(self, page: Page, script: str) -> Dict[str, Any]:
        try:
            result = page.evaluate(script)
            return {"action": "execute_script", "success": True, "message": "Script executed successfully", "result": result}
        except Exception as e:
            logger.error(f"Execute script failed: {e}")
            return {"action": "execute_script", "success": False, "message": f"Script execution failed: {str(e)}"}

    def _handle_drag_and_drop(self, page: Page, source_selector: str, target_selector: str) -> Dict[str, Any]:
        try:
            page.wait_for_selector(source_selector, timeout=self.default_timeout)
            page.wait_for_selector(target_selector, timeout=self.default_timeout)
            source = page.locator(source_selector)
            target = page.locator(target_selector)
            source.drag_to(target, timeout=self.default_timeout)
            return {"action": "drag_and_drop", "success": True, "message": f"Dragged element from {source_selector} to {target_selector}"}
        except Exception as e:
            logger.error(f"Drag and drop failed from {source_selector} to {target_selector}: {e}")
            return {"action": "drag_and_drop", "success": False, "message": f"Drag and drop failed: {str(e)}"}

    def _capture_failure_screenshot(self, page: Page, action: str):
        filename = f"failure_{action}_{int(time.time())}.png"
        try:
            page.screenshot(path=filename)
            logger.info(f"Failure screenshot captured: {filename}")
        except Exception as e:
            logger.error(f"Failed to capture screenshot: {e}")


    # Add this method inside your WebBrowserTool class
    def _advanced_find_element(self, page: Page, keyword: str):
        """
        Advanced fallback for finding an element.
        Searches candidate elements (input, textarea, button, a, div) and returns
        the element whose combined attributes best match the provided keyword.
        """
        try:
            # Get candidate elements
            candidates = page.query_selector_all("input, textarea, button, a, div")
            best_match = None
            best_ratio = 0.0
            for candidate in candidates:
                # Evaluate several attributes from the element
                attrs = page.evaluate(
                    """(el) => {
                        return {
                            id: el.id,
                            name: el.getAttribute('name'),
                            placeholder: el.getAttribute('placeholder'),
                            aria: el.getAttribute('aria-label'),
                            text: el.innerText
                        };
                    }""",
                    candidate,
                )
                combined_text = " ".join(
                    filter(None, [
                        attrs.get("id"),
                        attrs.get("name"),
                        attrs.get("placeholder"),
                        attrs.get("aria"),
                        attrs.get("text"),
                    ])
                )
                ratio = difflib.SequenceMatcher(None, combined_text.lower(), keyword.lower()).ratio()
                if ratio > best_ratio:
                    best_ratio = ratio
                    best_match = candidate
            if best_ratio > 0.5:
                self.logger.info(
                    f"Advanced fallback detected element with similarity {best_ratio:.2f} for keyword '{keyword}'"
                )
                return best_match
            return None
        except Exception as e:
            self.logger.error(f"Advanced find element error: {e}")
            return None
