from DrissionPage import Chromium, ChromiumOptions

class Kimi:
    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://kimi.moonshot.cn/')
        chat_input = self.tab.ele('.editorContentEditable___FZJd9')
        chat_input.wait.enabled()
        chat_input.input(input_message)
        self.tab.wait(10)
        self.tab.ele('#send-button').click()
        self.tab.wait(wait_time)
        chat_input.input(' ')
        self.tab.ele('#send-button').wait.enabled()
        output = self.tab.ele('#^chat-markdown').text
        return output

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

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

    print(response)