import asyncio
import os
import random
import urllib
import urllib.parse as urlparse
from datetime import datetime
from typing import Literal

from playwright._impl._errors import TargetClosedError
from app.helpers.tool_helpers import run_shell
from app.tools.file import upload_to_presigned_url
from app.logger import logger
from app.tools.base import DEFAULT_WORKING_DIR
from app.types.browser_types import BrowserActionResult
from app.types.messages import BrowserActionRequest
from browser_use.browser.browser import Browser, BrowserConfig, BrowserContextConfig
from browser_use.browser.context import BrowserContext
from browser_use.controller.service import Controller


class BrowserDeadError(Exception):
    """Exception raised when the browser is not available or has crashed."""
    pass


class PageDeadError(Exception):
    """Exception raised when a page is no longer accessible."""
    pass


class ScreenshotError(Exception):
    """Error raised when a screenshot fails"""
    pass


class BrowserManager:
    """
    browser agent 基于 browser-use 和 playwright, 用于执行所有浏览器相关操作
    (browser agent based on browser-use and playwright, used to execute all browser-related operations)
    """
    browser: Browser
    browser_context: BrowserContext
    controller: Controller
    include_attributes: list[str]
    status: Literal['started', 'initializing', 'ready'] = 'started'

    def __init__(self, *, chrome_instance_path=None, headless=False, highlight_elements=True):
        """
        Initialize the browser manager.
        
        Args:
            chrome_instance_path: Path to Chrome browser instance
            headless: Whether to run the browser in headless mode
            highlight_elements: Whether to highlight elements in the browser
        """
        if not chrome_instance_path:
            chrome_instance_path = os.getenv('CHROME_INSTANCE_PATH', None)
        self.chrome_instance_path = chrome_instance_path
        self.headless = headless
        self.highlight_elements = highlight_elements
        self.browser = None
        self.browser_context = None
        self.controller = None
        self.include_attributes = ["id", "tagName", "href", "src", "alt", "ariaLabel", "placeholder", "name"]
        self.status = "started"  # started, initializing, ready

        logger.info(f"Browser manager initialized with Chrome path: {chrome_instance_path}")

    async def initialize(self):
        """
        Initialize the browser and create a new page.
        """
        if self.status in ("initializing", "ready"):
            return

        self.status = "initializing"
        logger.info("Initializing browser...")
        
        # 获取exts目录下的所有插件
        # extensions_dir = os.path.join(DEFAULT_WORKING_DIR, "exts")
        # extension_paths = []
        
        try:
            # # 遍历exts目录，获取所有子目录作为扩展
            # for item in os.listdir(extensions_dir):
            #     item_path = os.path.join(extensions_dir, item)
            #     if os.path.isdir(item_path) and not item.startswith('.'):
            #         extension_paths.append(item_path)
            
            # if not extension_paths:
            #     logger.warning("No extensions found in the extensions directory")
            # else:
            #     logger.info(f"Found {len(extension_paths)} extensions: {', '.join(os.path.basename(path) for path in extension_paths)}")
            
            # 将所有扩展路径合并为一个以逗号分隔的字符串
            # extensions_str = ",".join(extension_paths)
            
            browser_config = BrowserConfig(
                headless=self.headless,
                # extra_chromium_args=[
                #   f'--disable-extensions-except={extensions_str}',
                #   f'--load-extension={extensions_str}'
                # ],
                # chrome_instance_path=self.chrome_instance_path,
                # args=[
                #     "--disable-web-security",
                #     "--disable-features=IsolateOrigins,site-per-process",
                #     "--disable-site-isolation-trials"
                # ]
            )
            browser_context_config = BrowserContextConfig(
                # default_timeout_ms=30000,
                # default_navigation_timeout_ms=45000,
                # user_agent="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/132.0.0.0 Safari/537.36"
                highlight_elements=self.highlight_elements
            )
            self.browser = Browser(browser_config)
            
            # if self.browser.playwright_browser:
            #     logger.info(f'{self.browser.playwright_browser} ----')
            #     self.browser.playwright_browser.on('disconnected', self.recreate_page)
            
            self.browser_context = BrowserContext(self.browser, browser_context_config)
            self.controller = Controller()
            
           
            # Navigate to a blank page to prepare the context
            # await self.browser_context.goto("about:blank")
            await self.browser_context.get_session()
            self.status = "ready"
            
            logger.info("Browser initialized successfully")
        except Exception as e:
            self.status = "started"
            logger.error(f"Failed to initialize browser: {e}")
            raise BrowserDeadError(f"Failed to initialize browser: {e}")

    async def recreate_page(self):
        """
        Close the current page and create a new one.
        """
        if self.status != "ready":
            await self.initialize()

        logger.info("Recreating browser page")
        try:
            await self.browser_context.recreate_page()
            await self.browser_context.goto("about:blank")
            logger.info("Browser page recreated successfully")
        except Exception as e:
            logger.error(f"Failed to recreate page: {e}")
            raise PageDeadError(f"Failed to recreate page: {e}")

    async def execute_action(self, cmd: BrowserActionRequest) -> BrowserActionResult:
        """
        Execute a browser action.
        This version waits for the browser to be ready, checks page availability,
        calls the controller's act method, updates state, takes screenshots, uploads them,
        and then returns a BrowserActionResult with extended fields.
        """
        # Ensure browser is ready
        if self.status == "started":
            logger.info("Browser not initialized, starting initialization")
            await self.initialize()
            logger.info(f'{self.browser_context.context}-------')
            if self.browser_context.context:
                    self.browser_context.context.on('close', self.close)

        elif self.status != "ready":
            logger.info("Browser not ready, waiting for initialization")
            while self.status != "ready":
                await asyncio.sleep(0.2)

        # Check page availability
        logger.info("Check page availability")

        # Determine action type
        action = cmd.action
        action_type = None
        for field in action.__fields__:
            if getattr(action, field) is not None:
                action_type = field
                break
        if not action_type:
            raise ValueError("No action specified in the request")

        logger.info("Page available, executing the action")
        try:
            # Execute the action using the controller's act method
            result = await self.controller.act(action, self.browser_context)
            logger.info(f"Action execution finish, result {repr(result)}")

            # Update state
            logger.info("Updating state...")
            session = await self.browser_context.get_session()
            session.cached_state = await self.browser_context.get_state()
            cached_state = session.cached_state
            
            elements = cached_state.element_tree.clickable_elements_to_string()
            # if hasattr(cached_state, "clickable_elements"):
            #     elements = "\n".join(f"{el.index}[:]{el.description}" 
            #                          for el in cached_state.clickable_elements.values())

            # screenshot_save_path = ''
            screenshot_save_path = self.get_screenshot_save_path(cached_state.url)
            # logger.info("Taking screenshots", screenshot_save_path)
            try:
                clean_screenshot = await self.browser_context.take_screenshot(
                    save_path=screenshot_save_path
                )
            except ScreenshotError as e:
                logger.error(f"Error taking clean screenshot: {e}")
                clean_screenshot = b""
            try:
                marked_screenshot = await self.browser_context.take_screenshot()
            except ScreenshotError as e:
                logger.error(f"Error taking marked screenshot: {e}")
                marked_screenshot = b""

            logger.info(f"Screenshot saved to {screenshot_save_path}, uploading screenshots")
            await self.upload_screenshots(cmd, clean_screenshot, marked_screenshot)

            # Construct and return the BrowserActionResult with extended fields.
            return BrowserActionResult(
                url=cached_state.url,
                title=cached_state.title,
                result=(result.extracted_content if hasattr(result, "extracted_content") and result.extracted_content
                        else "success"),
                error=(result.error if hasattr(result, "error") else ""),
                markdown=(result.article_markdown if hasattr(result, "should_show_markdown") and result.should_show_markdown and hasattr(result, "article_markdown")
                          else ""),
                elements= elements if result.include_elements else '',
                screenshot_uploaded=bool(cmd.screenshot_presigned_url),
                clean_screenshot_uploaded=bool(cmd.clean_screenshot_presigned_url),
                clean_screenshot_path=screenshot_save_path,
                pixels_above=getattr(cached_state, "pixels_above", 0),
                pixels_below=getattr(cached_state, "pixels_below", 0)
            )
        except TargetClosedError:
            logger.error("Browser page closed unexpectedly")
            raise PageDeadError("Browser page closed unexpectedly")
        except PageDeadError as e:
            logger.error(f"Page dead error: {e}")
            await self.recreate_page()
            raise
        except Exception as e:
            logger.error(f"Error executing browser action: {e}")
            raise

    async def restart_browser(self):
        """
        Restart the browser by closing and reinitializing it.
        This version also attempts to restart Chrome via a shell command.
        """
        logger.info("Restarting browser")
        if self.browser:
            try:
                # await self.browser.close()
                # self.browser = None
                if self.browser_context:
                    await self.browser_context.close()
                self.browser_context = None
                self.controller = None
                self.status = "started"
            except Exception as e:
                logger.error(f"Error closing browser: {e}")

        # Attempt to restart Chrome via supervisorctl
        logger.info("Try restart chrome")
        try:
            output = await run_shell("sudo supervisorctl restart chrome")
            logger.info(output)
            logger.info("Chrome restarted")
        except Exception as e:
            logger.error(f"Error restarting chrome: {e}")

        # Reinitialize the browser
        await self.initialize()
        logger.info("Browser restarted")

    async def health_check(self) -> bool:
        """
        Check if the browser is healthy.
        """
        if self.status != "ready":
            return False

        try:
            if not self.browser or not self.browser_context or not self.controller:
                return False

            await self.browser_context.evaluate_javascript("1 + 1")
            return True
        except Exception as e:
            logger.error(f"Browser health check failed: {e}")
            return False

    def get_screenshot_save_path(self, page_url: str) -> str:
        """
        Generate a path for saving a screenshot.
        """
        screenshots_dir = f"{DEFAULT_WORKING_DIR}/screenshots"
        os.makedirs(screenshots_dir, exist_ok=True)

        parsed_url = urlparse.urlparse(page_url)
        timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        random_id = random.randint(1000, 9999)

        if parsed_url.scheme in ('http', 'https') and parsed_url.hostname:
            hostname = parsed_url.hostname.replace('.com', '').replace('www.', '').replace('.', '_')
        else:
            hostname = page_url.split('/').pop().replace('.', '_')

        return f"{screenshots_dir}/{hostname}_{timestamp}_{random_id}.webp"

    async def upload_screenshots(self, cmd: BrowserActionRequest, clean_screenshot: bytes, marked_screenshot: bytes):
        """
        Upload screenshots to the provided presigned URLs.
        """
        if cmd.screenshot_presigned_url:
            try:
                await upload_to_presigned_url(
                    marked_screenshot,
                    cmd.screenshot_presigned_url
                )
                logger.info(f"Screenshot uploaded successfully to {cmd.screenshot_presigned_url}")
            except Exception as e:
                logger.error("Failed to upload screenshot")
        else:
            logger.info("No presigned URL provided for screenshot, skipped uploading")

        if cmd.clean_screenshot_presigned_url:
            try:
                await upload_to_presigned_url(
                    clean_screenshot,
                    cmd.clean_screenshot_presigned_url
                )
                logger.info(f"Clean screenshot uploaded successfully to {cmd.clean_screenshot_presigned_url}")
            except Exception as e:
                logger.error("Failed to upload clean screenshot")
        else:
            logger.info("No presigned clean URL provided for screenshot, skipped uploading")

    async def close(self):
        """
        关闭浏览器实例
        """
        logger.info("Closing browser")
        if self.browser:
            try:
                if self.browser_context:
                    await self.browser_context.close()
                # await self.browser.close()
                self.browser = None
                self.browser_context = None
                self.controller = None
                self.status = "started"
                logger.info("Browser closed successfully")
            except Exception as e:
                logger.error(f"Error closing browser: {e}")
        else:
            logger.info("Browser not initialized, no need to close")
