import os
from typing import Optional
from revChatGPT.V3 import Chatbot
from feathr.chat.notebook_utils import *


class FeathrChat(object):
    def __init__(self):
        key = self.get_api_key()
        self.chat_bot = None
        self.feathr_client = None
        if key:
            self.chat_bot = Chatbot(key)

    def get_api_key(self):
        return os.getenv("CHATGPT_API_KEY", None)

    def ask_llm_in_notebook(self, question_with_prompt: str):
        """Ask LLM model a question within the notebook"""
        if not self.chat_bot:
            key = self.get_api_key()
            if not key:
                raise RuntimeError(
                    "Please set environment variable CHATGPT_API_KEY before using Feathr Chat. You can get your API key for ChatGPT at https://platform.openai.com/account/api-keys. For example, run: os.environ['CHATGPT_API_KEY'] = 'your api key' and retry."
                )
            self.chat_bot = Chatbot(key)

        content = self.chat_bot.ask(question_with_prompt)
        if "```" in content and self.is_a_code_gen_question(question_with_prompt):
            # for debugging
            # res = content
            # print(re.sub(r"```.*```", "", res, flags=re.DOTALL))
            code = extract_code_from_string(content)
            create_new_cell(code)
        else:
            print(content)

    def is_a_code_gen_question(self, question_with_prompt):
        return "explain" not in question_with_prompt.lower() and "what" not in question_with_prompt.lower()