wsb-daemon 0.258.6441

use wsb::persist::global::models::{DBMessage, DBSession, MinimapEntry, SessionFilter, Stats};
use wsb::prelude::*;
use crate::web::machine_web::{place_filter, MachineWeb};
use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::sse::{Event, KeepAlive, Sse},
    Json,
};
use std::convert::Infallible;

#[derive(Deserialize)]
pub struct ListSessionsQuery {
    pub limit: Option<i64>,
    pub cursor: Option<String>,
    pub project: Option<String>,
    pub worktree: Option<String>,
    pub placed: Option<String>,
    pub machine: Option<String>,
    pub agent: Option<String>,
    pub date: Option<String>,
    pub date_from: Option<String>,
    pub date_to: Option<String>,
    pub active_since: Option<String>,
    pub min_messages: Option<i64>,
    pub max_messages: Option<i64>,
    pub min_user_messages: Option<i64>,
    pub include_one_shot: Option<String>,
    pub include_children: Option<String>,
}

#[derive(Serialize)]
pub struct SessionListResponse {
    pub sessions: Vec<DBSession>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub next_cursor: Option<String>,
    pub total: usize,
}

pub async fn list_sessions(
    State(state): State<MachineWeb>,
    Query(params): Query<ListSessionsQuery>,
) -> Result<Json<SessionListResponse>, (StatusCode, Json<ErrorResponse>)> {
    let limit = params.limit.unwrap_or(200).min(500);
    let filter = SessionFilter {
        place: place_filter(params.project, params.worktree, params.placed)?,
        machine: params
            .machine
            .map(|m| m.split(',').map(|s| s.to_string()).collect()),
        agent: params
            .agent
            .map(|a| a.split(',').map(|s| s.to_string()).collect()),
        date: params.date,
        date_from: params.date_from,
        date_to: params.date_to,
        active_since: params.active_since,
        min_messages: params.min_messages,
        max_messages: params.max_messages,
        min_user_messages: params.min_user_messages,
        include_one_shot: params.include_one_shot.as_deref() == Some("true"),
        include_children: params.include_children.as_deref() == Some("true"),
    };

    let cursor = params.cursor.as_ref().and_then(|c| {
        let parts: Vec<&str> = c.splitn(2, '|').collect();
        if parts.len() == 2 {
            Some((parts[0].to_string(), parts[1].to_string()))
        } else {
            None
        }
    });

    let db = state.db.clone();

    let sessions = tokio::task::spawn_blocking(move || {
        db.list_sessions(&filter, limit + 1, cursor)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let has_more = sessions.len() as i64 > limit;
    let sessions: Vec<_> = sessions.into_iter().take(limit as usize).collect();
    let total = sessions.len();

    let next_cursor = if has_more {
        sessions
            .last()
            .map(|s| format!("{}|{}", s.ended_at.as_deref().unwrap_or(""), s.id))
    } else {
        None
    };

    Ok(Json(SessionListResponse {
        sessions,
        next_cursor,
        total,
    }))
}

pub async fn get_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();
    let id_check = id.clone();

    let session = tokio::task::spawn_blocking(move || db.get_session(&id_check))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    match session {
        Some(s) => {
            let is_starred = tokio::task::spawn_blocking({
                let db = state.db.clone();
                move || db.is_starred(&id)
            })
            .await
            .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
            .unwrap_or(false);

            let mut val = serde_json::to_value(&s).unwrap();
            val.as_object_mut()
                .unwrap()
                .insert("is_starred".to_string(), serde_json::json!(is_starred));
            Ok(Json(val))
        }
        None => Err(not_found("session not found")),
    }
}

#[derive(Deserialize)]
pub struct MessagesQuery {
    pub limit: Option<i64>,
    pub from: Option<i64>,
    pub direction: Option<String>,
}

#[derive(Serialize)]
pub struct MessagesResponse {
    pub messages: Vec<DBMessage>,
    pub count: usize,
}

pub async fn get_messages(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
    Query(params): Query<MessagesQuery>,
) -> Result<Json<MessagesResponse>, (StatusCode, Json<ErrorResponse>)> {
    let limit = params.limit.unwrap_or(100).min(1000);
    let desc = params.direction.as_deref() == Some("desc");
    let from = params.from.unwrap_or(if desc { i64::MAX } else { 0 });

    let db = state.db.clone();
    let id_clone = id.clone();

    let messages = tokio::task::spawn_blocking(move || {
        db.get_messages(&id_clone, limit, from, desc)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let count = messages.len();
    Ok(Json(MessagesResponse { messages, count }))
}

pub async fn get_children(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<Json<Vec<DBSession>>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();

    let children = tokio::task::spawn_blocking(move || db.get_children(&id))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(children))
}

#[derive(Deserialize)]
pub struct MinimapQuery {
    pub from: Option<i64>,
    pub max: Option<i64>,
}

#[derive(Serialize)]
pub struct MinimapResponse {
    pub entries: Vec<MinimapEntry>,
    pub count: usize,
}

pub async fn get_minimap(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
    Query(params): Query<MinimapQuery>,
) -> Result<Json<MinimapResponse>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();

    let entries = tokio::task::spawn_blocking(move || {
        db.get_minimap(&id, params.from.unwrap_or(0), params.max)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let count = entries.len();
    Ok(Json(MinimapResponse { entries, count }))
}

#[derive(Deserialize)]
pub struct SessionSearchQuery {
    pub q: String,
}

pub async fn search_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
    Query(params): Query<SessionSearchQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();
    let id_clone = id.clone();
    let q = params.q.clone();

    let ordinals =
        tokio::task::spawn_blocking(move || db.search_session(&id_clone, &q))
            .await
            .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
            .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({ "ordinals": ordinals })))
}

pub async fn watch_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Sse<impl futures::Stream<Item = Result<Event, Infallible>>> {
    let stream = async_stream::stream! {
        let mut last_count: i64 = 0;
        let mut last_mtime: i64 = 0;
        let mut heartbeat_counter: u32 = 0;
        let db = state.db.clone();
        let id_clone = id.clone();

        if let Ok(Some((count, mtime))) = tokio::task::spawn_blocking(move || {
            db.get_session_version(&id_clone)
        }).await.unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e))) {
            last_count = count;
            last_mtime = mtime;
        }

        loop {
            tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
            heartbeat_counter += 1;

            let db = state.db.clone();
            let id_clone = id.clone();

            match tokio::task::spawn_blocking(move || {
                db.get_session_version(&id_clone)
            }).await {
                Ok(Ok(Some((count, mtime)))) => {
                    if count != last_count || mtime != last_mtime {
                        last_count = count;
                        last_mtime = mtime;
                        yield Ok(Event::default()
                            .event("session_updated")
                            .data(id.clone()));
                    }
                }
                Ok(Ok(None)) => {
                    break;
                }
                _ => {}
            }

            if heartbeat_counter >= 20 {
                heartbeat_counter = 0;
                yield Ok(Event::default()
                    .event("heartbeat")
                    .data(chrono::Utc::now().to_rfc3339()));
            }
        }
    };

    Sse::new(stream).keep_alive(KeepAlive::default())
}

#[derive(Deserialize)]
pub struct RenameRequest {
    pub display_name: Option<String>,
}

pub async fn rename_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
    Json(body): Json<RenameRequest>,
) -> Result<Json<DBSession>, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();
    let db_get = state.db.clone();
    let id_get = id.clone();

    tokio::task::spawn_blocking(move || {
        db.rename_session(&id, body.display_name.as_deref())
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let session = tokio::task::spawn_blocking(move || db_get.get_session(&id_get))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    match session {
        Some(s) => Ok(Json(s)),
        None => Err(not_found("session not found")),
    }
}

pub async fn delete_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<StatusCode, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();

    let deleted = tokio::task::spawn_blocking(move || db.soft_delete_session(&id))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    if deleted {
        Ok(StatusCode::NO_CONTENT)
    } else {
        Err(not_found("session not found"))
    }
}

pub async fn restore_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<StatusCode, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();

    let restored = tokio::task::spawn_blocking(move || db.restore_session(&id))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    if restored {
        Ok(StatusCode::NO_CONTENT)
    } else {
        Err(not_found("session not found or not in trash"))
    }
}

pub async fn permanent_delete(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<StatusCode, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();

    let deleted = tokio::task::spawn_blocking(move || db.permanent_delete(&id))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    if deleted {
        Ok(StatusCode::NO_CONTENT)
    } else {
        Err((
            StatusCode::CONFLICT,
            Json(ErrorResponse {
                error: "session not found or not in trash".to_string(),
            }),
        ))
    }
}

#[derive(Deserialize)]
pub struct OneShot {
    pub include_one_shot: Option<String>,
}

/// How many sessions ran in each project and worktree, and outside `$PRJ`.
pub async fn list_places(
    State(state): State<MachineWeb>,
    Query(params): Query<OneShot>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let include = params.include_one_shot.as_deref() == Some("true");
    let db = state.db.clone();

    let places = tokio::task::spawn_blocking(move || db.list_places(include))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!(places)))
}

#[derive(Deserialize)]
pub struct StatsQuery {
    pub include_one_shot: Option<String>,
    pub project: Option<String>,
    pub worktree: Option<String>,
    pub placed: Option<String>,
}

pub async fn list_machines(
    State(state): State<MachineWeb>,
    Query(params): Query<OneShot>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let include = params.include_one_shot.as_deref() == Some("true");
    let db = state.db.clone();

    let machines = tokio::task::spawn_blocking(move || db.list_machines(include))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({ "machines": machines })))
}

pub async fn list_agents(
    State(state): State<MachineWeb>,
    Query(params): Query<OneShot>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let include = params.include_one_shot.as_deref() == Some("true");
    let db = state.db.clone();

    let agents = tokio::task::spawn_blocking(move || db.list_agents(include))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(serde_json::json!({ "agents": agents })))
}

pub async fn get_stats(
    State(state): State<MachineWeb>,
    Query(params): Query<StatsQuery>,
) -> Result<Json<Stats>, (StatusCode, Json<ErrorResponse>)> {
    let include = params.include_one_shot.as_deref() == Some("true");
    let place = place_filter(params.project, params.worktree, params.placed)?;
    let db = state.db.clone();

    let stats = tokio::task::spawn_blocking(move || db.get_stats(include, &place))
        .await
        .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
        .map_err(|e| internal_error(&e.to_string()))?;

    Ok(Json(stats))
}

#[derive(Serialize, Debug)]
pub struct ErrorResponse {
    pub error: String,
}

fn internal_error(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(ErrorResponse {
            error: msg.to_string(),
        }),
    )
}

#[derive(Deserialize)]
pub struct CompactQuery {
    pub token_mode: Option<String>,
}

pub async fn get_compact_text(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
    Query(params): Query<CompactQuery>,
) -> Result<Json<serde_json::Value>, (StatusCode, Json<ErrorResponse>)> {
    let token_mode = params.token_mode.as_deref() == Some("true");
    let db = state.db.clone();
    let id_clone = id.clone();

    let messages = tokio::task::spawn_blocking(move || {
        db.get_messages(&id_clone, i64::MAX, 0, false)
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let text = messages
        .iter()
        .filter(|m| !m.content.trim().is_empty())
        .map(|m| {
            let role = match m.role.as_str() {
                "user" => "USER:",
                "assistant" => "ASSIST:",
                other => other,
            };
            format!("{}\n{}", role, m.content.trim())
        })
        .collect::<Vec<_>>()
        .join("\n\n");

    let compacted = wsb::compact::compact_text(&text, token_mode);

    Ok(Json(serde_json::json!({
        "text": compacted,
    })))
}

pub async fn export_session(
    State(state): State<MachineWeb>,
    Path(id): Path<String>,
) -> Result<axum::response::Response, (StatusCode, Json<ErrorResponse>)> {
    let db = state.db.clone();
    let id_for_messages = id.clone();
    let id_for_body = id.clone();

    let (session, messages) = tokio::task::spawn_blocking(move || {
        let session = db.get_session(&id_for_messages)?;
        let msgs = db
            .get_messages(&id_for_messages, i64::MAX, 0, false)?;
        anyhow::Ok((session, msgs))
    })
    .await
    .unwrap_or_else(|e| Err(anyhow::anyhow!("join error: {}", e)))
    .map_err(|e| internal_error(&e.to_string()))?;

    let session = session.ok_or_else(|| not_found("Session not found"))?;

    let title = session
        .display_name
        .clone()
        .or_else(|| {
            session.project.clone().map(|project| match &session.worktree {
                Some(worktree) => format!("{project}/{worktree}"),
                None => project,
            })
        })
        .unwrap_or_else(|| id_for_body.clone());
    let title = title.as_str();

    let mut body = format!("# {}\n\n", title);
    for m in &messages {
        let role = match m.role.as_str() {
            "user" => "User",
            "assistant" => "Assistant",
            other => other,
        };
        body.push_str(&format!("## {}\n\n{}\n\n", role, m.content));
    }

    let filename = format!(
        "{}-{}.md",
        title.replace(' ', "-"),
        &id_for_body[..8.min(id_for_body.len())]
    );

    Ok(axum::response::IntoResponse::into_response((
        [
            ("content-type", "text/markdown; charset=utf-8"),
            (
                "content-disposition",
                format!("attachment; filename=\"{}\"", filename).as_str(),
            ),
        ],
        body,
    )))
}

fn not_found(msg: &str) -> (StatusCode, Json<ErrorResponse>) {
    (
        StatusCode::NOT_FOUND,
        Json(ErrorResponse {
            error: msg.to_string(),
        }),
    )
}