"""
ChatBrowserUse - Client for browser-use cloud API

This wraps the BaseChatModel protocol and sends requests to the browser-use cloud API
for optimized browser automation LLM inference.
"""

import logging
import os
from typing import TypeVar, overload

import httpx
from pydantic import BaseModel

from openbrowser.llm.base import BaseChatModel
from openbrowser.llm.messages import BaseMessage
from openbrowser.llm.views import ChatInvokeCompletion
from openbrowser.observability import observe

T = TypeVar('T', bound=BaseModel)

logger = logging.getLogger(__name__)


class ChatBrowserUse(BaseChatModel):
	"""
	Client for browser-use cloud API.

	This sends requests to the browser-use cloud API which uses optimized models
	and prompts for browser automation tasks.

	Usage:
		agent = Agent(
			task="Find the number of stars of the browser-use repo",
			llm=ChatBrowserUse(model='bu-latest'),
		)
	"""

	def __init__(
		self,
		model: str = 'bu-latest',
		api_key: str | None = None,
		base_url: str | None = None,
		timeout: float = 120.0,
		**kwargs,
	):
		"""
		Initialize ChatBrowserUse client.

		Args:
			model: Model name to use. Options: 'bu-latest', 'bu-1-0'. Defaults to 'bu-latest'.
			api_key: API key for browser-use cloud. Defaults to BROWSER_USE_API_KEY env var.
			base_url: Base URL for the API. Defaults to BROWSER_USE_LLM_URL env var or production URL.
			timeout: Request timeout in seconds.
		"""
		# Validate model name
		valid_models = ['bu-latest', 'bu-1-0']
		if model not in valid_models:
			raise ValueError(f"Invalid model: '{model}'. Must be one of {valid_models}")

		self.model = 'bu-1-0' if model == 'bu-latest' else model  # must update on new model releases
		self.fast = False
		self.api_key = api_key or os.getenv('BROWSER_USE_API_KEY')
		self.base_url = base_url or os.getenv('BROWSER_USE_LLM_URL', 'https://llm.api.browser-use.com')
		self.timeout = timeout

		if not self.api_key:
			raise ValueError(
				'You need to set the BROWSER_USE_API_KEY environment variable. '
				'Get your key at https://cloud.browser-use.com/new-api-key'
			)

	@property
	def provider(self) -> str:
		return 'browser-use'

	@property
	def name(self) -> str:
		return self.model

	@overload
	async def ainvoke(self, messages: list[BaseMessage], output_format: None = None) -> ChatInvokeCompletion[str]: ...

	@overload
	async def ainvoke(self, messages: list[BaseMessage], output_format: type[T]) -> ChatInvokeCompletion[T]: ...

	@observe(name='chat_openbrowser_ainvoke')
	async def ainvoke(
		self, messages: list[BaseMessage], output_format: type[T] | None = None
	) -> ChatInvokeCompletion[T] | ChatInvokeCompletion[str]:
		"""
		Send request to browser-use cloud API.

		Args:
			messages: List of messages to send
			output_format: Expected output format (Pydantic model)

		Returns:
			ChatInvokeCompletion with structured response and usage info
		"""
		# Prepare request payload
		payload = {
			'messages': [self._serialize_message(msg) for msg in messages],
			'fast': self.fast,
		}

		# Add output format schema if provided
		if output_format is not None:
			payload['output_format'] = output_format.model_json_schema()

		# Make API request
		async with httpx.AsyncClient(timeout=self.timeout) as client:
			try:
				response = await client.post(
					f'{self.base_url}/v1/chat/completions',
					json=payload,
					headers={
						'Authorization': f'Bearer {self.api_key}',
						'Content-Type': 'application/json',
					},
				)
				response.raise_for_status()
				result = response.json()

			except httpx.HTTPStatusError as e:
				error_detail = ''
				try:
					error_data = e.response.json()
					error_detail = error_data.get('detail', str(e))
				except Exception:
					error_detail = str(e)

				error_msg = ''
				if e.response.status_code == 401:
					error_msg = f'Invalid API key. {error_detail}'
				elif e.response.status_code == 402:
					error_msg = f'Insufficient credits. {error_detail}'
				else:
					error_msg = f'API request failed: {error_detail}'

				raise ValueError(error_msg)

			except httpx.TimeoutException:
				error_msg = f'Request timed out after {self.timeout}s'
				raise ValueError(error_msg)

			except Exception as e:
				error_msg = f'Failed to connect to browser-use API: {e}'
				raise ValueError(error_msg)

			# Parse response - server returns structured data as dict
			if output_format is not None:
				# Server returns structured data as a dict, validate it
				completion_data = result['completion']
				logger.debug(
					f'📥 Got structured data from service: {list(completion_data.keys()) if isinstance(completion_data, dict) else type(completion_data)}'
				)

				# Convert action dicts to ActionModel instances if needed
				# llm-use returns dicts to avoid validation with empty ActionModel
				if isinstance(completion_data, dict) and 'action' in completion_data:
					actions = completion_data['action']
					if actions and isinstance(actions[0], dict):
						from typing import get_args

						# Get ActionModel type from output_format
						action_model_type = get_args(output_format.model_fields['action'].annotation)[0]

						# Convert dicts to ActionModel instances
						completion_data['action'] = [action_model_type.model_validate(action_dict) for action_dict in actions]

				completion = output_format.model_validate(completion_data)
			else:
				completion = result['completion']

			# Parse usage info
			usage = None
			if 'usage' in result:
				from openbrowser.llm.views import ChatInvokeUsage

				usage = ChatInvokeUsage(**result['usage'])

		return ChatInvokeCompletion(
			completion=completion,
			usage=usage,
		)

	def _serialize_message(self, message: BaseMessage) -> dict:
		"""Serialize a message to JSON format."""
		# Handle Union types by checking the actual message type
		msg_dict = message.model_dump()
		return {
			'role': msg_dict['role'],
			'content': msg_dict['content'],
		}
