lakehouse/crates/queryd/src/workspace.rs

/// Agent workspaces — named overlays for contract/search-specific work.
/// Each workspace tracks an agent's activity on a specific contract or search,
/// with daily/weekly/monthly tiers and instant handoff capability.

use arrow::array::{ArrayRef, RecordBatch, StringArray, Int64Array};
use arrow::datatypes::{DataType, Field, Schema};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

use crate::delta;
use object_store::ObjectStore;

/// Retention tier for workspace data.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Tier {
    Daily,   // expires end of day, active search scratch
    Weekly,  // expires end of week, active contract work
    Monthly, // expires end of month, contract lifecycle
    Pinned,  // never expires, manually managed
}

/// A saved query/filter within a workspace.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SavedSearch {
    pub name: String,
    pub sql: String,
    pub created_at: DateTime<Utc>,
}

/// A shortlisted candidate or record.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShortlistEntry {
    pub dataset: String,
    pub record_id: String,
    pub notes: String,
    pub added_at: DateTime<Utc>,
    pub added_by: String,
}

/// Activity log entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ActivityEntry {
    pub action: String,      // "search", "shortlist", "call", "email", "update", "ingest"
    pub detail: String,
    pub timestamp: DateTime<Utc>,
    pub agent: String,
}

/// A workspace — an agent's working context for a contract or search.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Workspace {
    pub id: String,
    pub name: String,            // "Apex Corp .NET Developers - Chicago"
    pub description: String,
    pub tier: Tier,
    pub owner: String,           // current agent
    pub previous_owners: Vec<HandoffRecord>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,

    // Work content
    pub saved_searches: Vec<SavedSearch>,
    pub shortlist: Vec<ShortlistEntry>,
    pub activity: Vec<ActivityEntry>,
    pub ingested_datasets: Vec<String>,  // datasets this workspace created
    pub delta_keys: Vec<String>,         // delta files specific to this workspace
    pub tags: Vec<String>,
}

/// Record of a handoff between agents.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HandoffRecord {
    pub from_agent: String,
    pub to_agent: String,
    pub reason: String,
    pub timestamp: DateTime<Utc>,
}

/// Workspace manager — in-memory registry with persistence.
#[derive(Clone)]
pub struct WorkspaceManager {
    workspaces: Arc<RwLock<HashMap<String, Workspace>>>,
    store: Arc<dyn ObjectStore>,
}

impl WorkspaceManager {
    pub fn new(store: Arc<dyn ObjectStore>) -> Self {
        Self {
            workspaces: Arc::new(RwLock::new(HashMap::new())),
            store,
        }
    }

    /// Rebuild from persisted workspace files on startup.
    pub async fn rebuild(&self) -> Result<usize, String> {
        let keys = storaged::ops::list(&self.store, Some("workspaces/")).await?;
        let mut ws_map = self.workspaces.write().await;
        ws_map.clear();

        for key in &keys {
            if !key.ends_with(".json") { continue; }
            let data = storaged::ops::get(&self.store, key).await?;
            match serde_json::from_slice::<Workspace>(&data) {
                Ok(ws) => { ws_map.insert(ws.id.clone(), ws); }
                Err(e) => tracing::warn!("failed to load workspace {key}: {e}"),
            }
        }

        let count = ws_map.len();
        if count > 0 {
            tracing::info!("loaded {count} workspaces");
        }
        Ok(count)
    }

    /// Create a new workspace.
    pub async fn create(&self, name: String, description: String, owner: String, tier: Tier) -> Result<Workspace, String> {
        let now = Utc::now();
        let id = format!("ws-{}", now.timestamp_millis());

        let ws = Workspace {
            id: id.clone(),
            name,
            description,
            tier,
            owner,
            previous_owners: vec![],
            created_at: now,
            updated_at: now,
            saved_searches: vec![],
            shortlist: vec![],
            activity: vec![],
            ingested_datasets: vec![],
            delta_keys: vec![],
            tags: vec![],
        };

        self.persist(&ws).await?;
        self.workspaces.write().await.insert(id.clone(), ws.clone());
        tracing::info!("created workspace: {} ({})", ws.name, ws.id);
        Ok(ws)
    }

    /// Handoff workspace to another agent. Instant — no data copy.
    pub async fn handoff(&self, workspace_id: &str, to_agent: &str, reason: &str) -> Result<Workspace, String> {
        let mut ws_map = self.workspaces.write().await;
        let ws = ws_map.get_mut(workspace_id)
            .ok_or_else(|| format!("workspace not found: {workspace_id}"))?;

        let record = HandoffRecord {
            from_agent: ws.owner.clone(),
            to_agent: to_agent.to_string(),
            reason: reason.to_string(),
            timestamp: Utc::now(),
        };

        ws.previous_owners.push(record);
        ws.owner = to_agent.to_string();
        ws.updated_at = Utc::now();

        ws.activity.push(ActivityEntry {
            action: "handoff".to_string(),
            detail: format!("handed off to {} — {}", to_agent, reason),
            timestamp: Utc::now(),
            agent: to_agent.to_string(),
        });

        let ws_clone = ws.clone();
        drop(ws_map);
        self.persist(&ws_clone).await?;

        tracing::info!("workspace '{}' handed off to {}", ws_clone.name, to_agent);
        Ok(ws_clone)
    }

    /// Add a saved search to a workspace.
    pub async fn add_search(&self, workspace_id: &str, name: String, sql: String, agent: &str) -> Result<(), String> {
        let mut ws_map = self.workspaces.write().await;
        let ws = ws_map.get_mut(workspace_id)
            .ok_or_else(|| format!("workspace not found: {workspace_id}"))?;

        ws.saved_searches.push(SavedSearch {
            name: name.clone(),
            sql,
            created_at: Utc::now(),
        });
        ws.activity.push(ActivityEntry {
            action: "search".into(),
            detail: format!("saved search: {name}"),
            timestamp: Utc::now(),
            agent: agent.to_string(),
        });
        ws.updated_at = Utc::now();

        let ws_clone = ws.clone();
        drop(ws_map);
        self.persist(&ws_clone).await
    }

    /// Add a candidate/record to the shortlist.
    pub async fn add_to_shortlist(&self, workspace_id: &str, dataset: String, record_id: String, notes: String, agent: &str) -> Result<(), String> {
        let mut ws_map = self.workspaces.write().await;
        let ws = ws_map.get_mut(workspace_id)
            .ok_or_else(|| format!("workspace not found: {workspace_id}"))?;

        ws.shortlist.push(ShortlistEntry {
            dataset: dataset.clone(),
            record_id: record_id.clone(),
            notes,
            added_at: Utc::now(),
            added_by: agent.to_string(),
        });
        ws.activity.push(ActivityEntry {
            action: "shortlist".into(),
            detail: format!("added {record_id} from {dataset}"),
            timestamp: Utc::now(),
            agent: agent.to_string(),
        });
        ws.updated_at = Utc::now();

        let ws_clone = ws.clone();
        drop(ws_map);
        self.persist(&ws_clone).await
    }

    /// Log an activity.
    pub async fn log_activity(&self, workspace_id: &str, action: String, detail: String, agent: &str) -> Result<(), String> {
        let mut ws_map = self.workspaces.write().await;
        let ws = ws_map.get_mut(workspace_id)
            .ok_or_else(|| format!("workspace not found: {workspace_id}"))?;

        ws.activity.push(ActivityEntry {
            action,
            detail,
            timestamp: Utc::now(),
            agent: agent.to_string(),
        });
        ws.updated_at = Utc::now();

        let ws_clone = ws.clone();
        drop(ws_map);
        self.persist(&ws_clone).await
    }

    /// Get a workspace.
    pub async fn get(&self, workspace_id: &str) -> Option<Workspace> {
        self.workspaces.read().await.get(workspace_id).cloned()
    }

    /// List all workspaces, optionally filtered by owner or tier.
    pub async fn list(&self, owner: Option<&str>, tier: Option<&Tier>) -> Vec<Workspace> {
        let ws_map = self.workspaces.read().await;
        ws_map.values()
            .filter(|ws| {
                owner.map_or(true, |o| ws.owner == o) &&
                tier.map_or(true, |t| ws.tier == *t)
            })
            .cloned()
            .collect()
    }

    /// Persist workspace to object storage.
    async fn persist(&self, ws: &Workspace) -> Result<(), String> {
        let key = format!("workspaces/{}.json", ws.id);
        let json = serde_json::to_vec_pretty(ws).map_err(|e| e.to_string())?;
        storaged::ops::put(&self.store, &key, json.into()).await
    }
}