from DrissionPage import Chromium, ChromiumOptions
from DrissionPage.errors import ElementNotFoundError


class Kimi:
    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://kimi.moonshot.cn/')
        self.first_time = True

    def chat(self, input_message, wait_time=10):
        """Send a message through the chat interface and retrieve the response."""
        try:
            chat_input = self.tab.ele('.editorContentEditable___FZJd9')
            chat_input.wait.enabled()
            chat_input.input(input_message)
            self.tab.wait(3)
            self.tab.ele('#send-button').wait.enabled()
            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
        except ElementNotFoundError:
            output = ''
        return output

    def generate(self, input_message, wait_time=10):
        if not self.first_time:
            new_session_btn = self.tab.ele('.myAgentTool___Y1_mC')
            new_session_btn.wait.enabled()
            new_session_btn.click()
            self.tab.wait(3)
        self.first_time = False
        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__":
    kimi = Kimi()
    message = '今日新闻'
    response = kimi.search(message)
    kimi.quit()

    print(response)