chat.py

import json
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Callable, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator

from models.feedback import MAX_COMMENT_LENGTH, FeedbackReason

# App display names are resolved by an injected callable, never by importing the database
# layer here: models/ must stay import-pure so a Pydantic module cannot drag the Firestore
# client into every import graph that touches a chat message.
AppNameResolver = Callable[[str], Optional[str]]


class MessageSender(str, Enum):
    ai = 'ai'
    human = 'human'


class MessageType(str, Enum):
    text = 'text'
    day_summary = 'day_summary'


class MessageConversationStructured(BaseModel):
    title: str
    emoji: str


class MessageConversation(BaseModel):
    id: str
    structured: MessageConversationStructured
    created_at: datetime


# Chat Completions `file` content parts (purpose=user_data) accept these document
# types. Images stay on the separate vision / image_url path and are not listed here.
# Source: https://platform.openai.com/docs/guides/pdf-files
# (File inputs → "Accepted file types" and "Full list of accepted file types").
# Read 2026-09-08 from a saved copy of that page. Non-PDF documents are
# text-extracted only; spreadsheets use a 1,000-row augmentation flow.
CHAT_FILE_DOCUMENT_EXTENSIONS: frozenset[str] = frozenset(
    {
        # PDF
        'pdf',
        # Spreadsheets
        'xla',
        'xlb',
        'xlc',
        'xlm',
        'xls',
        'xlsx',
        'xlt',
        'xlw',
        'csv',
        'tsv',
        'iif',
        # Rich documents
        'doc',
        'docx',
        'dot',
        'odt',
        'rtf',
        # Presentations
        'pot',
        'ppa',
        'pps',
        'ppt',
        'pptx',
        'pwz',
        'wiz',
        # Text and code
        'asm',
        'bat',
        'c',
        'cc',
        'conf',
        'cpp',
        'css',
        'cxx',
        'def',
        'dic',
        'eml',
        'h',
        'hh',
        'htm',
        'html',
        'ics',
        'ifb',
        'in',
        'js',
        'json',
        'ksh',
        'list',
        'log',
        'markdown',
        'md',
        'mht',
        'mhtml',
        'mime',
        'mjs',
        'nws',
        'pl',
        'py',
        'rst',
        's',
        'sql',
        'srt',
        'text',
        'txt',
        'vcf',
        'vtt',
        'xml',
        # The provider's Text-and-code MIME column also documents types whose
        # canonical suffixes are missing from its Extensions column (e.g.
        # 'text/x-rust' -> 'rs', 'text/x-yaml' -> 'yaml'). Keep them
        # allowlisted so the extension-first branch cannot reject a file whose
        # MIME the provider explicitly documents.
        'astro',
        'awk',
        'c++',
        'clj',
        'cmake',
        'cs',
        'dart',
        'diff',
        'dockerfile',
        'ejs',
        'elixir',
        'erb',
        'erlang',
        'go',
        'gradle',
        'graphql',
        'groovy',
        'hbs',
        'hcl',
        'hs',
        'ini',
        'j2',
        'jade',
        'java',
        'jl',
        'json5',
        'jsx',
        'kt',
        'lisp',
        'liquid',
        'lua',
        'm',
        'mk',
        'mustache',
        'ndjson',
        'patch',
        'php',
        'ps1',
        'proto',
        'pug',
        'r',
        'rb',
        'rs',
        'sass',
        'scala',
        'scss',
        'sh',
        'swift',
        'terraform',
        'tf',
        'tex',
        'tmpl',
        'toml',
        'ts',
        'tsx',
        'vbs',
        'yaml',
        'yml',
        'zsh',
    }
)
CHAT_FILE_DOCUMENT_MIME_TYPES: frozenset[str] = frozenset(
    {
        'application/csv',
        'application/graphql',
        'application/javascript',
        'application/json',
        'application/json5',
        'application/msword',
        'application/pdf',
        'application/rtf',
        'application/toml',
        'application/typescript',
        'application/vnd.apple.iwork',
        'application/vnd.apple.keynote',
        'application/vnd.apple.pages',
        'application/vnd.google-apps.document',
        'application/vnd.google-apps.presentation',
        'application/vnd.google-apps.spreadsheet',
        'application/vnd.ms-excel',
        'application/vnd.ms-powerpoint',
        'application/vnd.oasis.opendocument.text',
        'application/vnd.openxmlformats-officedocument.presentationml.presentation',
        'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
        'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        'application/x-awk',
        'application/x-bash',
        'application/x-graphql',
        'application/x-httpd-php',
        'application/x-httpd-php-source',
        'application/x-iif',
        'application/x-json5',
        'application/x-ndjson',
        'application/x-patch',
        'application/x-php',
        'application/x-powershell',
        'application/x-protobuf',
        'application/x-rust',
        'application/x-scala',
        'application/x-sql',
        'application/x-subrip',
        'application/x-terraform',
        'application/x-toml',
        'application/x-yaml',
        'application/yaml',
        'message/rfc822',
        'text/calendar',
        'text/css',
        'text/csv',
        'text/html',
        'text/javascript',
        'text/jsx',
        'text/markdown',
        'text/plain',
        'text/rtf',
        'text/srt',
        'text/tsv',
        'text/tsx',
        'text/vbscript',
        'text/vtt',
        'text/x-R',
        'text/x-asm',
        'text/x-astro',
        'text/x-awk',
        'text/x-bash',
        'text/x-c',
        'text/x-c++',
        'text/x-clojure',
        'text/x-cmake',
        'text/x-csharp',
        'text/x-dart',
        'text/x-diff',
        'text/x-dockerfile',
        'text/x-ejs',
        'text/x-elixir',
        'text/x-erb',
        'text/x-erlang',
        'text/x-go',
        'text/x-golang',
        'text/x-gradle',
        'text/x-graphql',
        'text/x-groovy',
        'text/x-handlebars',
        'text/x-haskell',
        'text/x-hcl',
        'text/x-iif',
        'text/x-ini',
        'text/x-jade',
        'text/x-java',
        'text/x-jinja2',
        'text/x-julia',
        'text/x-kotlin',
        'text/x-less',
        'text/x-lisp',
        'text/x-liquid',
        'text/x-lua',
        'text/x-makefile',
        'text/x-mustache',
        'text/x-objectivec',
        'text/x-objectivec++',
        'text/x-patch',
        'text/x-perl',
        'text/x-php',
        'text/x-properties',
        'text/x-protobuf',
        'text/x-pug',
        'text/x-python',
        'text/x-r',
        'text/x-rst',
        'text/x-ruby',
        'text/x-rust',
        'text/x-sass',
        'text/x-scala',
        'text/x-script.python',
        'text/x-scss',
        'text/x-sh',
        'text/x-shellscript',
        'text/x-sql',
        'text/x-subrip',
        'text/x-swift',
        'text/x-terraform',
        'text/x-tex',
        'text/x-tmpl',
        'text/x-toml',
        'text/x-twig',
        'text/x-typescript',
        'text/x-vcard',
        'text/x-yaml',
        'text/x-zsh',
        'text/xml',
    }
)


def _chat_file_normalized_mime(mime_type: Optional[str]) -> str:
    """Lowercase MIME, treating missing values and str(None)=='None' as empty."""
    if mime_type is None:
        return ''
    mime = str(mime_type).strip().lower()
    if mime in ('', 'none', 'null'):
        return ''
    return mime


def chat_file_is_document(name: str, mime_type: Optional[str]) -> bool:
    """True for any OpenAI-allowlisted chat document type, including PDF.

    Images are not in this allowlist. Extension is checked first and wins when
    present; MIME is the fallback so a missing MIME neither admits nor rejects
    on its own.
    """
    ext = Path(name or '').suffix.lower().lstrip('.')
    if ext:
        return ext in CHAT_FILE_DOCUMENT_EXTENSIONS
    mime = _chat_file_normalized_mime(mime_type)
    return bool(mime) and mime in CHAT_FILE_DOCUMENT_MIME_TYPES


class FileChat(BaseModel):
    id: str
    name: str
    thumbnail: Optional[str] = ""
    mime_type: str
    openai_file_id: str
    created_at: datetime
    thumb_name: Optional[str] = ""

    def is_image(self):
        return self.mime_type.startswith("image")

    def is_document(self) -> bool:
        return chat_file_is_document(self.name, self.mime_type)

    def model_dump(self, **kwargs):
        exclude_fields = {'thumb_name'}
        return super().model_dump(exclude=exclude_fields, **kwargs)


class ChartDataPoint(BaseModel):
    label: str
    value: float


class ChartDataset(BaseModel):
    label: str
    data_points: List[ChartDataPoint]
    color: Optional[str] = None  # hex color, e.g. "#4CAF50"


class ChartData(BaseModel):
    chart_type: Literal['line', 'bar']
    title: str
    x_label: Optional[str] = None
    y_label: Optional[str] = None
    datasets: List[ChartDataset]


class ChatEvidenceReference(BaseModel):
    """One bounded, optional source reference attached to a chat answer.

    The answer text remains authoritative.  Clients may render these references
    as supplemental chrome, but an unavailable or future reference must never
    make the answer itself unreadable.
    """

    id: str = Field(..., min_length=1, max_length=256)
    kind: str
    state: str
    title: Optional[str] = Field(None, max_length=160)
    summary: Optional[str] = Field(None, max_length=600)
    conversation_id: Optional[str] = Field(None, max_length=256)
    segment_id: Optional[str] = Field(None, max_length=256)
    frame_id: Optional[str] = Field(None, max_length=256)
    request_id: Optional[str] = Field(None, max_length=256)
    start_ms: Optional[int] = Field(None, ge=0)
    end_ms: Optional[int] = Field(None, ge=0)
    captured_at_ms: Optional[int] = Field(None, ge=0)
    error_code: Optional[str] = Field(None, max_length=128)
    error_message: Optional[str] = Field(None, max_length=600)
    metadata: Dict[str, Any] = Field(default_factory=dict)

    @field_validator('id')
    @classmethod
    def _normalize_evidence_id(cls, value: str) -> str:
        normalized = value.strip()
        if not normalized:
            raise ValueError('evidence id must not be blank')
        return normalized

    @field_validator('kind')
    @classmethod
    def _normalize_evidence_kind(cls, value: str) -> str:
        normalized = (value or '').strip().lower()
        return (
            normalized
            if normalized in {'conversation_summary', 'conversation_segment', 'screen', 'keyframe', 'request'}
            else 'unknown'
        )

    @field_validator('state')
    @classmethod
    def _normalize_evidence_state(cls, value: str) -> str:
        normalized = (value or '').strip().lower()
        return normalized if normalized in {'available', 'loading', 'offline', 'pruned', 'failed'} else 'unknown'

    @field_validator(
        'title',
        'summary',
        'conversation_id',
        'segment_id',
        'frame_id',
        'request_id',
        'error_code',
        'error_message',
    )
    @classmethod
    def _strip_evidence_strings(cls, value: Optional[str]) -> Optional[str]:
        if value is None:
            return None
        normalized = value.strip()
        return normalized or None

    @model_validator(mode='after')
    def _validate_evidence_identity(self) -> 'ChatEvidenceReference':
        try:
            serialized_metadata = json.dumps(self.metadata, sort_keys=True, separators=(',', ':'))
        except (TypeError, ValueError) as exc:
            raise ValueError('evidence metadata must be JSON serializable') from exc
        if len(self.metadata) > 16 or len(serialized_metadata) > 2_000:
            raise ValueError('evidence metadata exceeds the bounded transport limit')
        if self.end_ms is not None and self.start_ms is not None and self.end_ms < self.start_ms:
            raise ValueError('end_ms must be greater than or equal to start_ms')
        if self.kind == 'conversation_summary' and not self.conversation_id:
            raise ValueError('conversation_summary requires conversation_id')
        if self.kind == 'conversation_segment' and not (self.conversation_id and self.segment_id):
            raise ValueError('conversation_segment requires conversation_id and segment_id')
        if self.kind in {'screen', 'keyframe'} and not self.frame_id:
            raise ValueError(f'{self.kind} requires frame_id')
        if self.kind == 'request' and not self.request_id:
            raise ValueError('request requires request_id')
        return self


class ChatEvidenceEnvelope(BaseModel):
    """Versioned transport envelope for supplemental chat evidence."""

    schema_version: int = Field(default=1, ge=1, le=2_147_483_647)
    request_id: Optional[str] = Field(None, max_length=256)
    references: List[ChatEvidenceReference] = Field(default_factory=list, max_length=24)

    @field_validator('request_id')
    @classmethod
    def _strip_request_id(cls, value: Optional[str]) -> Optional[str]:
        if value is None:
            return None
        normalized = value.strip()
        return normalized or None

    @model_validator(mode='after')
    def _reject_duplicate_reference_ids(self) -> 'ChatEvidenceEnvelope':
        identities = [reference.id for reference in self.references]
        if len(identities) != len(set(identities)):
            raise ValueError('evidence reference ids must be unique')
        if self.schema_version != 1:
            self.references = [
                reference.model_copy(update={'kind': 'unknown', 'state': 'unknown'}) for reference in self.references
            ]
        return self


class Message(BaseModel):
    id: str
    text: str
    created_at: datetime
    sender: MessageSender
    app_id: Optional[str] = None
    # TODO: remove plugin_id after migration
    plugin_id: Optional[str] = None
    from_external_integration: bool = False
    type: MessageType
    memories_id: List[str] = []  # used in db
    memories: List[MessageConversation] = []  # used front facing
    reported: bool = False
    report_reason: Optional[str] = None
    files_id: List[str] = []
    files: List[FileChat] = []
    chat_session_id: Optional[str] = None
    session_id: Optional[str] = None
    data_protection_level: Optional[str] = None
    langsmith_run_id: Optional[str] = None  # LangSmith run ID for feedback tracking
    prompt_name: Optional[str] = None  # LangSmith prompt name for versioning
    prompt_commit: Optional[str] = None  # LangSmith prompt commit/version for traceability
    rating: Optional[int] = None  # User feedback: 1 = thumbs up, -1 = thumbs down, None = no rating
    # Desktop journal compatibility fields. These are optional so the existing
    # message response remains readable by older clients while a new client can
    # reconcile the canonical turn identity and structured payload exactly.
    metadata: Optional[str] = None
    content_blocks: List[dict[str, Any]] = Field(
        default_factory=list,
        description=(
            'Structured chat content blocks. New rows store these directly; '
            'legacy rows are projected from metadata.content_blocks.'
        ),
    )
    evidence: Optional[ChatEvidenceEnvelope] = None
    client_message_id: Optional[str] = None
    message_source: Optional[str] = None
    journal_revision: Optional[int] = None
    chart_data: Optional[Union[ChartData, dict]] = None  # Inline chart visualization data

    @model_validator(mode='before')
    @classmethod
    def _sync_app_and_plugin_ids(cls, data: Any) -> Any:
        if isinstance(data, dict):
            app_id_val = data.get('app_id')
            plugin_id_val = data.get('plugin_id')

            if app_id_val is not None:
                data['plugin_id'] = app_id_val
            elif plugin_id_val is not None:
                data['app_id'] = plugin_id_val

            if 'content_blocks' not in data:
                metadata = data.get('metadata')
                if isinstance(metadata, str):
                    try:
                        legacy_blocks = json.loads(metadata).get('content_blocks')
                    except (AttributeError, TypeError, ValueError):
                        legacy_blocks = None
                    if isinstance(legacy_blocks, list):
                        data['content_blocks'] = legacy_blocks
        return data

    @classmethod
    def deserialize_many_safe(cls, records, on_error=None) -> List['Message']:
        """Build Message objects from raw stored records, skipping any that fail
        validation so one malformed or legacy chat message cannot 500 a whole history
        load. on_error(record, exception), when provided, is called for each skip."""
        parsed: List['Message'] = []
        for record in records:
            try:
                parsed.append(cls(**record))
            except Exception as exc:  # noqa: BLE001 - one bad record must not break the history
                if on_error is not None:
                    on_error(record, exc)
        return parsed

    @staticmethod
    def _resolve_sender_name(
        message: 'Message',
        *,
        use_plugin_name_if_available: bool,
        app_name_by_id: Dict[str, Optional[str]],
        app_name_resolver: Optional[AppNameResolver],
    ) -> str:
        if message.sender == 'human':
            return 'User'
        if use_plugin_name_if_available and app_name_resolver and message.app_id and message.app_id.strip():
            app_id = message.app_id.strip()
            if app_id not in app_name_by_id:
                name = app_name_resolver(app_id)
                app_name_by_id[app_id] = name.strip() if isinstance(name, str) and name.strip() else None
            resolved_name = app_name_by_id[app_id]
            if resolved_name:
                return resolved_name
        return message.sender.upper()

    @staticmethod
    def get_messages_as_string(
        messages: List['Message'],
        use_user_name_if_available: bool = False,
        use_plugin_name_if_available: bool = False,
        include_file_info: bool = False,
        app_name_resolver: Optional[AppNameResolver] = None,
    ) -> str:
        sorted_messages = sorted(messages, key=lambda m: m.created_at)
        app_name_by_id: Dict[str, Optional[str]] = {}

        formatted_messages = []
        for message in sorted_messages:
            sender_name = Message._resolve_sender_name(
                message,
                use_plugin_name_if_available=use_plugin_name_if_available,
                app_name_by_id=app_name_by_id,
                app_name_resolver=app_name_resolver,
            )
            msg_text = f"({message.created_at.strftime('%d %b %Y at %H:%M UTC')}) {sender_name}: {message.text}"

            # Add file info if requested and files exist
            if include_file_info and message.files_id and len(message.files_id) > 0:
                file_info = f" [Files attached: {len(message.files_id)} file(s), IDs: {', '.join(message.files_id)}]"
                msg_text += file_info

            formatted_messages.append(msg_text)

        return '\n'.join(formatted_messages)

    @staticmethod
    def get_messages_as_xml(
        messages: List['Message'],
        use_user_name_if_available: bool = False,
        use_plugin_name_if_available: bool = False,
        include_file_info: bool = False,
        app_name_resolver: Optional[AppNameResolver] = None,
    ) -> str:
        sorted_messages = sorted(messages, key=lambda m: m.created_at)
        app_name_by_id: Dict[str, Optional[str]] = {}

        formatted_messages = []
        for message in sorted_messages:
            # Build file section if requested
            file_section = ""
            if include_file_info and message.files and len(message.files) > 0:
                file_section = '<attachments>\n'
                for file in message.files:
                    file_section += f'  <file id="{file.id}" name="{file.name}" type="{file.mime_type}"/>\n'
                file_section += '</attachments>'
            elif include_file_info and message.files_id and len(message.files_id) > 0:
                # Fallback if files not loaded but IDs exist
                file_section = '<attachments>\n'
                for file_id in message.files_id:
                    file_section += f'  <file id="{file_id}"/>\n'
                file_section += '</attachments>'
            elif message.files and len(message.files) > 0:
                # Original behavior when include_file_info is False
                file_section = (
                    '<attachments>' + ''.join(f"<file>{file.name}</file>" for file in message.files) + '</attachments>'
                )

            msg = f"""<message>
<created_at>{message.created_at.strftime('%d %b %Y at %H:%M UTC')}</created_at>
<sender>{Message._resolve_sender_name(message, use_plugin_name_if_available=use_plugin_name_if_available, app_name_by_id=app_name_by_id, app_name_resolver=app_name_resolver)}</sender>
<content>{message.text}</content>
{file_section}
</message>"""

            # Only strip the block's surrounding whitespace. The template above is flush-left, so a
            # .replace('    ', '') here would instead delete 4-space runs from message.text (code,
            # tables, aligned or pasted text), corrupting the history shown to the LLM.
            formatted_messages.append(msg.strip())

        return '\n'.join(formatted_messages)


class ResponseMessage(Message):
    ask_for_nps: Optional[bool] = False


class PageContext(BaseModel):
    """Page context for chat - indicates what the user is currently viewing.

    When ``type`` is ``conversation`` with an ``id``, and/or ``start_date`` /
    ``end_date`` are set, retrieval tools hard-scope to that conversation and/or
    timeframe (#4515). Dates must include a timezone offset
    (YYYY-MM-DDTHH:MM:SS+HH:MM).
    """

    type: Literal["conversation", "task", "memory", "recap"]
    id: Optional[str] = None
    title: Optional[str] = None
    start_date: Optional[str] = None
    end_date: Optional[str] = None

    @staticmethod
    def _require_aware_iso(value: str, field_name: str) -> Optional[str]:
        stripped = value.strip()
        if not stripped:
            return None
        try:
            parsed = datetime.fromisoformat(stripped.replace("Z", "+00:00"))
        except ValueError as exc:
            raise ValueError(
                f"{field_name} must be ISO-8601 with timezone " f"(YYYY-MM-DDTHH:MM:SS+HH:MM), got {value!r}"
            ) from exc
        if parsed.tzinfo is None:
            raise ValueError(
                f"{field_name} must include a timezone offset " f"(YYYY-MM-DDTHH:MM:SS+HH:MM), got {value!r}"
            )
        return stripped

    @field_validator("start_date", "end_date", mode="before")
    @classmethod
    def _validate_scope_dates(cls, value: Any, info: ValidationInfo) -> Optional[str]:
        if value is None:
            return None
        if not isinstance(value, str):
            raise ValueError(f"{info.field_name or 'date'} must be an ISO-8601 string with timezone offset")
        field_name = info.field_name if isinstance(info.field_name, str) else "date"
        return cls._require_aware_iso(value, field_name)


class SendMessageRequest(BaseModel):
    text: str
    file_ids: Optional[List[str]] = []
    context: Optional[PageContext] = None


class GenerateReplyTurn(BaseModel):
    """A prior turn supplied by the caller purely as generation context."""

    text: str = Field(..., min_length=1, max_length=100000)
    sender: MessageSender


class GenerateReplyRequest(BaseModel):
    text: str = Field(..., min_length=1, max_length=100000)
    history: List[GenerateReplyTurn] = Field(default_factory=list, max_length=50)
    app_id: Optional[str] = Field(None, max_length=200)


class GenerateReplyResponse(BaseModel):
    text: str
    app_id: Optional[str] = None


class RateMessageRequest(BaseModel):
    rating: Optional[int] = None
    # Optional so existing clients keep working; a thumbs-down that arrives
    # without a reason is recorded as "not captured", never as "no reason".
    reason: Optional[FeedbackReason] = None
    comment: Optional[str] = Field(None, max_length=MAX_COMMENT_LENGTH)


class ShareChatMessagesRequest(BaseModel):
    message_ids: list[str] = []


class ChatSession(BaseModel):
    id: str
    message_ids: Optional[List[str]] = []
    file_ids: Optional[List[str]] = []
    app_id: Optional[str] = None
    plugin_id: Optional[str] = None
    created_at: datetime
    # Legacy Assistants IDs remain readable on old session docs; nothing writes them.
    openai_thread_id: Optional[str] = None
    openai_assistant_id: Optional[str] = None

    @model_validator(mode='before')
    @classmethod
    def _sync_chat_session_app_and_plugin_ids(cls, data: Any) -> Any:
        if isinstance(data, dict):
            app_id_val = data.get('app_id')
            plugin_id_val = data.get('plugin_id')

            if app_id_val is not None:
                data['plugin_id'] = app_id_val
            elif plugin_id_val is not None:
                data['app_id'] = plugin_id_val
        return data

    def add_file_ids(self, new_file_ids: List[str]):
        if self.file_ids is None:
            self.file_ids = []
        for file_id in new_file_ids:
            if file_id not in self.file_ids:
                self.file_ids.append(file_id)

    def retrieve_new_file(self, file_ids) -> List:
        existing_files = set(self.file_ids or [])
        return list(set(file_ids) - existing_files)