from DrissionPage import Chromium, ChromiumOptions

class DouBao:
    def __init__(self, base_timeout=300):
        self.base_timeout = base_timeout
        """Initialize the Chromium browser with specified timeout settings."""
        co = ChromiumOptions().auto_port()
        co.set_timeouts(base=self.base_timeout)
        self.browser = Chromium(addr_or_opts=co)
        self.tab = self.browser.latest_tab

    def search(self, input_message, wait_time=10):
        """Send a message through the chat interface and retrieve the response."""
        self.tab.get('https://www.doubao.com/')
        chat_input = self.tab.ele('@data-testid=chat_input_input')
        chat_input.input(input_message)
        self.tab.ele('#flow-end-msg-send').click()
        self.tab.wait(wait_time)
        self.tab.ele('@data-testid=chat_input_local_break_button').wait.hidden()
        output = self.tab.eles('@data-testid=message_text_content')[1].text
        return output

    def quit(self):
        self.browser.quit()

if __name__ == "__main__":
    dou_bao = DouBao()
    message = '今日新闻'
    response = dou_bao.search(message)

    print(response)