from openai import OpenAI
import os
from pydantic import BaseModel, ConfigDict
from .printer import Printer
from .constants import (
    DEFAULT_ANSWER_LENGTH_PARAGRAPHS,
    DEFAULT_LINES_PER_PARAGRAPH,
    MAX_ANSWER_LENGTH_PARAGRAPHS,
    MAX_LINES_PER_PARAGRAPH,
    DEFAULT_PERPLEXITY_MODEL,
    MAX_CONTEXT_TOKENS,
    DEFAULT_AUTO_PROMPT_AUGMENT,
    DEFAULT_PROMPT_AUGMENTATION_TEMPERATURE,
    DEFAULT_PROMPT_AUGMENTATION_MODEL,
    DEFAULT_ENABLE_UNBIASED_REASONING,
    DEFAULT_DEEP_RESEARCH_ENABLED,
    DEFAULT_DEEP_RESEARCH_MODEL,
    DEFAULT_AUXKNOW_SYSTEM_PROMPT,
    DEFAULT_FAST_MODE_ENABLED,
)
from dotenv import load_dotenv
from copy import deepcopy
from typing import Generator, Optional, Union
from uuid import uuid4
import re


class AuxKnowAnswer(BaseModel):
    """AuxKnowAnswer class to store the response from the AuxKnow API.

    Attributes:
        is_final (bool): Indicates if the answer is final.
        answer (str): The answer text.
        citations (list[str]): List of citations for the answer.
    """

    is_final: bool = False
    answer: str
    citations: list[str]


class AuxKnowConfig(BaseModel):
    """AuxKnowConfig class to store the configuration for AuxKnow.

    Attributes:
        auto_model_routing (bool): Whether to automatically route queries to the appropriate model.
        auto_query_restructuring (bool): Whether to automatically restructure queries.
        answer_length_in_paragraphs (int): The length of the answer in paragraphs.
        lines_per_paragraph (int): The number of lines per paragraph.
        fast_mode (bool): When enabled, overrides all other settings to provide fastest possible response.
    """

    auto_model_routing: bool = True
    auto_query_restructuring: bool = True
    answer_length_in_paragraphs: int = DEFAULT_ANSWER_LENGTH_PARAGRAPHS
    lines_per_paragraph: int = DEFAULT_LINES_PER_PARAGRAPH
    auto_prompt_augment: bool = DEFAULT_AUTO_PROMPT_AUGMENT
    enable_unibiased_reasoning: bool = DEFAULT_ENABLE_UNBIASED_REASONING
    fast_mode: bool = DEFAULT_FAST_MODE_ENABLED


class AuxKnowSession(BaseModel):
    """AuxKnowSession class to manage a session with AuxKnow.

    Attributes:
        session_id (str): The unique identifier for the session.
        context (list[dict[str, str]]): The context of the session as a list of question-answer pairs.
        auxknow (AuxKnow): The AuxKnow instance associated with the session.
        closed (bool): Indicates if the session is closed.
    """

    session_id: str
    context: list[dict[str, str]] = []
    auxknow: "AuxKnow"
    closed: bool = False

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def ask(
        self, question: str, deep_research=False, fast_mode=DEFAULT_FAST_MODE_ENABLED
    ) -> AuxKnowAnswer:
        """Ask a question within this session to maintain context.

        Args:
            question (str): The question to ask.
            deep_research (bool): Whether to enable deep research mode. (Default: False)
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            AuxKnowAnswer: The answer.
        """
        if self.closed:
            raise ValueError("Cannot ask a question on a closed session.")
        return self.auxknow._ask_with_context(
            self, question, deep_research=deep_research, fast_mode=fast_mode
        )

    def ask_stream(
        self,
        question: str,
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
        fast_mode=DEFAULT_FAST_MODE_ENABLED,
    ) -> Generator[AuxKnowAnswer, None, None]:
        """Ask a question within this session to maintain context with streaming response.

        Args:
            question (str): The question to ask.
            deep_research (bool): Whether to enable deep research mode. (Default: False)
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            Generator[AuxKnowAnswer, None, None]: A generator of answers.
        """
        if self.closed:
            raise ValueError("Cannot ask a question on a closed session.")
        return self.auxknow._ask_with_context_stream(
            self, question, deep_research=deep_research, fast_mode=fast_mode
        )

    def close(self) -> None:
        """Close the session."""
        self.auxknow._close_session(self)


class AuxKnow:
    """AuxKnow a simpler Answer Engine built on top of Perplexity.

    Attributes:
        verbose (bool): Whether to enable verbose logging.
        config (AuxKnowConfig): The configuration for AuxKnow.
        sessions (dict): A dictionary to store active sessions.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        openai_api_key: Optional[str] = None,
        verbose: bool = False,
        auto_prompt_augment: bool = DEFAULT_AUTO_PROMPT_AUGMENT,
    ):
        """Initialize the AuxKnow instance.

        Args:
            api_key (Optional[str]): The API key for Perplexity.
            openai_api_key (Optional[str]): The API key for OpenAI.
            verbose (bool): Whether to enable verbose logging.
        """
        self.verbose = verbose
        self.config = AuxKnowConfig()

        if self.verbose:
            Printer.print_orange_message("🧠 Initializing AuxKnow API! 🤯")

        if self.verbose:
            Printer.print_light_grey_message("🌴 Loading environment variables... ")

        self.__load_environment_variables()

        perplexity_api_key = api_key or os.getenv("PERPLEXITY_API_KEY")
        self.initialized = False

        if not perplexity_api_key:
            Printer.print_yellow_message(
                "PERPLEXITY_API_KEY not found in environment variables. Cannot use AuxKnow."
            )
            return

        if not openai_api_key:
            openai_api_key = os.getenv("OPENAI_API_KEY")

        if not openai_api_key:
            Printer.print_yellow_message(
                "OPENAI_API_KEY not found in environment variables. Cannot use AuxKnow."
            )
            return

        self.config.auto_prompt_augment = auto_prompt_augment

        if self.config.auto_prompt_augment and self.verbose:
            Printer.print_yellow_message("🔥 Prompt Augmentation enabled!")

        self.llm = OpenAI(api_key=openai_api_key)
        llm_initialized = self.__ping_test_llm()

        self.client = OpenAI(
            api_key=perplexity_api_key, base_url="https://api.perplexity.ai"
        )
        self.initialized = self._ping_test()
        self.initialized = self.initialized and llm_initialized

        if self.verbose:
            self._print_initialization_status()

        self.sessions = {}

    def __load_environment_variables(self) -> None:
        """Load the environment variables."""
        cwd_path = os.getcwd()
        env_path = os.path.join(cwd_path, ".env")

        if self.verbose:
            Printer.print_light_grey_message(
                f"🌴 Loading environment variables from {env_path}..."
            )

        dotenv_loaded = load_dotenv(override=True, dotenv_path=env_path)

        if dotenv_loaded and self.verbose:
            Printer.print_green_message("🌴 Environment variables loaded.")
        elif not dotenv_loaded and self.verbose:
            Printer.print_red_message("🌴 Environment variables not loaded.")

    def __restructure_query(self, query: str) -> str:
        """Restructure the query so that it's fine-tuned enough to return a better quality answer.

        Args:
            query (str): The original query.

        Returns:
            str: The restructured query.
        """
        try:
            prompt = f"""
            Query: '''{query}'''
            RESPOND STRICTLY WITH THE RESTRUCTURED QUERY ONLY, NOTHING ELSE.
            """
            system = """
                You are AuxKnow, an advanced Answer Engine that provides answers to the user's questions.
                In this instance, you will be acting as a 'Query Restructurer' to fine-tune the query for better results.
            """
            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ]
            response = self.llm.chat.completions.create(
                messages=messages,
                model="gpt-4o-mini",
            )
            restructured_query = response.choices[0].message.content
            if self.verbose:
                Printer.print_yellow_message(
                    f"Restructured query: '{restructured_query}' "
                )
            return restructured_query
        except Exception as e:
            Printer.print_red_message(
                f"Error while restructuring query: {e}. Defaulting to the original query."
            )
            return query

    def __route_query_to_model(self, query: str) -> str:
        """Route the query to the appropriate model based on the query.

        Args:
            query (str): The original query.

        Returns:
            str: The model name to use for the query.
        """
        supported_models = ["sonar", "sonar-pro"]
        if self.config.enable_unibiased_reasoning:
            supported_models.append("r1-1776")
        try:
            prompt = f"""
            Query: '''{query}'''

            Determine the most suitable model for the query.

            Available models:
            1. **sonar** – Best for general queries, quick lookups, and simple factual questions.
            2. **sonar-pro** – Advanced model for complex, analytical, or research-heavy questions, providing citations.
            {"3. **r1-1776** – Uncensored, unbiased model for factual, unrestricted responses." if self.config.enable_unibiased_reasoning else ""} 

            Examples:
            - Query: "Where is Tesla headquartered?" → Response: "sonar"
            - Query: "What are the key factors affecting Tesla's Q4 revenue projections?" → Response: "sonar-pro"
            {'- Query: "Explain the geopolitical implications of BRICS expansion without censorship." → Response: "r1-1776"' if self.config.enable_unibiased_reasoning else ""}

            Strictly respond with **only** {', '.join(supported_models)}. 
            """
            system = """
                You are AuxKnow, an advanced Answer Engine that provides answers to the user's questions.
                In this instance, you will be acting as a 'Model Router' to determine which model to use for the given query.
            """

            messages = [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ]

            response = self.llm.chat.completions.create(
                messages=messages,
                model="gpt-4o-mini",
            )

            model = response.choices[0].message.content

            if model.lower() not in ["sonar", "sonar-pro", "r1-1776"]:
                Printer.print_red_message(
                    f"Invalid model name '{model}'. Please respond with either 'sonar' or 'sonar-pro' or 'r1-1776'. Defaulting to 'sonar'."
                )
                return "sonar"
            return model
        except Exception as e:
            Printer.print_red_message(
                f"Error while routing query to model: {e}. Defaulting to 'sonar'."
            )
            return "sonar"

    def __ping_test_llm(self) -> bool:
        """Perform a ping test on the LLM API.

        Returns:
            bool: True if the ping test is successful, False otherwise.
        """
        try:
            response = self.llm.chat.completions.create(
                messages=[
                    {
                        "role": "system",
                        "content": "This is a simple PING test to check the API, respond with PONG.",
                    },
                    {
                        "role": "user",
                        "content": "PING",
                    },
                ],
                model="gpt-4o-mini",
                max_tokens=200,
            )
            pong = response.choices[0].message.content
            if self.verbose:
                Printer.print_light_grey_message(f"LLM API ping test response: {pong}")
            if pong.lower().find("pong") == -1:
                Printer.print_red_message(
                    "LLM API ping test failed. Cannot use AuxKnow."
                )
                return False
            return True
        except Exception as e:
            Printer.print_red_message(
                f"LLM API ping test failed: {e}. Cannot use AuxKnow."
            )
            return False

    def _ping_test(self) -> bool:
        """Perform a ping test to check API connectivity.

        Returns:
            bool: True if the ping test is successful, False otherwise.
        """
        try:
            response = self.client.chat.completions.create(
                messages=[
                    {
                        "role": "system",
                        "content": "This is a simple PING test to check the API, respond with PONG.",
                    },
                    {
                        "role": "user",
                        "content": "PING",
                    },
                ],
                model="sonar",
                max_tokens=200,
            )
            pong = response.choices[0].message.content
            if self.verbose:
                Printer.print_light_grey_message(
                    f"Perplexity API ping test response: {pong}"
                )
            if pong.lower().find("pong") == -1:
                Printer.print_red_message(
                    "Perplexity API ping test failed. Cannot use AuxKnow."
                )
                return False
            return True
        except Exception as e:
            Printer.print_red_message(
                f"Perplexity API ping test failed: {e}. Cannot use AuxKnow."
            )
            return False

    def _print_initialization_status(self) -> None:
        """Print the initialization status."""
        if self.initialized:
            Printer.print_light_grey_message("🚀 AuxKnow ping test passed.")
            Printer.print_light_grey_message("🚀 AuxKnow API initialized successfully!")
        if self.verbose:
            Printer.print_light_grey_message("🗣️  Verbose: ON.")

    def set_config(self, config: dict) -> None:
        """Set the configuration for AuxKnow.

        Args:
            config (dict): The configuration dictionary.
        """
        answer_length_in_paragraphs = config.get("answer_length_in_paragraphs")
        auto_query_restructuring = config.get("auto_query_restructuring")
        auto_model_routing = config.get("auto_model_routing")
        lines_per_paragraph = config.get("lines_per_paragraph")
        auto_prompt_augment = config.get(
            "auto_prompt_augment", DEFAULT_AUTO_PROMPT_AUGMENT
        )
        self.config.auto_prompt_augment = auto_prompt_augment
        self.config.enable_unibiased_reasoning = config.get(
            "enable_unbiased_reasoning", DEFAULT_ENABLE_UNBIASED_REASONING
        )
        self.config.fast_mode = config.get("fast_mode", False)

        if self.config.auto_query_restructuring != auto_query_restructuring:
            self.config.auto_query_restructuring = auto_query_restructuring

        if self.config.auto_model_routing != auto_model_routing:
            self.config.auto_model_routing = auto_model_routing

        self.config.enable_unibiased_reasoning = config.get(
            "enable_unbiased_reasoning", DEFAULT_ENABLE_UNBIASED_REASONING
        )

        if self.config.answer_length_in_paragraphs > MAX_ANSWER_LENGTH_PARAGRAPHS:
            Printer.print_yellow_message(
                f"Answer length in paragraphs exceeds the maximum limit of {MAX_ANSWER_LENGTH_PARAGRAPHS}. Defaulting to {DEFAULT_ANSWER_LENGTH_PARAGRAPHS}."
            )
            self.config.answer_length_in_paragraphs = DEFAULT_ANSWER_LENGTH_PARAGRAPHS
        elif answer_length_in_paragraphs:
            self.config.answer_length_in_paragraphs = answer_length_in_paragraphs

        if self.config.lines_per_paragraph > MAX_LINES_PER_PARAGRAPH:
            Printer.print_yellow_message(
                f"Lines per paragraph exceeds the maximum limit of {MAX_LINES_PER_PARAGRAPH}. Defaulting to {DEFAULT_LINES_PER_PARAGRAPH}."
            )
            self.config.lines_per_paragraph = DEFAULT_LINES_PER_PARAGRAPH
        elif lines_per_paragraph:
            self.config.lines_per_paragraph = lines_per_paragraph

    def get_config(self) -> AuxKnowConfig:
        """Get the configuration for AuxKnow.

        Returns:
            AuxKnowConfig: The current configuration.
        """
        if not self.config:
            self.config = AuxKnowConfig()
        return deepcopy(self.config)

    def _get_prompt_augmentation_segment(self, question: str, context: str = "") -> str:
        """
        Augments the prompt and returnes the supporting prompt.

        Args:
            question (str): The question to ask.
            context (str): The context for the question.

        Returns:
            str: The supporting prompt.
        """
        try:
            user_prompt = f"""
                Your job is to provide a detailed and comprehensive supporting prompt to the given prompt. 

                The supporting prompt should be a detailed and comprehensive explanation of the given prompt. 
                It should provide a thorough and in-depth explanation of the given prompt, including its context, 
                background, and any relevant details. The supporting prompt should be written in a clear and concise 
                manner, using appropriate language and terminology to ensure clarity and understanding.

                The supporting prompt should be structured in a way that is easy to read and understand, with clear 
                headings and subheadings to organize the information. It should also be written in a way that is 
                easy to follow and understand, with clear and concise language that is easy to understand.

                The supporting prompt should be written in a way that is easy to read and understand, with clear 
                headings and subheadings to organize the information. It should also be written in a way that is 
                easy to follow and understand, with clear and concise language that is easy to understand.   

                We are giving you the prompt / question and context to provide the best, most factual and comprhensive response.

                Don't respond to the prompt / question, simply provide the supporting prompt.

                Prompt / Question: {question}
                Context: {context}
            """
            response = self.llm.chat.completions.create(
                model=DEFAULT_PROMPT_AUGMENTATION_MODEL,
                messages=[
                    {"role": "user", "content": user_prompt},
                ],
                temperature=DEFAULT_PROMPT_AUGMENTATION_TEMPERATURE,
            )
            return response.choices[0].message.content
        except:
            Printer.print_red_message(
                "Error while getting prompt augmentation segment. Returning empty string."
            )
            return ""

    def _augment_prompt(self, user_prompt: str, augmentation_segment: str) -> str:
        """
        Augments the prompt and returnes the supporting prompt.

        Args:
            user_prompt (str): The user prompt.
            augmentation_segment (str): The augmentation segment.

        Returns:
            str: The augmented prompt.
        """
        try:
            if augmentation_segment.strip() == "":
                return user_prompt
            augmented_prompt = f"""
                {user_prompt}
                {augmentation_segment}
            """
            return augmented_prompt
        except:
            Printer.print_red_message(
                "Error while augmenting prompt. Returning original prompt."
            )
            return user_prompt

    def _extract_citations_from_response(self, response: dict) -> list[str]:
        """Extract citations from the response.

        Args:
            response (dict): The response from the API.

        Returns:
            list[str]: The list of citations.
        """
        citations = []
        try:
            if hasattr(response, "citations") and response.citations:
                citations.extend(response.citations)
        except:
            pass
        return list(set(citations))  # Remove duplicates

    def _get_model(
        self, question: str, deep_research: bool, fast_mode: bool = False
    ) -> str:
        """Get the model to use for the query.

        Args:
            question (str): The question being asked
            deep_research (bool): Whether deep research mode is enabled
            fast_mode (bool): Whether fast mode is enabled (overrides other settings)

        Returns:
            str: The model name to use
        """
        if fast_mode:
            return "sonar"  # Fast mode always uses sonar
        model = DEFAULT_PERPLEXITY_MODEL
        if self.config.auto_model_routing and not deep_research:
            model = self.__route_query_to_model(question)
        elif deep_research:
            model = DEFAULT_DEEP_RESEARCH_MODEL
        return model

    def _build_user_ask_prompt(
        self,
        question: str,
        context: str = "",
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
    ) -> str:
        """
        Constructs the user prompt for asking a question.

        Args:
            question (str): The question to ask.
            context (str): The context for the question.
            deep_research (bool): Whether to enable deep research mode. (Default: False)

        Returns:
            str: The user prompt.
        """
        return f"""
                Question: {question}
                Respond in {self.config.answer_length_in_paragraphs} paragraphs with {self.config.lines_per_paragraph} lines per paragraph.
                Important: Do not include any thinking process or planning in your response.
                Provide only the final answer.
                {"Conduct a deep research like a PhD researcher and provide a detailed, factual, accurate and comprehensive response." if deep_research else ""}
                {"Context: " + context if context and context.strip() != "" else ""}
            """

    def ask(
        self,
        question: str,
        context: str = "",
        for_citations=False,
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
        fast_mode=DEFAULT_FAST_MODE_ENABLED,
    ) -> AuxKnowAnswer:
        """Ask a question to AuxKnow.

        Args:
            question (str): The question to ask.
            context (str): The context for the question.
            deep_research (bool): Whether to enable deep research mode.
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            AuxKnowAnswer: The answer.
        """
        try:
            if self.config.fast_mode:
                fast_mode = True

            if not fast_mode and self.config.auto_query_restructuring:
                question = self.__restructure_query(question)

            model = self._get_model(
                question=question, deep_research=deep_research, fast_mode=fast_mode
            )

            if self.verbose and not for_citations:
                Printer.print_yellow_message(
                    f"🧠 Asking question: '{question}' with model: '{model}'."
                )

            if not self.initialized:
                Printer.print_red_message(
                    "AuxKnow API not initialized. Cannot ask questions."
                )
                return AuxKnowAnswer(
                    answer="AuxKnow API not initialized. Cannot ask questions.",
                    citations=[],
                    is_final=True,
                )

            system_prompt = DEFAULT_AUXKNOW_SYSTEM_PROMPT

            user_prompt = self._build_user_ask_prompt(
                question, context, deep_research=deep_research
            )

            if not fast_mode and self.config.auto_prompt_augment:
                prompt_augmentation_segment = self._get_prompt_augmentation_segment(
                    question, context
                )
                user_prompt = self._augment_prompt(
                    user_prompt, prompt_augmentation_segment
                )

            messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {"role": "user", "content": user_prompt},
            ]

            response = self.client.chat.completions.create(
                messages=messages, model=model, stream=False
            )

            answer = response.choices[0].message.content
            clean_answer = re.sub(
                r"<think>.*?</think>", "", answer, flags=re.DOTALL
            ).strip()
            clean_answer = re.sub(r"\n{3,}", "\n\n", clean_answer)

            citations = self._extract_citations_from_response(response)

            if len(citations) == 0:
                citations, _ = self.get_citations(question, clean_answer)

            return AuxKnowAnswer(
                answer=clean_answer,
                citations=citations,
                is_final=True,
            )
        except Exception as e:
            Printer.print_red_message(f"Error while asking question: {e}.")
            return AuxKnowAnswer(
                answer="Sorry, can't provide an answer right now. Please try again later!",
                citations=[],
                is_final=True,
            )

    def ask_stream(
        self,
        question: str,
        context: str = "",
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
        fast_mode=DEFAULT_FAST_MODE_ENABLED,
    ) -> Generator[AuxKnowAnswer, None, None]:
        """Ask a question to AuxKnow with streaming response.

        Args:
            question (str): The question to ask.
            context (str): The context for the question.
            deep_research (bool): Whether to enable deep research mode.
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            Generator[AuxKnowAnswer, None, None]: A generator of answers.
        """
        try:
            if not fast_mode and self.config.auto_query_restructuring:
                question = self.__restructure_query(question)

            model = self._get_model(
                question=question, deep_research=deep_research, fast_mode=fast_mode
            )

            if self.verbose:
                Printer.print_yellow_message(
                    f"🧠 Asking question: '{question}' with model: '{model}'."
                )

            if not self.initialized:
                Printer.print_red_message(
                    "AuxKnow API not initialized. Cannot ask questions."
                )
                yield AuxKnowAnswer(
                    answer="AuxKnow API not initialized. Cannot ask questions.",
                    citations=[],
                    is_final=True,
                )
                return

            system_prompt = DEFAULT_AUXKNOW_SYSTEM_PROMPT

            user_prompt = self._build_user_ask_prompt(
                question, context, deep_research=deep_research
            )

            if not fast_mode and self.config.auto_prompt_augment:
                prompt_augmentation_segment = self._get_prompt_augmentation_segment(
                    question, context
                )
                user_prompt = self._augment_prompt(
                    user_prompt, prompt_augmentation_segment
                )

            messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {"role": "user", "content": user_prompt},
            ]

            response_stream = self.client.chat.completions.create(
                messages=messages, model=model, stream=True
            )

            full_answer = ""
            citations = []
            buffer = ""
            is_in_think_block = False

            for response in response_stream:
                chunk: str = response.choices[0].delta.content
                if not chunk:
                    continue

                buffer += chunk

                while True:
                    if is_in_think_block:
                        end_idx = buffer.find("</think>")
                        if end_idx == -1:
                            break
                        buffer = buffer[end_idx + 8 :]  # Remove think block
                        is_in_think_block = False
                    else:
                        start_idx = buffer.find("<think>")
                        if start_idx == -1:
                            if buffer:
                                new_citations = self._extract_citations_from_response(
                                    response
                                )
                                if new_citations:
                                    citations.extend(new_citations)
                                    citations = list(set(citations))

                                full_answer += buffer
                                yield AuxKnowAnswer(
                                    answer=buffer, citations=citations, is_final=False
                                )
                            buffer = ""
                            break

                        if start_idx > 0:
                            pre_think = buffer[:start_idx]
                            full_answer += pre_think
                            yield AuxKnowAnswer(
                                answer=pre_think, citations=citations, is_final=False
                            )
                        buffer = buffer[start_idx + 7 :]
                        is_in_think_block = True

            if buffer and not is_in_think_block:
                full_answer += buffer
                yield AuxKnowAnswer(answer=buffer, citations=citations, is_final=False)

            if len(citations) == 0:
                citations, _ = self.get_citations(question, full_answer)

            yield AuxKnowAnswer(
                answer=full_answer,
                citations=citations,
                is_final=True,
            )
        except Exception as e:
            Printer.print_red_message(f"Error while asking question: {e}.")
            yield AuxKnowAnswer(
                answer="Sorry, can't provide an answer right now. Please try again later!",
                citations=[],
                is_final=True,
            )

    def _ask_with_context_stream(
        self,
        session: AuxKnowSession,
        question: str,
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
        fast_mode=DEFAULT_FAST_MODE_ENABLED,
    ) -> Generator[AuxKnowAnswer, None, None]:
        """Ask a question within a session to maintain context with streaming.

        Args:
            session (AuxKnowSession): The session in which to ask the question.
            question (str): The question to ask.
            deep_research (bool): Whether deep research mode is enabled.
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            Generator[AuxKnowAnswer, None, None]: A generator of answers.
        """
        context_string = self._build_context_string(session.context)
        answer_stream = self.ask_stream(
            question, context_string, deep_research=deep_research, fast_mode=fast_mode
        )
        for partial_answer in answer_stream:
            session.context.append(
                {"question": question, "answer": partial_answer.answer}
            )
            yield partial_answer

    def _ask_with_context(
        self,
        session: AuxKnowSession,
        question: str,
        deep_research=DEFAULT_DEEP_RESEARCH_ENABLED,
        fast_mode=DEFAULT_FAST_MODE_ENABLED,
    ) -> AuxKnowAnswer:
        """Ask a question within a session to maintain context.

        Args:
            session (AuxKnowSession): The session in which to ask the question.
        Args:
            session (AuxKnowSession): The session in which to ask the question.
            question (str): The question to ask.
            deep_research (bool): Whether deep research mode is enabled.
            fast_mode (bool): When True, overrides other settings for fastest response.

        Returns:
            AuxKnowAnswer: The answers object.
        """
        context_string = self._build_context_string(session.context)
        return self.ask(
            question=question,
            context=context_string,
            deep_research=deep_research,
            fast_mode=fast_mode,
        )

    def create_session(self) -> AuxKnowSession:
        """Create a new session and return the session object.

        Returns:
            AuxKnowSession: The created session.
        """
        session_id = str(uuid4())
        session = AuxKnowSession(session_id=session_id, auxknow=self)
        self.sessions[session_id] = session
        return session

    def _close_session(self, session: AuxKnowSession) -> None:
        """Mark the session as closed.

        Args:
            session (AuxKnowSession): The session to close.
        """
        session.closed = True
        if session.session_id in self.sessions:
            del self.sessions[session.session_id]

    def _build_context_string(self, context: list[dict[str, str]]) -> str:
        """Build the context string from the list of question-answer pairs.

        Args:
            context (list[dict[str, str]]): The list of question-answer pairs.

        Returns:
            str: The context string.
        """
        context_string = ""
        for qa in context[-10:]:  # Take the last 10 question-answer pairs
            qa_string = f"Q: {qa['question']}\nA: {qa['answer']}\n"
            if len((context_string + qa_string).split()) <= MAX_CONTEXT_TOKENS:
                context_string += qa_string
            else:
                break
        return context_string

    def get_citations(
        self, query: str, query_response: str
    ) -> tuple[Union[list[str], None], str]:
        """
        Gets the citations for the given query and response.

        Args:
            llm (OpenAI): The Large Language model to generate metadata for citation generation.
            query (str): The query to search for.
            query_response (str): The response to the query.

        Returns:
            list[str]: The citations which is a list of URLs.
        """
        try:
            question = f"""
                Can you please generate a detailed list of citations for the given query and response?

                Query: '''{query}'''
                Response: '''{query_response}'''

            """
            response = self.ask(question, for_citations=True)
            return response.citations, ""
        except Exception as e:
            if self.verbose:
                Printer.print_red_message(f"Error while getting citations: {e}")
            return [], str(e)
