squigit-storage 0.1.0

// Copyright 2026 a7mddra
// SPDX-License-Identifier: Apache-2.0

//! Type definitions for thread storage.

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use uuid::Uuid;

pub const EMPTY_STATE_ASSET_ID: &str = "__empty_state_asset__";
pub const DEFAULT_THREAD_TITLE: &str = "New thread";
pub const DEFAULT_SIDE_CHAT_TITLE: &str = "New side chat";

fn new_thread_id() -> String {
    let now = Utc::now();
    let date_part = now.format("%Y%m%d-%H%M%S").to_string();
    let uuid_part = Uuid::new_v4().to_string();
    format!("{}-{}", date_part, &uuid_part[..8])
}

/// Metadata for a thread session stored in the thread index.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThreadMetadata {
    /// Unique identifier for the thread.
    pub id: String,
    /// Display title for the thread.
    pub title: String,
    /// When the thread was created.
    pub created_at: DateTime<Utc>,
    /// When the thread was last updated.
    pub updated_at: DateTime<Utc>,
    /// BLAKE3 hash of the associated image.
    pub image_hash: String,
    /// When the thread was pinned, or `None` when it is not pinned.
    pub pinned_at: Option<DateTime<Utc>>,
}

impl ThreadMetadata {
    /// Create new thread metadata with a generated ID.
    pub fn new(title: String, image_hash: String) -> Self {
        let now = Utc::now();

        Self {
            id: new_thread_id(),
            title,
            created_at: now,
            updated_at: now,
            image_hash,
            pinned_at: None,
        }
    }
}

/// Metadata for a SideChat thread stored outside Explorer workspaces.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SideChatMetadata {
    pub id: String,
    pub title: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

impl SideChatMetadata {
    pub fn new(title: String) -> Self {
        let now = Utc::now();
        Self {
            id: new_thread_id(),
            title,
            created_at: now,
            updated_at: now,
        }
    }
}

/// A workspace groups threads and the directories its AI sandbox can read.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkspaceMetadata {
    /// Unique identifier for the workspace.
    pub id: String,
    /// Workspace name displayed in the sidebar.
    pub name: String,
    /// When the workspace was created.
    pub created_at: DateTime<Utc>,
    /// Directories available to threads in this workspace.
    pub directories: Vec<String>,
    /// Thread metadata keyed by thread ID.
    pub threads: BTreeMap<String, ThreadMetadata>,
}

impl WorkspaceMetadata {
    /// Create a workspace with a generated ID.
    pub fn new(name: String, directories: Vec<String>) -> Self {
        Self {
            id: format!("workspace-{}", Uuid::new_v4()),
            name,
            created_at: Utc::now(),
            directories,
            threads: BTreeMap::new(),
        }
    }
}

/// A CAS object referenced by one user message.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MessageAttachment {
    pub attachment_hash: String,
    pub source_path: Option<String>,
}

/// A persisted message with a strict role-specific JSON shape.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "lowercase")]
pub enum ThreadMessage {
    User {
        id: String,
        content: String,
        timestamp: DateTime<Utc>,
        attachments: Vec<MessageAttachment>,
    },
    Assistant {
        id: String,
        content: String,
        timestamp: DateTime<Utc>,
        citations: Vec<CitationSource>,
        tool_steps: Vec<ToolStep>,
    },
}

/// Structured citation source metadata persisted with a message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CitationSource {
    pub title: String,
    pub url: String,
    pub summary: String,
    #[serde(default)]
    pub favicon_url: Option<String>,
    #[serde(default)]
    pub favicon_base64: Option<String>,
}

/// Tool-step metadata persisted with a message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolStep {
    pub id: String,
    pub name: String,
    pub status: String,
    #[serde(default)]
    pub args: serde_json::Value,
    #[serde(default)]
    pub message: Option<String>,
    #[serde(default, rename = "startedAtMs")]
    pub started_at_ms: Option<u64>,
    #[serde(default, rename = "endedAtMs")]
    pub ended_at_ms: Option<u64>,
}

impl ThreadMessage {
    fn new_id() -> String {
        format!("msg-{}", Uuid::new_v4())
    }

    pub fn is_valid_id(id: &str) -> bool {
        id.strip_prefix("msg-")
            .is_some_and(|uuid| Uuid::parse_str(uuid).is_ok())
    }

    /// Create a new user message with structured CAS attachment hashes.
    pub fn user_with_attachments(content: String, attachments: Vec<MessageAttachment>) -> Self {
        Self::User {
            id: Self::new_id(),
            content,
            timestamp: Utc::now(),
            attachments,
        }
    }

    pub fn id(&self) -> &str {
        match self {
            Self::User { id, .. } | Self::Assistant { id, .. } => id,
        }
    }

    pub fn content(&self) -> &str {
        match self {
            Self::User { content, .. } | Self::Assistant { content, .. } => content,
        }
    }

    pub fn attachments(&self) -> &[MessageAttachment] {
        match self {
            Self::User { attachments, .. } => attachments,
            Self::Assistant { .. } => &[],
        }
    }
}

/// OCR data for an image region.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrRegion {
    /// Extracted text.
    pub text: String,
    /// Bounding box coordinates.
    #[serde(default)]
    pub bbox: Vec<Vec<i32>>,
}

/// OCR output for a single model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrModelAnnotation {
    /// When this OCR model last finished scanning this thread image.
    #[serde(default)]
    pub scanned_at: Option<DateTime<Utc>>,
    /// Cached OCR results for this model.
    #[serde(default)]
    pub ocr_data: Vec<OcrRegion>,
}

/// OCR annotations entry.
///
/// The empty-state sentinel is stored as an empty array while real model IDs
/// store timestamped OCR data.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum OcrAnnotationEntry {
    EmptyState(Vec<OcrRegion>),
    Model(OcrModelAnnotation),
}

/// OCR annotations keyed by sentinel/model ID.
pub type OcrAnnotations = HashMap<String, OcrAnnotationEntry>;

pub fn default_ocr_annotations() -> OcrAnnotations {
    HashMap::from([(
        EMPTY_STATE_ASSET_ID.to_string(),
        OcrAnnotationEntry::EmptyState(Vec::new()),
    )])
}

/// LLM context window state for a thread.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ContextWindow {
    pub tokens_used: u32,
    #[serde(default)]
    pub compacted_at: Option<DateTime<Utc>>,
    #[serde(default)]
    pub compacted_context: Option<String>,
}

/// Per-thread model context for a CAS attachment.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct AttachmentManifestEntry {
    pub attachment_hash: String,
    pub display_name: String,
    pub file_type: crate::cas::AttachmentFileType,
    pub file_brief: Option<String>,
    pub last_mention_at: DateTime<Utc>,
}

pub type AttachmentManifest = Vec<AttachmentManifestEntry>;

/// Complete thread data including messages and OCR.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThreadData {
    /// Thread metadata.
    pub metadata: ThreadMetadata,
    /// Thread messages.
    #[serde(default)]
    pub messages: Vec<ThreadMessage>,
    /// OCR annotations keyed by sentinel/model ID.
    #[serde(default = "default_ocr_annotations")]
    pub ocr_data: OcrAnnotations,
    /// LLM context window state.
    #[serde(default)]
    pub context_window: ContextWindow,
    /// Reverse image search cache for the core thread image.
    #[serde(default)]
    pub reverse_image_search: Option<crate::cas::ReverseImageSearchCache>,
    /// Per-thread attachment context persisted in attachment_manifest.json.
    pub attachment_manifest: AttachmentManifest,
    /// Image tone resolved from the initial object's manifest.
    pub image_tone: Option<String>,
}

impl ThreadData {
    /// Create new thread data with metadata.
    pub fn new(metadata: ThreadMetadata, initial_attachment: AttachmentManifestEntry) -> Self {
        Self {
            metadata,
            messages: Vec::new(),
            ocr_data: default_ocr_annotations(),
            context_window: ContextWindow::default(),
            reverse_image_search: None,
            attachment_manifest: vec![initial_attachment],
            image_tone: None,
        }
    }
}

/// A SideChat thread uses the common message/context/attachment files without OCR state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SideChatData {
    pub metadata: SideChatMetadata,
    #[serde(default)]
    pub messages: Vec<ThreadMessage>,
    #[serde(default)]
    pub context_window: ContextWindow,
    pub attachment_manifest: AttachmentManifest,
}

impl SideChatData {
    pub fn new(metadata: SideChatMetadata, first_message: ThreadMessage) -> Self {
        Self {
            metadata,
            messages: vec![first_message],
            context_window: ContextWindow::default(),
            attachment_manifest: Vec::new(),
        }
    }
}