from DrissionPage import Chromium, ChromiumOptions

class TongYi:
    def __init__(self, base_timeout=30):
        self.base_timeout = base_timeout
        """Initialize the Chromium browser with specified timeout settings."""
        # co = ChromiumOptions().auto_port()
        co = ChromiumOptions()
        co.set_timeouts(base=self.base_timeout)
        self.browser = Chromium(addr_or_opts=co)
        self.tab = self.browser.latest_tab
        self.tab.get('https://tongyi.aliyun.com/')

    def chat(self, input_message, wait_time=10):
        """Send a message through the chat interface and retrieve the response."""
        chat_input = self.tab.ele('@placeholder=千事不决问通义')
        chat_input.input(input_message)
        self.tab.ele('.operateBtn--zFx6rSR0').click()
        self.tab.wait(wait_time)
        chat_input.input('over')
        self.tab.ele('.operateBtn--zFx6rSR0').wait.enabled()
        output = self.tab.ele('.containerWrap--lFLVsVCe').text
        return output

    def generate(self, input_message, wait_time=10):
        return self.chat(input_message, wait_time)

    def search(self, input_message, wait_time=10):
        return self.chat(input_message, wait_time)

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


if __name__ == "__main__":
    tong_yi = TongYi()
    message = '今日新闻'
    response = tong_yi.search(message)
    tong_yi.quit()

    print(response)