import asyncio
import logging
from importlib import resources
from typing import TYPE_CHECKING, Optional, Callable, Awaitable
from urllib.parse import urlparse

if TYPE_CHECKING:
	from browser_use.browser.types import Page


from browser_use.dom.views import (
	DOMBaseNode,
	DOMElementNode,
	DOMState,
	DOMTextNode,
	SelectorMap,
	ViewportInfo,
)
from browser_use.observability import observe_debug
from browser_use.utils import is_new_tab_page, time_execution_async
from browser_use.dom.dom_utils import DomUtils, FramesDescriptorDict
from playwright.async_api import Frame, JSHandle

# @dataclass
# class ViewportInfo:
# 	width: int
# 	height: int


class DomService:
	logger: logging.Logger

	def __init__(self, page: 'Page', logger: logging.Logger | None = None):
		self.page = page
		self.xpath_cache = {}
		self.logger = logger or logging.getLogger(__name__)

		self.js_code = resources.files('browser_use.dom.dom_tree').joinpath('index.js').read_text()

	# region - Clickable elements
	@observe_debug(ignore_input=True, ignore_output=True, name='get_clickable_elements')
	@time_execution_async('--get_clickable_elements')
	async def get_clickable_elements(
		self,
		highlight_elements: bool = True,
		focus_element: int = -1,
		viewport_expansion: int = 0,
	) -> DOMState:
		element_tree, selector_map = await self._build_dom_tree(highlight_elements, focus_element, viewport_expansion)
		return DOMState(element_tree=element_tree, selector_map=selector_map)

	@time_execution_async('--get_cross_origin_iframes')
	async def get_cross_origin_iframes(self) -> list[str]:
		# invisible cross-origin iframes are used for ads and tracking, dont open those
		hidden_frame_urls = await self.page.locator('iframe').filter(visible=False).evaluate_all('e => e.map(e => e.src)')

		is_ad_url = lambda url: any(
			domain in urlparse(url).netloc for domain in ('doubleclick.net', 'adroll.com', 'googletagmanager.com')
		)

		return [
			frame.url
			for frame in self.page.frames
			if urlparse(frame.url).netloc  # exclude data:urls and new tab pages
			and urlparse(frame.url).netloc != urlparse(self.page.url).netloc  # exclude same-origin iframes
			and frame.url not in hidden_frame_urls  # exclude hidden frames
			and not is_ad_url(frame.url)  # exclude most common ad network tracker frame URLs
		]

	@time_execution_async('--build_dom_tree')
	async def _build_dom_tree(
		self,
		highlight_elements: bool,
		focus_element: int,
		viewport_expansion: int,
	target_frame: Optional[Frame] = None, highlight_index: int = 0, initial_root_node: Optional[JSHandle] = None) -> tuple[DOMElementNode, SelectorMap]:
		if await self.page.evaluate('1+1') != 2:
			raise ValueError('The page cannot evaluate javascript code properly')

		if is_new_tab_page(self.page.url) or self.page.url.startswith('chrome://'):
			# short-circuit if the page is a new empty tab or chrome:// page for speed, no need to inject buildDomTree.js
			return (
				DOMElementNode(
					tag_name='body',
					xpath='',
					attributes={},
					children=[],
					is_visible=False,
					parent=None,
				),
				{},
			)

		# NOTE: We execute JS code in the browser to extract important DOM information.
		#       The returned hash map contains information about the DOM tree and the
		#       relationship between the DOM elements.
		debug_mode = self.logger.getEffectiveLevel() == logging.DEBUG
		args = {
			'doHighlightElements': highlight_elements,
			'focusHighlightIndex': focus_element,
			'viewportExpansion': viewport_expansion,
			'debugMode': debug_mode,'initialRootNode': initial_root_node, 'highlightIndex': highlight_index
		}

		try:
			self.logger.debug(f'🔧 Starting JavaScript DOM analysis for {self.page.url[:50]}...')
			if target_frame:
				eval_page: dict = await target_frame.evaluate(self.js_code, args)
			else:
				eval_page: dict = await self.page.evaluate(self.js_code, args)
			self.logger.debug('✅ JavaScript DOM analysis completed')
		except Exception as e:
			self.logger.error('Error evaluating JavaScript: %s', e)
			raise

		# Only log performance metrics in debug mode
		if debug_mode and 'perfMetrics' in eval_page:
			perf = eval_page['perfMetrics']

			# Get key metrics for summary
			total_nodes = perf.get('nodeMetrics', {}).get('totalNodes', 0)
			# processed_nodes = perf.get('nodeMetrics', {}).get('processedNodes', 0)

			# Count interactive elements from the DOM map
			interactive_count = 0
			if 'map' in eval_page:
				for node_data in eval_page['map'].values():
					if isinstance(node_data, dict) and node_data.get('isInteractive'):
						interactive_count += 1

			# Create concise summary
			url_short = self.page.url[:50] + '...' if len(self.page.url) > 50 else self.page.url
			self.logger.debug(
				'🔎 Ran buildDOMTree.js interactive element detection on: %s interactive=%d/%d\n',
				url_short,
				interactive_count,
				total_nodes,
				# processed_nodes,
			)

		self.logger.debug('🔄 Starting Python DOM tree construction...')
		result = await self._construct_dom_tree(eval_page)
		self.logger.debug('✅ Python DOM tree construction completed')
		return result

	@time_execution_async('--construct_dom_tree')
	async def _construct_dom_tree(
		self,
		eval_page: dict,
	) -> tuple[DOMElementNode, SelectorMap]:
		js_node_map = eval_page['map']
		js_root_id = eval_page['rootId']

		selector_map = {}
		node_map = {}

		for id, node_data in js_node_map.items():
			node, children_ids = self._parse_node(node_data)
			if node is None:
				continue

			node_map[id] = node

			if isinstance(node, DOMElementNode) and node.highlight_index is not None:
				selector_map[node.highlight_index] = node

			# NOTE: We know that we are building the tree bottom up
			#       and all children are already processed.
			if isinstance(node, DOMElementNode):
				for child_id in children_ids:
					if child_id not in node_map:
						continue

					child_node = node_map[child_id]

					child_node.parent = node
					node.children.append(child_node)

		html_to_dict = node_map[str(js_root_id)]

		del node_map
		del js_node_map
		del js_root_id

		if html_to_dict is None or not isinstance(html_to_dict, DOMElementNode):
			raise ValueError('Failed to parse HTML to dictionary')

		return html_to_dict, selector_map

	def _parse_node(
		self,
		node_data: dict,
	) -> tuple[DOMBaseNode | None, list[int]]:
		if not node_data:
			return None, []

		# Process text nodes immediately
		if node_data.get('type') == 'TEXT_NODE':
			text_node = DOMTextNode(
				text=node_data['text'],
				is_visible=node_data['isVisible'],
				parent=None,
			)
			return text_node, []

		# Process coordinates if they exist for element nodes

		viewport_info = None

		if 'viewport' in node_data:
			viewport_info = ViewportInfo(
				width=node_data['viewport']['width'],
				height=node_data['viewport']['height'],
			)

		element_node = DOMElementNode(
			tag_name=node_data['tagName'],
			xpath=node_data['xpath'],
			attributes=node_data.get('attributes', {}),
			children=[],
			is_visible=node_data.get('isVisible', False),
			is_interactive=node_data.get('isInteractive', False),
			is_top_element=node_data.get('isTopElement', False),
			is_in_viewport=node_data.get('isInViewport', False),
			highlight_index=node_data.get('highlightIndex'),
			shadow_root=node_data.get('shadowRoot', False),
			parent=None,
			viewport_info=viewport_info,
		)

		children_ids = node_data.get('children', [])

		return element_node, children_ids
	
	@time_execution_async('--get_multitarget_clickable_elements')
	async def get_multitarget_clickable_elements(
	  self,
	  highlight_elements: bool = True,
	  focus_element: int = -1,
	  viewport_expansion: int = 0,
	  remove_highlights: Optional[Callable[..., Awaitable[None]]] = None,
	) -> DOMState:
		dom_utils = DomUtils()

		frames_descriptor_dict:FramesDescriptorDict = await dom_utils.build_frames_descriptor_dict(self.page)

		if remove_highlights:
			tasks = [] # Trying to minimize the ugly visual effect by parallelizing the execution ...
			for frame in frames_descriptor_dict.keys():
				tasks.append(remove_highlights(frame))
			await asyncio.gather(*tasks)

		final_dom_element_node, dom_element_node, final_selector_map, highlight_index = None, None, {}, 0
		for frame, closed_shadow_roots in frames_descriptor_dict.items():
			# look in 'final_dom_element_node' for the point to link this new 'document.body' ...
			iframe_element = await DomUtils.get_insertion_point_for_body(final_dom_element_node, frame)
			if frame == self.page.main_frame or iframe_element:
				# If there is no iframe_element there is no point in doing anything ...
				# Always evaluating in document.body ...
				self.logger.info(f"Evaluating in frame with url=[{frame.url}] using document.body ...")
				dom_element_node, selector_map =         await self._build_dom_tree(highlight_elements, focus_element, viewport_expansion, frame, highlight_index)
				highlight_index += len(selector_map)
				final_selector_map.update(selector_map)
				if frame == self.page.main_frame:
					final_dom_element_node = dom_element_node
				else:
					assert final_dom_element_node is not None
					if iframe_element:
						# Verify if iframe_element has a 'html' child, which in turn has a 'body' child.
						body = await DomUtils.traverse_and_filter(iframe_element,
						                                          lambda node: asyncio.sleep(0, result=(node.xpath == "html/body")),
						                                          just_first_found=True)
					if body:
						DomUtils.copy_children(dom_element_node, body[0])
					else:
						# We link here the document.body itself ... it's more elegant ;-|
						dom_element_node.parent = iframe_element
						iframe_element.children.append(dom_element_node)

				# Dealing with closed ShadowRoot objects in the Frame ...
				for closed_shadow_root in closed_shadow_roots:
					self.logger.info(f"Evaluating in frame with url=[{frame.url}] using specific root node {closed_shadow_root.element_handle_to_shadow_root} ...")
					dom_element_node, selector_map =           await self._build_dom_tree(highlight_elements, focus_element, viewport_expansion, frame, highlight_index,
					                closed_shadow_root.element_handle_to_shadow_root)
					highlight_index += len(selector_map)
					final_selector_map.update(selector_map)
					# Look in 'final_dom_element_node' for the point to link the 'dom_element_node' corresponding to the closed ShadowRoot
					# HERE THE MATCHING IS EASY: LOOK FOR A MATCHING "xpath" IN 'final_dom_element_node' AND ADD TO THE FOUND
					# DOMElementNode THE CHILDREN OF 'dom_element_node'
					host_elements: list[DOMElementNode] =           await DomUtils.traverse_and_filter(final_dom_element_node,
					                  lambda node, target_xpath: asyncio.sleep(0, result=(node.xpath == target_xpath)),
					                  # This is passed as an argument to the lambda (not needed it's here as an example)
					                  dom_element_node.xpath)
					if host_elements and len(host_elements) > 1:
						# If there is more than one matching xpath the Frame must match also ...
						host_elements = [host for host in host_elements if DomUtils.is_matching_iframe(frame, await DomUtils.find_parent_iframe(host))]
					assert len(host_elements) == 1, (
					    f"There should be one and only one element matching the xpath [{dom_element_node.xpath}] for the closed shadow root...")
					host = host_elements[0]
					host.shadow_root = True
					DomUtils.copy_children(dom_element_node, host)
					await closed_shadow_root.element_handle_to_shadow_root.dispose()

		# After connecting the different element trees we return the root one ...
		assert final_dom_element_node is not None
		return DOMState(element_tree=final_dom_element_node, selector_map=final_selector_map)
