import base64
import json
import logging
import mimetypes
import os
import re
import uuid
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from io import BytesIO
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union

import httpx
import PIL.Image
import pyjson5
import requests
import yaml
from anthropic import APIError as AnthropicAPIError
from anthropic import RateLimitError as AnthropicRateLimitError
from google.genai import errors as gerrors
from langchain_core.messages.chat import ChatMessage
from pydantic import BaseModel, ConfigDict, ValidationError

from khoj.database.adapters import ConversationAdapters
from khoj.database.models import (
    ChatMessageModel,
    ClientApplication,
    Intent,
    KhojUser,
    UserMemory,
)
from khoj.processor.conversation import prompts
from khoj.search_filter.base_filter import BaseFilter
from khoj.search_filter.date_filter import DateFilter
from khoj.search_filter.file_filter import FileFilter
from khoj.search_filter.word_filter import WordFilter
from khoj.utils.helpers import (
    ConversationCommand,
    count_tokens,
    get_encoder,
    is_none_or_empty,
    is_promptrace_enabled,
    merge_dicts,
)
from khoj.utils.rawconfig import FileAttachment
from khoj.utils.yaml import yaml_dump

logger = logging.getLogger(__name__)

try:
    import importlib.util

    git_spec = importlib.util.find_spec("git")
    if git_spec is None:
        raise ImportError
except ImportError:
    if is_promptrace_enabled():
        logger.warning("GitPython not installed. `pip install gitpython` to use prompt tracer.")

model_to_prompt_size = {
    # OpenAI Models
    "gpt-4o": 60000,
    "gpt-4o-mini": 60000,
    "gpt-4.1": 60000,
    "gpt-4.1-mini": 120000,
    "gpt-4.1-nano": 120000,
    "o1-mini": 90000,
    "o1": 30000,
    "o3-mini": 90000,
    "o3": 60000,
    "o3-pro": 30000,
    "o4-mini": 90000,
    "gpt-5-2025-08-07": 120000,
    "gpt-5-mini-2025-08-07": 120000,
    "gpt-5-nano-2025-08-07": 120000,
    # Google Models
    "gemini-3-pro-preview": 120000,
    "gemini-2.5-flash": 120000,
    "gemini-2.5-flash-lite": 120000,
    "gemini-2.5-pro": 60000,
    "gemini-2.0-flash": 120000,
    "gemini-2.0-flash-lite": 120000,
    "gemini-1.5-flash": 120000,
    "gemini-1.5-pro": 60000,
    # Anthropic Models
    "claude-3-5-sonnet-20241022": 60000,
    "claude-3-5-sonnet-latest": 60000,
    "claude-3-7-sonnet-20250219": 60000,
    "claude-3-7-sonnet-latest": 60000,
    "claude-3-5-haiku-20241022": 60000,
    "claude-haiku-4-5-20251001": 60000,
    "claude-sonnet-4-0": 60000,
    "claude-sonnet-4-20250514": 60000,
    "claude-opus-4-0": 60000,
    "claude-opus-4-20250514": 60000,
}
model_to_tokenizer: Dict[str, str] = {}


class RetryableModelError(Exception):
    """
    Exception raised when a chat model fails with a retryable error.
    This is used to trigger fallback to the next model in the priority list.

    Wraps provider-specific retryable errors like:
    - OpenAI: RateLimitError, InternalServerError, APITimeoutError
    - Anthropic: RateLimitError, APIError
    - Google/Gemini: API errors with codes 429, 502, 503, 504
    """

    def __init__(self, message: str, original_exception: Exception = None, model_name: str = None):
        super().__init__(message)
        self.original_exception = original_exception
        self.model_name = model_name

    def __str__(self):
        model_info = f" (model: {self.model_name})" if self.model_name else ""
        return f"{super().__str__()}{model_info}"


def is_retryable_exception(exception: BaseException) -> bool:
    """
    Check if an exception is retryable and should trigger fallback to another model.
    """
    # OpenAI exceptions
    if hasattr(exception, "__module__") and exception.__module__ and "openai" in exception.__module__:
        import openai

        if isinstance(
            exception,
            (
                openai._exceptions.APITimeoutError,
                openai._exceptions.RateLimitError,
                openai._exceptions.InternalServerError,
            ),
        ):
            return True

    # Anthropic exceptions
    if isinstance(exception, (AnthropicRateLimitError, AnthropicAPIError)):
        return True

    # Google/Gemini exceptions
    if isinstance(exception, (gerrors.APIError, gerrors.ClientError)):
        # Check for specific error codes that are retryable
        if hasattr(exception, "code") and exception.code in [429, 500, 502, 503, 504]:
            return True

    # Network errors
    if isinstance(exception, (httpx.TimeoutException, httpx.NetworkError)):
        return True

    # Empty or no response by model over API results in ValueError
    if isinstance(exception, ValueError):
        return True

    return False


class AgentMessage(BaseModel):
    role: Literal["user", "assistant", "system", "environment"]
    content: Union[str, List]


class OperatorRun:
    def __init__(
        self,
        query: str,
        trajectory: list[AgentMessage] | list[dict] = None,
        response: str = None,
        webpages: list[dict] = None,
    ):
        self.query = query
        self.response = response
        self.webpages = webpages or []
        self.trajectory: list[AgentMessage] = []
        if trajectory:
            for item in trajectory:
                if isinstance(item, dict):
                    self.trajectory.append(AgentMessage(**item))
                elif hasattr(item, "role") and hasattr(item, "content"):  # Heuristic for AgentMessage like object
                    self.trajectory.append(item)
                else:
                    logger.warning(f"Unexpected item type in trajectory: {type(item)}")

    def to_dict(self) -> dict:
        # Ensure AgentMessage instances in trajectory are also dicts
        serialized_trajectory = []
        for msg in self.trajectory:
            if hasattr(msg, "model_dump"):  # Check if it's a Pydantic model
                serialized_trajectory.append(msg.model_dump())
            elif isinstance(msg, dict):
                serialized_trajectory.append(msg)  # Already a dict
        return {
            "query": self.query,
            "response": self.response,
            "trajectory": serialized_trajectory,
            "webpages": self.webpages,
        }


class ToolCall:
    def __init__(self, name: str, args: dict, id: str):
        self.name = name
        self.args = args
        self.id = id


class ResearchIteration:
    def __init__(
        self,
        query: ToolCall | dict | str,
        context: list = None,
        onlineContext: dict = None,
        codeContext: dict = None,
        operatorContext: dict | OperatorRun = None,
        summarizedResult: str = None,
        warning: str = None,
        raw_response: list = None,
    ):
        self.query = ToolCall(**query) if isinstance(query, dict) else query
        self.context = context
        self.onlineContext = onlineContext
        self.codeContext = codeContext
        self.operatorContext = OperatorRun(**operatorContext) if isinstance(operatorContext, dict) else operatorContext
        self.summarizedResult = summarizedResult
        self.warning = warning
        self.raw_response = raw_response

    def to_dict(self) -> dict:
        data = vars(self).copy()
        data["query"] = self.query.__dict__ if isinstance(self.query, ToolCall) else self.query
        data["operatorContext"] = self.operatorContext.to_dict() if self.operatorContext else None
        return data


def construct_iteration_history(
    previous_iterations: List[ResearchIteration],
    query: str = None,
    query_images: List[str] = None,
    query_files: str = None,
) -> list[ChatMessageModel]:
    iteration_history: list[ChatMessageModel] = []
    query_message_content = construct_structured_message(query, query_images, attached_file_context=query_files)
    if query_message_content:
        iteration_history.append(ChatMessageModel(by="you", message=query_message_content))

    # Group iterations: parallel tool calls share the same raw_response (only first has it)
    # We need to group them so one assistant message has all tool_use blocks and
    # one user message has all tool_results
    current_group_raw_response = None
    current_group_tool_results = []

    def flush_group():
        """Output the current group as assistant message + user message with tool results"""
        nonlocal current_group_raw_response, current_group_tool_results
        if current_group_raw_response and current_group_tool_results:
            iteration_history.append(
                ChatMessageModel(
                    by="khoj",
                    message=current_group_raw_response,
                    intent=Intent(type="tool_call", query=query),
                )
            )
            iteration_history.append(
                ChatMessageModel(
                    by="you",
                    intent=Intent(type="tool_result"),
                    message=current_group_tool_results,
                )
            )
        current_group_raw_response = None
        current_group_tool_results = []

    for iteration in previous_iterations:
        if not iteration.query or isinstance(iteration.query, str):
            # Flush any pending group before adding non-tool message
            flush_group()
            iteration_history.append(
                ChatMessageModel(
                    by="you",
                    message=iteration.summarizedResult
                    or iteration.warning
                    or "Please specify what you want to do next.",
                )
            )
            continue

        # If this iteration has raw_response, it starts a new group of parallel tool calls
        if iteration.raw_response:
            # Flush previous group if exists
            flush_group()
            current_group_raw_response = iteration.raw_response

        # If no raw_response and no current group, create a fallback single-tool response
        elif not current_group_raw_response:
            current_group_raw_response = [
                {
                    "type": "tool_use",
                    "id": iteration.query.id,
                    "name": iteration.query.name,
                    "input": iteration.query.args,
                }
            ]

        # Add tool result to current group
        current_group_tool_results.append(
            {
                "type": "tool_result",
                "id": iteration.query.id,
                "name": iteration.query.name,
                "content": iteration.summarizedResult,
            }
        )

    # Flush any remaining group
    flush_group()

    return iteration_history


def construct_chat_history(chat_history: list[ChatMessageModel], n: int = 4, agent_name="AI") -> str:
    chat_history_str = ""
    for chat in chat_history[-n:]:
        intent_type = chat.intent.type if chat.intent and chat.intent.type else ""
        inferred_queries = chat.intent.inferred_queries if chat.intent else None
        if chat.by == "khoj" and intent_type in ["remember", "reminder", "summarize"]:
            if inferred_queries:
                chat_history_str += f'{agent_name}: {{"queries": {inferred_queries}}}\n'
            chat_history_str += f"{agent_name}: {chat.message}\n\n"
        elif chat.by == "khoj" and chat.images:
            chat_history_str += f"{agent_name}: [generated image redacted for space]\n"
        elif chat.by == "khoj" and ("excalidraw" in intent_type):
            chat_history_str += f"{agent_name}: {inferred_queries[0]}\n"
        elif chat.by == "you":
            chat_history_str += f"User: {chat.message}\n"
            raw_query_files = chat.queryFiles
            if raw_query_files:
                query_files: Dict[str, str] = {}
                for file in raw_query_files:
                    query_files[file["name"]] = file["content"]

                query_file_context = gather_raw_query_files(query_files)
                chat_history_str += f"User: {query_file_context}\n"

    return chat_history_str


def construct_question_history(
    conversation_log: list[ChatMessageModel],
    include_query: bool = True,
    lookback: int = 6,
    query_prefix: str = "Q",
    agent_name: str = "Khoj",
) -> str:
    """
    Constructs a chat history string formatted for query extraction purposes.
    """
    history_parts = ""
    original_query = None
    for chat in conversation_log[-lookback:]:
        if chat.by == "you":
            original_query = json.dumps(chat.message)
            history_parts += f"{query_prefix}: {original_query}\n"
        if chat.by == "khoj":
            if original_query is None:
                continue

            message = chat.message
            inferred_queries_list = chat.intent.inferred_queries or [] if chat.intent else []

            # Ensure inferred_queries_list is a list, defaulting to the original query in a list
            if not inferred_queries_list:
                inferred_queries_list = [original_query]
            # If it's a string (though unlikely based on usage), wrap it in a list
            elif isinstance(inferred_queries_list, str):
                inferred_queries_list = [inferred_queries_list]

            if include_query:
                # Ensure 'type' exists and is a string before checking 'to-image'
                intent_type = chat.intent.type if chat.intent and chat.intent.type else ""
                if "to-image" not in intent_type:
                    history_parts += f'{agent_name}: {{"queries": {inferred_queries_list}}}\n'
                    history_parts += f"A: {message}\n\n"
            else:
                history_parts += f"{agent_name}: {message}\n\n"

            # Reset original_query for the next turn
            original_query = None

    return history_parts


def construct_chat_history_for_operator(conversation_history: List[ChatMessageModel], n: int = 6) -> list[AgentMessage]:
    """
    Construct chat history for operator agent in conversation log.
    Only include last n completed turns (i.e with user and khoj message).
    """
    chat_history: list[AgentMessage] = []
    user_message: Optional[AgentMessage] = None

    for chat in conversation_history:
        if len(chat_history) >= n:
            break
        if chat.by == "you" and chat.message:
            content = [{"type": "text", "text": chat.message}]
            for file in chat.queryFiles or []:
                content += [{"type": "text", "text": f"## File: {file['name']}\n\n{file['content']}"}]
            user_message = AgentMessage(role="user", content=content)
        elif chat.by == "khoj" and chat.message:
            chat_history += [user_message, AgentMessage(role="assistant", content=chat.message)]
    return chat_history


def construct_tool_chat_history(
    previous_iterations: List[ResearchIteration], tool: ConversationCommand = None
) -> List[ChatMessageModel]:
    """
    Construct chat history from previous iterations for a specific tool

    If a tool is provided, only the inferred queries for that tool is added.
    If no tool is provided inferred query for all tools used are added.
    """
    chat_history: list = []

    def base_extractor(iteration: ResearchIteration) -> List[str]:
        return []

    extract_inferred_query_map: Dict[ConversationCommand, Callable[[ResearchIteration], List[str]]] = {
        ConversationCommand.SemanticSearchFiles: (
            lambda iteration: [c["query"] for c in iteration.context] if iteration.context else []
        ),
        ConversationCommand.SearchWeb: (
            lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
        ),
        ConversationCommand.ReadWebpage: (
            lambda iteration: list(iteration.onlineContext.keys()) if iteration.onlineContext else []
        ),
        ConversationCommand.PythonCoder: (
            lambda iteration: list(iteration.codeContext.keys()) if iteration.codeContext else []
        ),
    }
    for iteration in previous_iterations:
        if not iteration.query or isinstance(iteration.query, str):
            chat_history.append(
                ChatMessageModel(
                    by="you",
                    message=iteration.summarizedResult
                    or iteration.warning
                    or "Please specify what you want to do next.",
                )
            )
            continue

        # If a tool is provided use the inferred query extractor for that tool if available
        # If no tool is provided, use inferred query extractor for the tool used in the iteration
        # Fallback to base extractor if the tool does not have an inferred query extractor
        inferred_query_extractor = extract_inferred_query_map.get(
            tool or ConversationCommand(iteration.query.name), base_extractor
        )
        chat_history += [
            ChatMessageModel(
                by="you",
                message=yaml.dump(iteration.query.args, default_flow_style=False),
            ),
            ChatMessageModel(
                by="khoj",
                intent=Intent(
                    type="remember",
                    query=yaml.dump(iteration.query.args, default_flow_style=False),
                    inferred_queries=inferred_query_extractor(iteration),
                    memory_type="notes",
                ),
                message=iteration.summarizedResult,
            ),
        ]

    return chat_history


class ChatEvent(Enum):
    START_LLM_RESPONSE = "start_llm_response"
    END_LLM_RESPONSE = "end_llm_response"
    MESSAGE = "message"
    REFERENCES = "references"
    GENERATED_ASSETS = "generated_assets"
    STATUS = "status"
    THOUGHT = "thought"
    METADATA = "metadata"
    USAGE = "usage"
    END_RESPONSE = "end_response"
    INTERRUPT = "interrupt"
    END_EVENT = "␃🔚␗"


def message_to_log(
    user_message,
    chat_response,
    user_message_metadata={},
    khoj_message_metadata={},
    chat_history: List[ChatMessageModel] = [],
) -> List[ChatMessageModel]:
    """Create json logs from messages, metadata for conversation log"""
    default_khoj_message_metadata = {
        "intent": {"type": "remember", "memory-type": "notes", "query": user_message},
    }
    khoj_response_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # Filter out any fields that are set to None
    user_message_metadata = {k: v for k, v in user_message_metadata.items() if v is not None}
    khoj_message_metadata = {k: v for k, v in khoj_message_metadata.items() if v is not None}

    # Create json log from Human's message
    human_log = merge_dicts({"message": user_message, "by": "you"}, user_message_metadata)

    # Create json log from GPT's response
    khoj_log = merge_dicts(khoj_message_metadata, default_khoj_message_metadata)
    khoj_log = merge_dicts({"message": chat_response, "by": "khoj", "created": khoj_response_time}, khoj_log)

    # Validate message logs
    # Only validates top-level fields, not nested fields, defined in ChatMessageModel
    class StrictChatMessageModel(ChatMessageModel):
        model_config = ConfigDict(extra="forbid", strict=True)

    try:
        StrictChatMessageModel(**human_log)
    except ValidationError as e:
        logger.error(f"Validation error in user chat message: {e}\nUser Message: {human_log}\n")
    try:
        StrictChatMessageModel(**khoj_log)
    except ValidationError as e:
        logger.error(f"Validation error in khoj chat message: {e}\nKhoj Message: {khoj_log}\n")

    human_message = ChatMessageModel(**human_log)
    khoj_message = ChatMessageModel(**khoj_log)
    chat_history.extend([human_message, khoj_message])
    return chat_history


async def save_to_conversation_log(
    q: str,
    chat_response: str,
    user: KhojUser,
    user_message_time: str = None,
    compiled_references: List[Dict[str, Any]] = [],
    online_results: Dict[str, Any] = {},
    code_results: Dict[str, Any] = {},
    operator_results: List[OperatorRun] = None,
    inferred_queries: List[str] = [],
    intent_type: str = "remember",
    client_application: ClientApplication = None,
    conversation_id: str = None,
    automation_id: str = None,
    relevant_memories: List[UserMemory] = [],
    query_images: List[str] = None,
    raw_query_files: List[FileAttachment] = [],
    generated_images: List[str] = [],
    generated_mermaidjs_diagram: str = None,
    research_results: Optional[List[ResearchIteration]] = None,
    train_of_thought: List[Any] = [],
    tracer: Dict[str, Any] = {},
):
    from khoj.routers.helpers import ai_update_memories

    user_message_time = user_message_time or datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    turn_id = tracer.get("mid") or str(uuid.uuid4())

    user_message_metadata = {"created": user_message_time, "images": query_images, "turnId": turn_id}

    if raw_query_files and len(raw_query_files) > 0:
        user_message_metadata["queryFiles"] = [file.model_dump(mode="json") for file in raw_query_files]

    khoj_message_metadata = {
        "context": compiled_references,
        "intent": {"inferred-queries": inferred_queries, "type": intent_type},
        "onlineContext": online_results,
        "codeContext": code_results,
        "operatorContext": [o.to_dict() for o in operator_results] if operator_results and not chat_response else None,
        "researchContext": [r.to_dict() for r in research_results] if research_results and not chat_response else None,
        "automationId": automation_id,
        "trainOfThought": train_of_thought,
        "turnId": turn_id,
        "images": generated_images,
    }

    if generated_mermaidjs_diagram:
        khoj_message_metadata["mermaidjsDiagram"] = generated_mermaidjs_diagram

    try:
        new_messages = message_to_log(
            user_message=q,
            chat_response=chat_response,
            user_message_metadata=user_message_metadata,
            khoj_message_metadata=khoj_message_metadata,
            chat_history=[],
        )
    except ValidationError as e:
        new_messages = None
        logger.error(f"Error constructing chat history: {e}")

    db_conversation = None
    if new_messages:
        db_conversation = await ConversationAdapters.save_conversation(
            user,
            new_messages,
            client_application=client_application,
            conversation_id=conversation_id,
            user_message=q,
        )

    if not automation_id:
        # Don't update memories from automations, as this could get noisy.
        await ai_update_memories(
            user=user,
            conversation_history=new_messages or [],
            memories=relevant_memories,
            agent=db_conversation.agent if db_conversation else None,
            tracer=tracer,
        )

    if is_promptrace_enabled():
        merge_message_into_conversation_trace(q, chat_response, tracer)

    logger.info(
        f"""
Saved Conversation Turn ({db_conversation.id if db_conversation else "N/A"}):
You ({user.username}): "{q}"

Khoj: "{chat_response}"
""".strip()
    )


def construct_structured_message(
    message: list[dict] | str,
    images: list[str] = None,
    model_type: str = None,
    vision_enabled: bool = True,
    attached_file_context: str = None,
):
    """
    Format messages into appropriate multimedia format for supported chat model types.

    Assume vision is enabled and chat model provider supports messages in chatml format, unless specified otherwise.
    """
    constructed_messages: List[dict[str, Any]] = []
    if not is_none_or_empty(message):
        constructed_messages += [{"type": "text", "text": message}] if isinstance(message, str) else message
    # Drop image message passed by caller if chat model does not have vision enabled
    if not vision_enabled:
        constructed_messages = [m for m in constructed_messages if m.get("type") != "image_url"]
    if not is_none_or_empty(attached_file_context):
        constructed_messages += [{"type": "text", "text": attached_file_context}]
    if vision_enabled and images:
        for image in images:
            constructed_messages += [{"type": "image_url", "image_url": {"url": image}}]
    return constructed_messages


def gather_raw_query_files(
    query_files: Dict[str, str],
):
    """
    Gather contextual data from the given (raw) files
    """

    if len(query_files) == 0:
        return ""

    contextual_data = " ".join(
        [f"File: {file_name}\n\n{file_content}\n\n" for file_name, file_content in query_files.items()]
    )
    return f"I have attached the following files:\n\n{contextual_data}"


def generate_chatml_messages_with_context(
    # Context
    user_message: str,
    query_files: str = None,
    query_images=None,
    context_message="",
    relevant_memories: List[UserMemory] = None,
    generated_asset_results: Dict[str, Dict] = {},
    program_execution_context: List[str] = [],
    chat_history: list[ChatMessageModel] = [],
    system_message: str = None,
    # Model Config
    model_name="gpt-4o-mini",
    model_type="",
    max_prompt_size=None,
    tokenizer_name=None,
    vision_enabled=False,
):
    """Generate chat messages with appropriate context from previous conversation to send to the chat model"""
    # Set max prompt size from user config or based on pre-configured for model and machine specs
    if not max_prompt_size:
        max_prompt_size = model_to_prompt_size.get(model_name, 10000)

    # Scale lookback turns proportional to max prompt size supported by model
    lookback_turns = max_prompt_size // 750

    # Extract Chat History for Context
    chatml_messages: List[ChatMessage] = []
    for chat in chat_history:
        message_context = []
        message_attached_files = ""

        generated_assets = {}

        chat_message = chat.message
        role = "user" if chat.by == "you" else "assistant"

        # Legacy code to handle excalidraw diagrams prior to Dec 2024
        if chat.by == "khoj" and chat.intent and "excalidraw" in chat.intent.type:
            chat_message = (chat.intent.inferred_queries or [])[0]

        # Add search and action context
        if not is_none_or_empty(chat.onlineContext):
            message_context += [
                {
                    "type": "text",
                    "text": f"{prompts.online_search_conversation.format(online_results=chat.onlineContext)}",
                }
            ]

        if not is_none_or_empty(chat.codeContext):
            message_context += [
                {
                    "type": "text",
                    "text": f"{prompts.code_executed_context.format(code_results=chat.codeContext)}",
                }
            ]

        if not is_none_or_empty(chat.operatorContext):
            operator_context = chat.operatorContext
            operator_content = "\n\n".join([f"## Task: {oc['query']}\n{oc['response']}\n" for oc in operator_context])
            message_context += [
                {
                    "type": "text",
                    "text": f"{prompts.operator_execution_context.format(operator_results=operator_content)}",
                }
            ]

        if not is_none_or_empty(chat.context):
            references = "\n\n".join(
                {
                    f"# URI: {item.uri or item.file}\n## {item.compiled}\n"
                    for item in chat.context or []
                    if isinstance(item, dict)
                }
            )
            message_context += [{"type": "text", "text": f"{prompts.notes_conversation.format(references=references)}"}]

        if not is_none_or_empty(message_context):
            reconstructed_context_message = ChatMessage(content=message_context, role="user")
            chatml_messages.append(reconstructed_context_message)

        # Add generated assets
        if not is_none_or_empty(chat.images) and role == "assistant":
            generated_assets["image"] = {
                "description": (chat.intent.inferred_queries or [user_message])[0],
            }

        if not is_none_or_empty(chat.mermaidjsDiagram) and role == "assistant":
            generated_assets["diagram"] = {
                "query": (chat.intent.inferred_queries or [user_message])[0],
            }

        if not is_none_or_empty(generated_assets):
            chatml_messages.append(
                ChatMessage(
                    content=f"{prompts.generated_assets_context.format(generated_assets=yaml_dump(generated_assets))}\n",
                    role="user",
                )
            )

        # Add user query with attached file, images or khoj response
        if chat.queryFiles:
            raw_query_files = chat.queryFiles
            query_files_dict = dict()
            for file in raw_query_files:
                query_files_dict[file["name"]] = file["content"]

            message_attached_files = gather_raw_query_files(query_files_dict)

        message_content = construct_structured_message(
            chat_message, chat.images if role == "user" else [], model_type, vision_enabled, message_attached_files
        )

        reconstructed_message = ChatMessage(
            content=message_content,
            role=role,
            additional_kwargs={"message_type": chat.intent.type if chat.intent else None},
        )
        chatml_messages.append(reconstructed_message)

        if len(chatml_messages) >= 3 * lookback_turns:
            break

    messages: list[ChatMessage] = []

    if not is_none_or_empty(system_message):
        messages.append(ChatMessage(content=system_message, role="system"))

    if len(chatml_messages) > 0:
        messages += chatml_messages

    if program_execution_context:
        program_context_text = "\n".join(program_execution_context)
        context_message += f"{prompts.additional_program_context.format(context=program_context_text)}\n"

    if not is_none_or_empty(context_message):
        messages.append(ChatMessage(content=context_message, role="user"))

    if not is_none_or_empty(generated_asset_results):
        messages.append(
            ChatMessage(
                content=prompts.generated_assets_context.format(generated_assets=yaml_dump(generated_asset_results)),
                role="user",
            ),
        )

    if not is_none_or_empty(relevant_memories):
        memory_context = "Your memory system retrieved the following memories about me based on our previous conversations. Ignore them if they are not relevant to the query.\n<retrieved_memories>\n"
        for memory in relevant_memories:
            friendly_dt = memory.created_at.strftime("%Y-%m-%d %H:%M:%S")
            memory_context += f"- [{friendly_dt}]: {memory.raw}\n"
        memory_context += "</retrieved_memories>"
        messages.append(ChatMessage(content=memory_context, role="user"))

    if not is_none_or_empty(user_message):
        messages.append(
            ChatMessage(
                content=construct_structured_message(
                    user_message, query_images, model_type, vision_enabled, query_files
                ),
                role="user",
            )
        )

    # Normalize message content to list of chatml dictionaries
    for message in messages:
        if isinstance(message.content, str):
            message.content = [{"type": "text", "text": message.content}]

    # Truncate oldest messages from conversation history until under max supported prompt size by model
    messages = truncate_messages(messages, max_prompt_size, model_name, tokenizer_name)

    # Return messages in chronological order
    return messages


def count_total_tokens(
    messages: list[ChatMessage], encoder, system_message: Optional[list[ChatMessage]] = None
) -> Tuple[int, int]:
    """Count total tokens in messages including system message"""
    system_message_tokens = (
        sum([count_tokens(message.content, encoder) for message in system_message]) if system_message else 0
    )
    message_tokens = sum([count_tokens(message.content, encoder) for message in messages])
    # Reserves 4 tokens to demarcate each message (e.g <|im_start|>user, <|im_end|>, <|endoftext|> etc.)
    total_tokens = message_tokens + system_message_tokens + 4 * len(messages)
    return total_tokens, system_message_tokens


def truncate_messages(
    messages: list[ChatMessage],
    max_prompt_size: int,
    model_name: str,
    tokenizer_name=None,
) -> list[ChatMessage]:
    """Truncate messages to fit within max prompt size supported by model"""
    encoder = get_encoder(model_name, tokenizer_name)

    # Extract system message from messages
    system_message = []
    non_system_messages = []
    for message in messages:
        if message.role == "system":
            system_message.append(message)
        else:
            non_system_messages.append(message)

    # New message list without system messages
    messages = non_system_messages

    # Drop older messages until under max supported prompt size by model
    total_tokens, system_message_tokens = count_total_tokens(messages, encoder, system_message)

    while total_tokens > max_prompt_size and (len(messages) > 1 or len(messages[0].content) > 1):
        # If the last message has more than one content part, pop the oldest content part.
        # For tool calls, the whole message should dropped, assistant's tool call content being truncated annoys AI APIs.
        if len(messages[0].content) > 1 and messages[0].additional_kwargs.get("message_type") != "tool_call":
            # The oldest content part is earlier in content list. So pop from the front.
            messages[0].content.pop(0)
        # Otherwise, pop the last message if it has only one content part or is a tool call.
        else:
            # The oldest message is the last one. So pop from the back.
            dropped_message = messages.pop(0)
            # Drop tool result pair of tool call, if tool call message has been removed
            if (
                dropped_message.additional_kwargs.get("message_type") == "tool_call"
                and messages
                and messages[0].additional_kwargs.get("message_type") == "tool_result"
            ):
                messages.pop(0)

        total_tokens, _ = count_total_tokens(messages, encoder, system_message)

    # Truncate current message if still over max supported prompt size by model
    total_tokens, _ = count_total_tokens(messages, encoder, system_message)
    if total_tokens > max_prompt_size:
        # At this point, a single message with a single content part of type dict should remain
        assert len(messages) == 1 and len(messages[0].content) == 1 and isinstance(messages[0].content[0], dict), (
            "Expected a single message with a single content part remaining at this point in truncation"
        )

        # Collate message content into single string to ease truncation
        part = messages[0].content[0]
        message_content: str = part["text"] if part["type"] == "text" else json.dumps(part)
        message_role = messages[0].role

        remaining_context = "\n".join(message_content.split("\n")[:-1])
        original_question = "\n" + "\n".join(message_content.split("\n")[-1:])

        original_question_tokens = count_tokens(original_question, encoder)
        remaining_tokens = max_prompt_size - system_message_tokens
        if remaining_tokens > original_question_tokens:
            remaining_tokens -= original_question_tokens
            truncated_context = encoder.decode(encoder.encode(remaining_context)[:remaining_tokens]).strip()
            truncated_content = truncated_context + original_question
        else:
            truncated_content = encoder.decode(encoder.encode(original_question)[:remaining_tokens]).strip()
        messages = [ChatMessage(content=[{"type": "text", "text": truncated_content}], role=message_role)]

        truncated_snippet = (
            f"{truncated_content[:1000]}\n...\n{truncated_content[-1000:]}"
            if len(truncated_content) > 2000
            else truncated_content
        )
        logger.debug(
            f"Truncate current message to fit within max prompt size of {max_prompt_size} supported by {model_name} model:\n {truncated_snippet}"
        )

    return system_message + messages if system_message else messages


def reciprocal_conversation_to_chatml(message_pair):
    """Convert a single back and forth between user and assistant to chatml format"""
    return [ChatMessage(content=message, role=role) for message, role in zip(message_pair, ["user", "assistant"])]


def clean_json(response: str):
    """Remove any markdown json codeblock and newline formatting if present. Useful for non schema enforceable models"""
    return response.strip().replace("\n", "").removeprefix("```json").removesuffix("```")


def clean_mermaidjs(response: str):
    """Remove any markdown mermaidjs codeblock and newline formatting if present. Useful for non schema enforceable models"""
    return response.strip().removeprefix("```mermaid").removesuffix("```")


def clean_code_python(code: str):
    """Remove any markdown codeblock and newline formatting if present. Useful for non schema enforceable models"""
    return code.strip().removeprefix("```python").removesuffix("```")


def load_complex_json(json_str):
    """
    Preprocess a raw JSON string to
    - escape unescaped double quotes within value strings while preserving the JSON structure and already escaped quotes.
    - remove suffix after the first valid JSON object,
    """

    def replace_unescaped_quotes(match):
        # Get the content between colons and commas/end braces
        content = match.group(1)
        # Replace unescaped double, single quotes that aren't already escaped
        # Uses negative lookbehind to avoid replacing already escaped quotes
        # Replace " with \"
        processed_dq = re.sub(r'(?<!\\)"', '\\"', content)
        # Replace \' with \\'
        processed_final = re.sub(r"(?<!\\)\\'", r"\\\\'", processed_dq)
        return f': "{processed_final}"'

    # Match content between : and either , or }
    # This pattern looks for ': ' followed by any characters until , or }
    pattern = r':\s*"(.*?)(?<!\\)"(?=[,}])'

    # Process the JSON string
    cleaned = clean_json(rf"{json_str}")
    processed = re.sub(pattern, replace_unescaped_quotes, cleaned)

    # See which json loader can load the processed JSON as valid
    errors = []
    json_loaders_to_try = [json.loads, pyjson5.loads]
    for loads in json_loaders_to_try:
        try:
            return loads(processed)
        except (json.JSONDecodeError, pyjson5.Json5Exception) as e_load:
            loader_name = loads.__name__
            errors.append(f"{loader_name} (initial parse): {type(e_load).__name__}: {str(e_load)}")

            # Handle plain text suffixes by slicing at error position
            if hasattr(e_load, "pos") and 0 < e_load.pos < len(processed):
                try:
                    sliced = processed[: e_load.pos].strip()
                    if sliced:
                        return loads(sliced)
                except Exception as e_slice:
                    errors.append(
                        f"{loader_name} after slice at {e_load.pos}: {type(e_slice).__name__}: {str(e_slice)}"
                    )
    # If all loaders fail, raise the aggregated error
    raise ValueError(
        f"Failed to load JSON with errors: {'; '.join(errors)}\n\n"
        f"While attempting to load this cleaned JSON:\n{processed}"
    )


def defilter_query(query: str):
    """Remove any query filters in query"""
    defiltered_query = query
    filters: List[BaseFilter] = [WordFilter(), FileFilter(), DateFilter()]
    for filter in filters:
        defiltered_query = filter.defilter(defiltered_query)
    return defiltered_query


@dataclass
class ImageWithType:
    content: Any
    type: str


def get_image_from_base64(image: str, type="b64"):
    # Extract image type and base64 data from inline image data
    image_base64 = image.split(",", 1)[1]
    image_type = image.split(";", 1)[0].split(":", 1)[1]

    # Convert image to desired format
    if type == "b64":
        return ImageWithType(content=image_base64, type=image_type)
    elif type == "pil":
        image_data = base64.b64decode(image_base64)
        image_pil = PIL.Image.open(BytesIO(image_data))
        return ImageWithType(content=image_pil, type=image_type)
    elif type == "bytes":
        image_data = base64.b64decode(image_base64)
        return ImageWithType(content=image_data, type=image_type)


def get_image_from_url(image_url: str, type="pil"):
    try:
        response = requests.get(image_url)
        response.raise_for_status()  # Check if the request was successful

        # Get content type from response or infer from URL
        content_type = response.headers.get("content-type") or mimetypes.guess_type(image_url)[0] or "image/webp"

        # Convert image to desired format
        image_data: Any = None
        if type == "b64":
            image_data = base64.b64encode(response.content).decode("utf-8")
        elif type == "pil":
            image_data = PIL.Image.open(BytesIO(response.content))
        elif type == "bytes":
            image_data = response.content
        else:
            raise ValueError(f"Invalid image type: {type}")

        return ImageWithType(content=image_data, type=content_type)
    except requests.exceptions.RequestException as e:
        logger.error(f"Failed to get image from URL {image_url}: {e}")
        return ImageWithType(content=None, type=None)


def commit_conversation_trace(
    session: list[ChatMessage],
    response: str | list[dict],
    tracer: dict,
    system_message: str | list[dict] = "",
    repo_path: str = None,
) -> str:
    """
    Save trace of conversation step using git. Useful to visualize, compare and debug traces.
    Returns the path to the repository.
    """
    try:
        from git import Repo
    except ImportError:
        return None

    # Infer repository path from environment variable or provided path
    repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
    if not repo_path:
        return None

    # Serialize session, system message and response to yaml
    system_message_yaml = json.dumps(system_message, ensure_ascii=False, sort_keys=False)
    response_yaml = json.dumps(response, ensure_ascii=False, sort_keys=False)
    formatted_session = [{"role": message.role, "content": message.content} for message in session]
    session_yaml = json.dumps(formatted_session, ensure_ascii=False, sort_keys=False)
    query = (
        json.dumps(session[-1].content, ensure_ascii=False, sort_keys=False).strip().removeprefix("'").removesuffix("'")
    )  # Extract serialized query from chat session

    # Extract chat metadata for session
    uid, cid, mid = tracer.get("uid", "main"), tracer.get("cid", "main"), tracer.get("mid")

    try:
        # Prepare git repository
        os.makedirs(repo_path, exist_ok=True)
        repo = Repo.init(repo_path)

        # Remove post-commit hook if it exists
        hooks_dir = os.path.join(repo_path, ".git", "hooks")
        post_commit_hook = os.path.join(hooks_dir, "post-commit")
        if os.path.exists(post_commit_hook):
            os.remove(post_commit_hook)

        # Configure git user if not set
        if not repo.config_reader().has_option("user", "email"):
            repo.config_writer().set_value("user", "name", "Prompt Tracer").release()
            repo.config_writer().set_value("user", "email", "promptracer@khoj.dev").release()

        # Create an initial commit if the repository is newly created
        if not repo.head.is_valid():
            repo.index.commit("And then there was a trace")

        # Check out the initial commit
        initial_commit = repo.commit("HEAD~0")
        repo.head.reference = initial_commit
        repo.head.reset(index=True, working_tree=True)

        # Create or switch to user branch from initial commit
        user_branch = f"u_{uid}"
        if user_branch not in repo.branches:
            repo.create_head(user_branch)
        repo.heads[user_branch].checkout()

        # Create or switch to conversation branch from user branch
        conv_branch = f"c_{cid}"
        if conv_branch not in repo.branches:
            repo.create_head(conv_branch)
        repo.heads[conv_branch].checkout()

        # Create or switch to message branch from conversation branch
        msg_branch = f"m_{mid}" if mid else None
        if msg_branch and msg_branch not in repo.branches:
            repo.create_head(msg_branch)
        if msg_branch:
            repo.heads[msg_branch].checkout()

        # Include file with content to commit
        files_to_commit = {"query": session_yaml, "response": response_yaml, "system_prompt": system_message_yaml}

        # Write files and stage them
        for filename, content in files_to_commit.items():
            file_path = os.path.join(repo_path, filename)
            # Unescape special characters in content for better readability
            content = content.strip().replace("\\n", "\n").replace("\\t", "\t")
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(content)
            repo.index.add([filename])

        # Create commit
        metadata_yaml = yaml.dump(tracer, allow_unicode=True, sort_keys=False, default_flow_style=False)
        commit_message = f"""
{query[:250]}

Response:
---
{response[:500]}...

Metadata
---
{metadata_yaml}
""".strip()

        repo.index.commit(commit_message)

        logger.debug(f"Saved conversation trace to repo at {repo_path}")
        return repo_path
    except Exception as e:
        logger.error(f"Failed to add conversation trace to repo: {str(e)}", exc_info=True)
        return None


def merge_message_into_conversation_trace(query: str, response: str, tracer: dict, repo_path=None) -> bool:
    """
    Merge the message branch into its parent conversation branch.

    Args:
        query: User query
        response: Assistant response
        tracer: Dictionary containing uid, cid and mid
        repo_path: Path to the git repository

    Returns:
        bool: True if merge was successful, False otherwise
    """
    try:
        from git import Repo
    except ImportError:
        return False
    try:
        # Extract branch names
        msg_branch = f"m_{tracer['mid']}"
        conv_branch = f"c_{tracer['cid']}"

        # Infer repository path from environment variable or provided path
        repo_path = repo_path if not is_none_or_empty(repo_path) else os.getenv("PROMPTRACE_DIR")
        if not repo_path:
            return None
        repo = Repo(repo_path)

        # Checkout conversation branch
        repo.heads[conv_branch].checkout()

        # Create commit message
        metadata_yaml = yaml.dump(tracer, allow_unicode=True, sort_keys=False, default_flow_style=False)
        commit_message = f"""
{query[:250]}

Response:
---
{response[:500]}...

Metadata
---
{metadata_yaml}
""".strip()

        # Merge message branch into conversation branch
        repo.git.merge(msg_branch, no_ff=True, m=commit_message)

        # Delete message branch after merge
        repo.delete_head(msg_branch, force=True)

        logger.debug(f"Successfully merged {msg_branch} into {conv_branch}")
        return True
    except Exception as e:
        logger.error(f"Failed to merge message {msg_branch} into conversation {conv_branch}: {str(e)}", exc_info=True)
        return False


def messages_to_print(messages: list[ChatMessage], max_length: int = 70) -> str:
    """
    Format and truncate messages to print, ensuring JSON serializable content
    """

    def safe_serialize(content: Any) -> str:
        try:
            # Try JSON serialization
            json.dumps(content)
            return content
        except (TypeError, json.JSONDecodeError):
            # Handle non-serializable types
            if hasattr(content, "format") and content.format == "WEBP":
                return "[WebP Image]"
            elif hasattr(content, "__dict__"):
                return str(content.__dict__)
            return str(content)

    return "\n".join([f"{json.dumps(safe_serialize(message.content))[:max_length]}..." for message in messages])


class StructuredOutputSupport(int, Enum):
    NONE = 0
    OBJECT = 1
    SCHEMA = 2
    TOOL = 3


class ResponseWithThought:
    def __init__(self, text: str = None, thought: str = None, raw_content: list = None):
        self.text = text
        self.thought = thought
        self.raw_content = raw_content
