"""
Service for logging A2A tasks and events to the database.
"""

import copy
import json
import logging
import math
import uuid
from typing import Any, Callable, Dict, Union

from a2a.types import (
    A2ARequest,
    JSONRPCError,
    JSONRPCResponse,
    Task as A2ATask,
    TaskArtifactUpdateEvent,
    TaskStatusUpdateEvent,
)
from sqlalchemy.orm import Session as DBSession

from ....common import a2a
from ..repository.entities import Task, TaskEvent
from ..repository.task_repository import TaskRepository
from solace_agent_mesh.shared.utils.timestamp_utils import now_epoch_ms

log = logging.getLogger(__name__)

class TaskLoggerService:
    """Service for logging A2A tasks and events to the database."""

    def __init__(
        self, session_factory: Callable[[], DBSession] | None, config: Dict[str, Any]
    ):
        self.session_factory = session_factory
        self.config = config
        self.log_identifier = "[TaskLoggerService]"
        log.info(f"{self.log_identifier} Initialized.")

    def log_event(self, event_data: Dict[str, Any]):
        """
        Parses a raw A2A message and logs it as a task event.
        Creates or updates the master task record as needed.
        """
        if not self.config.get("enabled", False):
            return

        if not self.session_factory:
            log.warning(
                f"{self.log_identifier} Task logging is enabled but no database is configured. Skipping event."
            )
            return

        topic = event_data.get("topic")
        payload = event_data.get("payload")
        user_properties = event_data.get("user_properties", {})

        if not topic or not payload:
            log.warning(
                f"{self.log_identifier} Received event with missing topic or payload."
            )
            return

        if "/a2a/v1/discovery/" in topic:
            # Ignore discovery messages
            return

        if "/a2a/v1/trust/" in topic:
            # Ignore trust messages early to avoid queue buildup
            return

        # Parse the event into a Pydantic model first.
        parsed_event = self._parse_a2a_event(topic, payload)
        if parsed_event is None:
            # Parsing failed or event should be ignored.
            return

        db = self.session_factory()
        try:
            repo = TaskRepository()

            # Infer details from the parsed event
            direction, task_id, user_id, session_id = self._infer_event_details(
                topic, parsed_event, user_properties
            )

            if not task_id:
                log.debug(
                    f"{self.log_identifier} Could not determine task_id for event on topic {topic}. Skipping."
                )
                return

            # Check if we should log this event type
            if not self._should_log_event(topic, parsed_event):
                log.debug(
                    f"{self.log_identifier} Event on topic {topic} is configured to be skipped."
                )
                return

            # Sanitize the original raw payload before storing
            sanitized_payload = self._sanitize_payload(payload)

            # Check for existing task or create a new one
            task = repo.find_by_id(db, task_id)
            if not task:
                # Extract parent_task_id and background execution metadata
                parent_task_id = None
                background_execution_enabled = False
                max_execution_time_ms = None
                
                log.info(
                    f"{self.log_identifier} Creating new task {task_id}: direction={direction}, "
                    f"parsed_event_type={type(parsed_event).__name__}"
                )
                
                if direction == "request" and isinstance(parsed_event, A2ARequest):
                    message = a2a.get_message_from_send_request(parsed_event)
                    log.info(f"{self.log_identifier} Message extracted: {message is not None}")
                    
                    if message:
                        log.info(f"{self.log_identifier} Message metadata: {message.metadata}")
                        
                        if message.metadata:
                            parent_task_id = message.metadata.get("parentTaskId")
                            background_execution_enabled = message.metadata.get("backgroundExecutionEnabled", False)
                            # Default to 1 hour (3600000ms) if background execution is enabled but no timeout specified
                            max_execution_time_ms = message.metadata.get("maxExecutionTimeMs")
                            if background_execution_enabled and max_execution_time_ms is None:
                                max_execution_time_ms = 3600000  # 1 hour default
                        else:
                            log.warning(
                                f"{self.log_identifier} Message has no metadata for task {task_id}"
                            )
                    else:
                        log.warning(
                            f"{self.log_identifier} Could not extract message from request for task {task_id}"
                        )

                if direction == "request":
                    initial_text = self._extract_initial_text(parsed_event)
                    current_time = now_epoch_ms()
                    new_task = Task(
                        id=task_id,
                        user_id=user_id or "unknown",
                        parent_task_id=parent_task_id,
                        start_time=current_time,
                        initial_request_text=(
                            initial_text[:1024] if initial_text else None
                        ),  # Truncate
                        execution_mode="background" if background_execution_enabled else "foreground",
                        last_activity_time=current_time,
                        background_execution_enabled=background_execution_enabled,
                        max_execution_time_ms=max_execution_time_ms,
                        session_id=session_id,  # Store session_id for persistent event buffering
                    )
                    repo.save_task(db, new_task)
                    log.info(
                        f"{self.log_identifier} Created new task record for ID: {task_id}"
                        + (f" with parent: {parent_task_id}" if parent_task_id else "")
                        + (f" (background execution enabled)" if background_execution_enabled else "")
                        + (f" (session: {session_id})" if session_id else "")
                    )
                else:
                    # We received an event for a task we haven't seen the start of.
                    # This can happen if the logger starts mid-conversation. Create a placeholder.
                    current_time = now_epoch_ms()
                    placeholder_task = Task(
                        id=task_id,
                        user_id=user_id or "unknown",
                        parent_task_id=parent_task_id,
                        start_time=current_time,
                        initial_request_text="[Task started before logger was active]",
                        execution_mode="foreground",
                        last_activity_time=current_time,
                        background_execution_enabled=False,
                        session_id=session_id,  # Store session_id for persistent event buffering
                    )
                    repo.save_task(db, placeholder_task)
                    log.info(
                        f"{self.log_identifier} Created placeholder task record for ID: {task_id}"
                    )
            else:
                # Update last activity time for existing task
                # This is a non-critical update that can fail due to cross-process SQLite concurrency
                try:
                    task.last_activity_time = now_epoch_ms()
                    repo.save_task(db, task)
                except Exception as activity_update_error:
                    # StaleDataError or other concurrency issues - log and continue
                    # The task may have been modified/deleted by another process (FastAPI vs SAC)
                    log.debug(
                        f"{self.log_identifier} Non-critical: Failed to update last_activity_time for task {task_id}: {activity_update_error}"
                    )
                    # Rollback and begin a new transaction so subsequent operations can continue
                    db.rollback()
                    db.begin()

            # Create and save the event using the sanitized raw payload
            task_event = TaskEvent(
                id=str(uuid.uuid4()),
                task_id=task_id,
                user_id=user_id,
                created_time=now_epoch_ms(),
                topic=topic,
                direction=direction,
                payload=sanitized_payload,
            )
            repo.save_event(db, task_event)

            # If it's a final event, update the master task record
            final_status = self._get_final_status(parsed_event)
            if final_status:
                task_to_update = repo.find_by_id(db, task_id)
                if task_to_update:
                    current_time = now_epoch_ms()
                    task_to_update.end_time = current_time
                    task_to_update.status = final_status
                    task_to_update.last_activity_time = current_time
                    
                    # Extract and store token usage if present
                    if isinstance(parsed_event, A2ATask) and parsed_event.metadata:
                        token_usage = parsed_event.metadata.get("token_usage")
                        if token_usage and isinstance(token_usage, dict):
                            task_to_update.total_input_tokens = token_usage.get("total_input_tokens")
                            task_to_update.total_output_tokens = token_usage.get("total_output_tokens")
                            task_to_update.total_cached_input_tokens = token_usage.get("total_cached_input_tokens")
                            task_to_update.token_usage_details = token_usage
                            log.info(
                                f"{self.log_identifier} Stored token usage for task {task_id}: "
                                f"input={token_usage.get('total_input_tokens')}, "
                                f"output={token_usage.get('total_output_tokens')}, "
                                f"cached={token_usage.get('total_cached_input_tokens')}"
                            )

                    repo.save_task(db, task_to_update)
                    log.info(
                        f"{self.log_identifier} Finalized task record for ID: {task_id} with status: {final_status}"
                    )
                    
            
            db.commit()
        except Exception as e:
            log.exception(
                f"{self.log_identifier} Error logging event on topic {topic}: {e}"
            )
            db.rollback()
        finally:
            db.close()

    def _parse_a2a_event(self, topic: str, payload: dict) -> Union[
        A2ARequest,
        A2ATask,
        TaskStatusUpdateEvent,
        TaskArtifactUpdateEvent,
        JSONRPCError,
        None,
    ]:
        """
        Safely parses a raw A2A message payload into a Pydantic model.
        Returns the parsed model or None if parsing fails or is not applicable.
        """
        # Ignore discovery messages (agents and gateways)
        if "/discovery/" in topic:
            return None
        # Ignore trust manager trust card messages
        if "/trust/" in topic:
            return None

        try:
            # Check if it's a response (has 'result' or 'error')
            if "result" in payload or "error" in payload:
                rpc_response = JSONRPCResponse.model_validate(payload)
                error = a2a.get_response_error(rpc_response)
                if error:
                    return error
                result = a2a.get_response_result(rpc_response)
                if result:
                    # The result is already a parsed Pydantic model
                    return result
            # Check if it's a request
            elif "method" in payload:
                return A2ARequest.model_validate(payload)

            log.warning(
                f"{self.log_identifier} Payload for topic '{topic}' is not a recognizable JSON-RPC request or response. Payload: {payload}"
            )
            return None

        except Exception as e:
            log.error(
                f"{self.log_identifier} Failed to parse A2A event for topic '{topic}': {e}. Payload: {payload}"
            )
            return None

    def _infer_event_details(
        self, topic: str, parsed_event: Any, user_props: Dict | None
    ) -> tuple[str, str | None, str | None, str | None]:
        """Infers direction, task_id, user_id, and session_id from a parsed A2A event.
        
        Returns:
            Tuple of (direction, task_id, user_id, session_id)
        """
        direction = "unknown"
        task_id = None
        session_id = None  # Will be extracted from context_id
        # Ensure user_props is a dict, not None
        user_props = user_props or {}
        user_id = user_props.get("userId")

        if isinstance(parsed_event, A2ARequest):
            direction = "request"
            task_id = a2a.get_request_id(parsed_event)
            # Extract session_id from context_id in the message
            message = a2a.get_message_from_send_request(parsed_event)
            if message:
                session_id = a2a.get_context_id(message)
        elif isinstance(
            parsed_event, (A2ATask, TaskStatusUpdateEvent, TaskArtifactUpdateEvent)
        ):
            direction = "response" if isinstance(parsed_event, A2ATask) else "status"
            task_id = getattr(parsed_event, "task_id", None) or getattr(
                parsed_event, "id", None
            )
            # Extract session_id from context_id
            session_id = getattr(parsed_event, "context_id", None)
        elif isinstance(parsed_event, JSONRPCError):
            direction = "error"
            if isinstance(parsed_event.data, dict):
                task_id = parsed_event.data.get("taskId")
                session_id = parsed_event.data.get("contextId")

        if not user_id:
            user_config = user_props.get("a2aUserConfig") or user_props.get("a2a_user_config")
            if isinstance(user_config, dict):
                user_profile = user_config.get("user_profile", {})
                if isinstance(user_profile, dict):
                    user_id = user_profile.get("id")

        return direction, str(task_id) if task_id else None, user_id, session_id

    def _extract_initial_text(self, parsed_event: Any) -> str | None:
        """Extracts the initial text from a send message request."""
        try:
            if isinstance(parsed_event, A2ARequest):
                message = a2a.get_message_from_send_request(parsed_event)
                if message:
                    return a2a.get_text_from_message(message)
        except Exception:
            return None
        return None

    def _get_final_status(self, parsed_event: Any) -> str | None:
        """Checks if a parsed event represents a final task status and returns the state."""
        if isinstance(parsed_event, A2ATask):
            return parsed_event.status.state.value
        elif isinstance(parsed_event, JSONRPCError):
            return "failed"
        return None

    def _should_log_event(self, topic: str, parsed_event: Any) -> bool:
        """Determines if an event should be logged based on configuration."""
        if not self.config.get("log_status_updates", True):
            if "status" in topic:
                return False
        if not self.config.get("log_artifact_events", True):
            if isinstance(parsed_event, TaskArtifactUpdateEvent):
                return False
        return True

    @staticmethod
    def _sanitize_non_finite_floats(value: Any) -> Any:
        """
        Recursively sanitize a value, replacing non-finite floats (NaN, Infinity, -Infinity)
        with None since PostgreSQL JSON type doesn't support these values.
        """
        if isinstance(value, float):
            if math.isnan(value) or math.isinf(value):
                return None
            return value
        elif isinstance(value, dict):
            return {k: TaskLoggerService._sanitize_non_finite_floats(v) for k, v in value.items()}
        elif isinstance(value, list):
            return [TaskLoggerService._sanitize_non_finite_floats(item) for item in value]
        else:
            return value

    def _sanitize_payload(self, payload: Dict) -> Dict:
        """
        Sanitizes payload for database storage:
        1. Strips or truncates file content based on configuration
        2. Replaces non-finite floats (NaN, Infinity, -Infinity) with None
           since PostgreSQL JSON type doesn't support these values
        """
        new_payload = copy.deepcopy(payload)

        def walk_and_sanitize(node):
            if isinstance(node, dict):
                for key, value in list(node.items()):
                    # Sanitize non-finite floats using the helper
                    node[key] = self._sanitize_non_finite_floats(value)
                    
                    if key == "parts" and isinstance(node[key], list):
                        new_parts = []
                        for part in node[key]:
                            if isinstance(part, dict) and "file" in part:
                                if not self.config.get("log_file_parts", True):
                                    continue  # Skip this part entirely

                                file_dict = part.get("file")
                                if isinstance(file_dict, dict) and "bytes" in file_dict:
                                    max_bytes = self.config.get(
                                        "max_file_part_size_bytes", 102400
                                    )
                                    file_bytes_b64 = file_dict.get("bytes")
                                    if isinstance(file_bytes_b64, str):
                                        if (len(file_bytes_b64) * 3 / 4) > max_bytes:
                                            file_dict["bytes"] = (
                                                f"[Content stripped, size > {max_bytes} bytes]"
                                            )
                                new_parts.append(part)
                            else:
                                walk_and_sanitize(part)
                                new_parts.append(part)
                        node["parts"] = new_parts
                    elif isinstance(node[key], (dict, list)):
                        walk_and_sanitize(node[key])
            elif isinstance(node, list):
                for i, item in enumerate(node):
                    node[i] = self._sanitize_non_finite_floats(item)
                    if isinstance(node[i], (dict, list)):
                        walk_and_sanitize(node[i])

        walk_and_sanitize(new_payload)
        return new_payload

    def _save_chat_messages_for_background_task(
        self, db: DBSession, task_id: str, task: Task, repo: TaskRepository
    ) -> None:
        """
        Save chat messages for a completed background task by reconstructing them from task events.
        This ensures chat history is available when users return to a session after a background task completes.
        Uses upsert to avoid duplicates.
        
        NOTE: Even if SSE events are buffered (task.events_buffered=True), we still save to chat_tasks
        from task_events. The chat_tasks data will have unresolved embeds, but this serves as a fallback.
        The frontend should prefer replaying from sse_event_buffer when available.
        """
        try:
            # Get all events for this task
            task_with_events = repo.find_by_id_with_events(db, task_id)
            if not task_with_events:
                log.warning(
                    f"{self.log_identifier} Could not find task {task_id} with events for chat message saving"
                )
                return
            
            _, events = task_with_events
            
            # Extract session_id and user_id from the task's initial request
            session_id = None
            user_id = task.user_id
            agent_name = None
            user_message_text = task.initial_request_text
            
            # Parse events to extract session context and reconstruct messages
            message_bubbles = []
            artifacts = []  # Track artifacts - only from final task response to avoid duplicates
            rag_data = []  # Track RAG metadata from tool results
            accumulated_agent_text = []  # Accumulate text from all status updates
            accumulated_agent_parts = []  # Accumulate parts from all status updates
            
            for event in events:
                try:
                    payload = event.payload
                    
                    # Extract session_id from the first request event
                    if event.direction == "request" and not session_id:
                        if "params" in payload and isinstance(payload["params"], dict):
                            message = payload["params"].get("message", {})
                            if isinstance(message, dict):
                                session_id = message.get("contextId")
                                # Extract agent name from metadata
                                metadata = message.get("metadata", {})
                                if isinstance(metadata, dict):
                                    agent_name = metadata.get("agent_name")
                                
                                # Add user message bubble
                                parts = message.get("parts", [])
                                
                                # Filter out the gateway timestamp part (first part if it starts with "Request received by gateway")
                                filtered_parts = []
                                for i, part in enumerate(parts):
                                    if part.get("kind") == "text":
                                        text = part.get("text", "")
                                        # Skip the first part if it's the gateway timestamp
                                        if i == 0 and text.startswith("Request received by gateway at:"):
                                            continue
                                        filtered_parts.append(part)
                                    else:
                                        filtered_parts.append(part)
                                
                                text_parts = [p.get("text", "") for p in filtered_parts if p.get("kind") == "text"]
                                combined_text = "".join(text_parts)
                                
                                if combined_text or any(p.get("kind") == "file" for p in filtered_parts):
                                    message_bubbles.append({
                                        "id": f"msg-{uuid.uuid4()}",
                                        "type": "user",
                                        "text": combined_text,
                                        "parts": filtered_parts,
                                    })
                    
                    # Collect text content AND RAG metadata from status events
                    # Text is accumulated to reconstruct the full agent response
                    elif event.direction == "status":
                        if "result" in payload:
                            result = payload["result"]
                            
                            # Extract text content from status updates
                            status = result.get("status", {})
                            if isinstance(status, dict):
                                message = status.get("message", {})
                                if isinstance(message, dict):
                                    parts = message.get("parts", [])
                                    for part in parts:
                                        if isinstance(part, dict):
                                            part_kind = part.get("kind")
                                            if part_kind == "text":
                                                text = part.get("text", "")
                                                if text:
                                                    accumulated_agent_text.append(text)
                                                    accumulated_agent_parts.append(part)
                                            elif part_kind == "data":
                                                # Extract RAG metadata from tool_result data parts
                                                data = part.get("data", {})
                                                if isinstance(data, dict):
                                                    data_type = data.get("type")
                                                    
                                                    if data_type == "tool_result":
                                                        result_data = data.get("result_data", {})
                                                        if isinstance(result_data, dict) and "rag_metadata" in result_data:
                                                            rag_metadata = result_data["rag_metadata"]
                                                            if isinstance(rag_metadata, dict):
                                                                # Add taskId to the RAG metadata
                                                                rag_metadata["taskId"] = task_id
                                                                rag_data.append(rag_metadata)
                                                                log.info(
                                                                    f"{self.log_identifier} Extracted RAG metadata for task {task_id}: "
                                                                    f"searchType={rag_metadata.get('searchType')}, "
                                                                    f"sources_count={len(rag_metadata.get('sources', []))}"
                                                                )
                                                    elif data_type == "artifact_creation_progress":
                                                        # Handle cancelled artifacts with rolled back text
                                                        if data.get("status") == "cancelled":
                                                            rolled_back_text = data.get("rolled_back_text")
                                                            if rolled_back_text:
                                                                accumulated_agent_text.append(rolled_back_text)
                                                                log.info(
                                                                    f"{self.log_identifier} Extracted rolled_back_text from cancelled artifact event for task {task_id}"
                                                                )
                                                
                                                accumulated_agent_parts.append(part)
                                            else:
                                                # Accumulate other non-text, non-data parts
                                                accumulated_agent_parts.append(part)
                    
                    # Extract artifacts and any additional text from final task response
                    elif event.direction == "response":
                        if "result" in payload:
                            result = payload["result"]
                            
                            # Only process final task response (kind="task")
                            if isinstance(result, dict) and result.get("kind") == "task":
                                # Extract artifacts from task metadata
                                metadata = result.get("metadata", {})
                                if isinstance(metadata, dict):
                                    # Try both 'produced_artifacts' and 'artifact_manifest'
                                    artifact_list = metadata.get("produced_artifacts") or metadata.get("artifact_manifest", [])
                                    if isinstance(artifact_list, list):
                                        for artifact_info in artifact_list:
                                            if isinstance(artifact_info, dict):
                                                # Handle both 'name' and 'filename' keys
                                                artifact_name = artifact_info.get("name") or artifact_info.get("filename")
                                                # Skip web_content_ artifacts (temporary files from deep research)
                                                if artifact_name and not artifact_name.startswith("web_content_"):
                                                    artifacts.append({
                                                        "kind": "artifact",
                                                        "status": "completed",
                                                        "name": artifact_name,
                                                        "file": {
                                                            "name": artifact_name,
                                                            "mime_type": artifact_info.get("mime_type"),
                                                            "uri": f"artifact://{session_id}/{artifact_name}" if session_id else f"artifact://unknown/{artifact_name}"
                                                        }
                                                    })
                                
                                # Final task object - extract any additional text not in status updates
                                status = result.get("status", {})
                                if isinstance(status, dict):
                                    message = status.get("message", {})
                                    if isinstance(message, dict):
                                        parts = message.get("parts", [])
                                        
                                        # Extract RAG metadata from tool_result data parts in final response
                                        for part in parts:
                                            if isinstance(part, dict) and part.get("kind") == "data":
                                                data = part.get("data", {})
                                                if isinstance(data, dict) and data.get("type") == "tool_result":
                                                    result_data = data.get("result_data", {})
                                                    if isinstance(result_data, dict) and "rag_metadata" in result_data:
                                                        rag_metadata = result_data["rag_metadata"]
                                                        if isinstance(rag_metadata, dict):
                                                            # Add taskId to the RAG metadata
                                                            rag_metadata["taskId"] = task_id
                                                            # Avoid duplicates
                                                            if rag_metadata not in rag_data:
                                                                rag_data.append(rag_metadata)
                                                                log.info(
                                                                    f"{self.log_identifier} Extracted RAG metadata from final response for task {task_id}: "
                                                                    f"searchType={rag_metadata.get('searchType')}, "
                                                                    f"sources_count={len(rag_metadata.get('sources', []))}"
                                                                )
                                        
                                        
                except Exception as e:
                    log.warning(
                        f"{self.log_identifier} Error parsing event for chat message reconstruction: {e}"
                    )
                    continue
            
            # After processing all events, create the agent message bubble from accumulated content
            if accumulated_agent_text or artifacts:
                combined_text = "".join(accumulated_agent_text).strip()
                
                # Check if artifact markers are already in the texts
                import re
                existing_markers = set()
                marker_pattern = r'«artifact_return:([^»]+)»'
                for match in re.finditer(marker_pattern, combined_text):
                    # Normalize the artifact name (strip version suffix)
                    artifact_ref = match.group(1)
                    if ':' in artifact_ref:
                        base_name = artifact_ref.rsplit(':', 1)[0]
                        try:
                            int(artifact_ref.rsplit(':', 1)[1])
                            # Add the base name (without version) to existing markers
                            # so we don't add a duplicate marker later
                            existing_markers.add(base_name)
                        except ValueError:
                            # Not a version number, treat the whole thing as the artifact name
                            existing_markers.add(artifact_ref)
                    else:
                        existing_markers.add(artifact_ref)
                
                # Only add artifact markers if they're not already present
                for artifact in artifacts:
                    artifact_name = artifact['name']
                    if artifact_name not in existing_markers:
                        combined_text += f"«artifact_return:{artifact_name}»"
                        log.info(
                            f"{self.log_identifier} Adding artifact marker for {artifact_name}"
                        )
                    else:
                        log.info(
                            f"{self.log_identifier} Skipping duplicate artifact marker for {artifact_name} (already in text)"
                        )
                
                # Filter out data parts from accumulated parts
                content_parts = [p for p in accumulated_agent_parts if p.get("kind") != "data"]
                
                message_bubbles.append({
                    "id": f"msg-{uuid.uuid4()}",
                    "type": "agent",
                    "text": combined_text,
                    "parts": content_parts,  # Only content parts, no artifacts
                })
                
            
            # Only save if we have a session_id and at least one message
            if not session_id:
                log.warning(
                    f"{self.log_identifier} Could not extract session_id for task {task_id}, skipping chat message save"
                )
                return
            
            if not message_bubbles:
                log.warning(
                    f"{self.log_identifier} No message bubbles reconstructed for task {task_id}, skipping chat message save"
                )
                return
            
            # Import here to avoid circular dependency
            from ..repository.chat_task_repository import ChatTaskRepository
            from ..repository.entities import ChatTask
            from ..repository.session_repository import SessionRepository
            
            # Check if the session exists in this database
            session_repo = SessionRepository()
            if not session_repo.exists(db, session_id):
                log.debug(
                    f"{self.log_identifier} Session {session_id} not found in webui_gateway database "
                    f"Skipping chat message save for task {task_id}"
                )
                return
            
            # Check if a chat task already exists (frontend may have saved it first with frontend-only fields)
            # If so, preserve frontend-only fields like contextQuote and displayHtml from the user message
            chat_task_repo = ChatTaskRepository()
            existing_chat_task = chat_task_repo.find_by_id(db, task_id)
            if existing_chat_task:
                try:
                    existing_bubbles = json.loads(existing_chat_task.message_bubbles) if isinstance(existing_chat_task.message_bubbles, str) else existing_chat_task.message_bubbles
                    # Find the existing user message bubble
                    existing_user_bubble = next((b for b in existing_bubbles if b.get("type") == "user"), None)
                    if existing_user_bubble:
                        # Extract frontend-only fields
                        frontend_only_fields = {}
                        if existing_user_bubble.get("contextQuote"):
                            frontend_only_fields["contextQuote"] = existing_user_bubble["contextQuote"]
                        if existing_user_bubble.get("contextQuoteSourceId"):
                            frontend_only_fields["contextQuoteSourceId"] = existing_user_bubble["contextQuoteSourceId"]
                        if existing_user_bubble.get("displayHtml"):
                            frontend_only_fields["displayHtml"] = existing_user_bubble["displayHtml"]
                        
                        if frontend_only_fields:
                            # Find the reconstructed user message bubble and merge frontend-only fields
                            for bubble in message_bubbles:
                                if bubble.get("type") == "user":
                                    bubble.update(frontend_only_fields)
                                    log.info(
                                        f"{self.log_identifier} Preserved frontend-only fields for task {task_id}: "
                                        f"contextQuote={bool(frontend_only_fields.get('contextQuote'))}, "
                                        f"contextQuoteSourceId={bool(frontend_only_fields.get('contextQuoteSourceId'))}, "
                                        f"displayHtml={bool(frontend_only_fields.get('displayHtml'))}"
                                    )
                                    break
                except Exception as e:
                    log.warning(
                        f"{self.log_identifier} Failed to extract frontend-only fields from existing chat task {task_id}: {e}"
                    )
            
            # Build task metadata including RAG data if present
            task_metadata_dict = {
                "schema_version": 1,
                "status": task.status,
                "agent_name": agent_name,
            }
            
            # Include RAG data if we found any
            if rag_data:
                task_metadata_dict["rag_data"] = rag_data
                log.info(
                    f"{self.log_identifier} Including {len(rag_data)} RAG data entries in task metadata for {task_id}"
                )
            
            # Create and save the chat task
            chat_task = ChatTask(
                id=task_id,
                session_id=session_id,
                user_id=user_id,
                user_message=user_message_text,
                message_bubbles=json.dumps(message_bubbles),
                task_metadata=json.dumps(task_metadata_dict),
                created_time=task.start_time,
                updated_time=task.end_time,
            )
            
            # chat_task_repo was already created above when checking for existing task
            chat_task_repo.save(db, chat_task)
            
            log.info(
                f"{self.log_identifier} Saved chat messages for background task {task_id} "
                f"(session: {session_id}, {len(message_bubbles)} message bubbles)"
            )
            
        except Exception as e:
            log.error(
                f"{self.log_identifier} Failed to save chat messages for background task {task_id}: {e}",
                exc_info=True
            )
