import os
import sys
from pathlib import Path

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import asyncio

from langchain_openai import ChatOpenAI
from langchain.callbacks import get_openai_callback
from typing import Dict, List

from browser_use import Agent, Controller
from browser_use.browser.browser import Browser, BrowserConfig
from browser_use.browser.context import BrowserContext

# Initialize controller first
browser = Browser(config=BrowserConfig(headless=False))
controller = Controller()

# Track downloads
downloaded_files: List[str] = []

async def handle_download(download):
    # Get the executable directory
    exe_dir = Path(os.path.dirname(os.path.abspath(__file__)))
    
    # Get original download path
    original_path = await download.path()
    if original_path:
        # Create new path in executable directory
        new_path = exe_dir / os.path.basename(original_path)
        
        # Move the file to executable directory
        os.rename(original_path, new_path)
        
        # Add the new path to downloaded files list
        downloaded_files.append(str(new_path))
        print(f"Downloaded and moved to: {new_path}")

@controller.action(
	'Upload file - the file name is inside the function - you only need to call this with the  correct index',
	requires_browser=True,
)
async def upload_file(index: int, browser: BrowserContext):
	element = await browser.get_element_by_index(index)
	my_file = Path.cwd() / 'examples/test_cv.txt'
	if not element:
		raise Exception(f'Element with index {index} not found')

	await element.set_input_files(str(my_file.absolute()))
	return f'Uploaded file to index {index}'


@controller.action('Close file dialog', requires_browser=True)
async def close_file_dialog(browser: BrowserContext):
	page = await browser.get_current_page()
	await page.keyboard.press('Escape')

def handle_page(new_page):
	print("New page created!")
	new_page.on("download", lambda download: asyncio.create_task(handle_download(download)))

async def print_contexts_after_delay(playwright_browser):
	await asyncio.sleep(5)
	if (len(playwright_browser.contexts) < 1):
		raise Exception('No contexts found')

	#  up download handler at Playwright browser level
	playwright_browser.contexts[0].on("page", handle_page)
	print('contexts after 5 sec:', len(playwright_browser.contexts))

async def main():
	sites = 'https://some.url'
	task = f"""go to {sites}
	Find documents and download them one-by-one.
	Once files are downloaded, or if downloading is not accessible without login, stop"""

	model = ChatOpenAI(model='gpt-4o')
	agent = Agent(
		task=task,
		llm=model,
		controller=controller,
		browser=browser,
	)

	# Get the underlying Playwright browser instance
	playwright_browser = await browser.get_playwright_browser()
	print('contexts initial:', len(playwright_browser.contexts))
	
	# Create task for delayed context printing
	asyncio.create_task(print_contexts_after_delay(playwright_browser))
	
	# Track token usage and cost
	with get_openai_callback() as cb:
		await agent.run()
		
		# Print usage statistics
		print("\nOpenAI API Usage Stats:")
		print(f"Total Tokens: {cb.total_tokens}")
		print(f"Prompt Tokens: {cb.prompt_tokens}")
		print(f"Completion Tokens: {cb.completion_tokens}")
		print(f"Total Cost (USD): ${cb.total_cost:.4f}")

	history_file_path = 'AgentHistoryList.json'
	agent.save_history(file_path=history_file_path)

	await browser.close()

	# Print downloaded files
	if downloaded_files:
		print("\nDownloaded files:")
		for file_path in downloaded_files:
			print(file_path)
			print(f"- {os.path.basename(file_path)}")
	else:
		print("\nNo files were downloaded during this session.")

	input('Press Enter to close...')


if __name__ == '__main__':
	asyncio.run(main())