use diesel::prelude::*;
use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
use std::collections::HashMap;
use std::fs;
use std::path::Path;

use crate::schema::{messages, preferences, projects, sessions};
use crate::settings::{
    DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, KvCachePreferences,
    ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences,
    StreamingCacheBudget,
};

pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = preferences)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct AppPreferences {
    pub id: i32,
    pub selected_model: String,
    pub dspark_enabled: bool,
    pub idle_timeout_minutes: i32,
    pub context_tokens: i32,
    pub max_generated_tokens: i32,
    pub system_prompt: String,
    pub temperature: Option<f32>,
    pub top_p: Option<f32>,
    pub min_p: Option<f32>,
    pub seed: Option<String>,
    pub reasoning_mode: String,
    pub cpu_threads: Option<i32>,
    pub power_percent: Option<i32>,
    pub prefill_chunk: Option<i32>,
    pub quality: bool,
    pub warm_weights: bool,
    pub mtp_draft_tokens: i32,
    pub mtp_margin: f32,
    pub glm_mtp: bool,
    pub glm_mtp_timing: bool,
    pub dspark_confidence_threshold: Option<f32>,
    pub dspark_strict: bool,
    pub ssd_streaming: bool,
    pub ssd_streaming_cold: bool,
    pub ssd_cache_experts: Option<i64>,
    pub ssd_cache_gib: Option<i64>,
    pub ssd_full_layers: Option<i32>,
    pub ssd_preload_experts: Option<i32>,
    pub directional_steering_file: Option<String>,
    pub directional_steering_ffn: Option<f32>,
    pub directional_steering_attn: Option<f32>,
    pub simulated_used_memory_gib: Option<i64>,
    pub expert_profile_path: Option<String>,
    pub endpoint_port: i32,
    pub endpoint_enabled: bool,
    pub endpoint_cors: bool,
    pub sidebar_collapsed: bool,
    /// Project the app reopens on. Cleared when that project goes away.
    pub last_project_id: Option<i32>,
    pub sidebar_width: i32,
    pub kv_budget_gib: Option<i64>,
    pub kv_min_tokens: Option<i32>,
    pub kv_cold_max_tokens: Option<i32>,
    pub kv_continued_interval_tokens: Option<i32>,
}

impl Default for AppPreferences {
    fn default() -> Self {
        Self {
            id: 1,
            selected_model: "deepseek-v4-flash".into(),
            dspark_enabled: false,
            idle_timeout_minutes: 10,
            context_tokens: 32_768,
            max_generated_tokens: 50_000,
            system_prompt: "You are a helpful assistant".into(),
            temperature: None,
            top_p: None,
            min_p: None,
            seed: None,
            reasoning_mode: "high".into(),
            cpu_threads: None,
            power_percent: None,
            prefill_chunk: None,
            quality: false,
            warm_weights: false,
            mtp_draft_tokens: 1,
            mtp_margin: 3.0,
            glm_mtp: false,
            glm_mtp_timing: false,
            dspark_confidence_threshold: None,
            dspark_strict: false,
            ssd_streaming: false,
            ssd_streaming_cold: false,
            ssd_cache_experts: None,
            ssd_cache_gib: None,
            ssd_full_layers: None,
            ssd_preload_experts: None,
            directional_steering_file: None,
            directional_steering_ffn: None,
            directional_steering_attn: None,
            simulated_used_memory_gib: None,
            expert_profile_path: None,
            endpoint_port: 4000,
            endpoint_enabled: true,
            endpoint_cors: false,
            sidebar_collapsed: false,
            last_project_id: None,
            sidebar_width: 276,
            kv_budget_gib: None,
            kv_min_tokens: None,
            kv_cold_max_tokens: None,
            kv_continued_interval_tokens: None,
        }
    }
}

impl AppPreferences {
    pub(crate) fn generation(&self) -> Result<GenerationPreferences, String> {
        let preferences = GenerationPreferences {
            context_tokens: self.context_tokens,
            max_generated_tokens: self.max_generated_tokens,
            system_prompt: self.system_prompt.clone(),
            temperature: self.temperature,
            top_p: self.top_p,
            min_p: self.min_p,
            seed: self
                .seed
                .as_deref()
                .map(str::parse)
                .transpose()
                .map_err(|_| "Saved seed is not a valid positive whole number.".to_owned())?,
            reasoning_mode: ReasoningMode::from_id(&self.reasoning_mode)
                .ok_or_else(|| format!("Unsupported reasoning mode: {}", self.reasoning_mode))?,
        };
        preferences.validate()?;
        Ok(preferences)
    }

    pub(crate) fn execution(&self) -> Result<ExecutionPreferences, String> {
        Ok(ExecutionPreferences {
            cpu_threads: self
                .cpu_threads
                .map(u32::try_from)
                .transpose()
                .map_err(|_| "Saved CPU helper threads is invalid.".to_owned())?,
            power_percent: self
                .power_percent
                .map(u8::try_from)
                .transpose()
                .map_err(|_| "Saved GPU power is invalid.".to_owned())?,
            prefill_chunk: self
                .prefill_chunk
                .map(u32::try_from)
                .transpose()
                .map_err(|_| "Saved prefill chunk is invalid.".to_owned())?,
            quality: self.quality,
            warm_weights: self.warm_weights,
        })
    }

    pub(crate) fn speculative(&self) -> SpeculativePreferences {
        SpeculativePreferences {
            mtp_draft_tokens: self.mtp_draft_tokens,
            mtp_margin: self.mtp_margin,
            glm_mtp: self.glm_mtp,
            glm_mtp_timing: self.glm_mtp_timing,
            dspark_enabled: self.dspark_enabled,
            dspark_confidence_threshold: self.dspark_confidence_threshold,
            dspark_strict: self.dspark_strict,
        }
    }

    pub(crate) fn runtime(&self) -> Result<RuntimePreferences, String> {
        let cache = match (self.ssd_cache_experts, self.ssd_cache_gib) {
            (None, None) => None,
            (Some(experts), None) => Some(StreamingCacheBudget::Experts(
                u32::try_from(experts).map_err(|_| "Saved SSD expert count is invalid.")?,
            )),
            (None, Some(gib)) => Some(StreamingCacheBudget::Gib(
                u64::try_from(gib).map_err(|_| "Saved SSD cache budget is invalid.")?,
            )),
            (Some(_), Some(_)) => {
                return Err("Saved SSD cache budget has both count and GiB values.".into());
            }
        };
        Ok(RuntimePreferences {
            execution: self.execution()?,
            speculative: self.speculative(),
            ssd: SsdPreferences {
                enabled: self.ssd_streaming,
                cold: self.ssd_streaming_cold,
                cache,
                full_layers: self
                    .ssd_full_layers
                    .map(u32::try_from)
                    .transpose()
                    .map_err(|_| "Saved SSD full-layer count is invalid.")?,
                preload_experts: self
                    .ssd_preload_experts
                    .map(u32::try_from)
                    .transpose()
                    .map_err(|_| "Saved SSD preload count is invalid.")?,
            },
            steering: SteeringPreferences {
                file: self.directional_steering_file.clone(),
                ffn_scale: self.directional_steering_ffn,
                attention_scale: self.directional_steering_attn,
            },
            diagnostics: DiagnosticPreferences {
                simulated_used_memory_gib: self
                    .simulated_used_memory_gib
                    .map(u64::try_from)
                    .transpose()
                    .map_err(|_| "Saved simulated memory is invalid.")?,
                expert_profile_path: self.expert_profile_path.clone(),
            },
            kv_cache: KvCachePreferences {
                budget_gib: self
                    .kv_budget_gib
                    .map(u64::try_from)
                    .transpose()
                    .map_err(|_| "Saved KV cache budget is invalid.")?,
                min_tokens: self
                    .kv_min_tokens
                    .map(u32::try_from)
                    .transpose()
                    .map_err(|_| "Saved KV cache minimum tokens is invalid.")?,
                cold_max_tokens: self
                    .kv_cold_max_tokens
                    .map(u32::try_from)
                    .transpose()
                    .map_err(|_| "Saved KV cache cold maximum is invalid.")?,
                continued_interval_tokens: self
                    .kv_continued_interval_tokens
                    .map(u32::try_from)
                    .transpose()
                    .map_err(|_| "Saved KV cache continued interval is invalid.")?,
            },
        })
    }
}

#[derive(AsChangeset)]
#[diesel(table_name = preferences)]
struct PreferenceChanges<'a> {
    selected_model: &'a str,
    dspark_enabled: bool,
    idle_timeout_minutes: i32,
    context_tokens: i32,
    max_generated_tokens: i32,
    system_prompt: &'a str,
    temperature: Option<f32>,
    top_p: Option<f32>,
    min_p: Option<f32>,
    seed: Option<&'a str>,
    reasoning_mode: &'a str,
    cpu_threads: Option<i32>,
    power_percent: Option<i32>,
    prefill_chunk: Option<i32>,
    quality: bool,
    warm_weights: bool,
    mtp_draft_tokens: i32,
    mtp_margin: f32,
    glm_mtp: bool,
    glm_mtp_timing: bool,
    dspark_confidence_threshold: Option<f32>,
    dspark_strict: bool,
    ssd_streaming: bool,
    ssd_streaming_cold: bool,
    ssd_cache_experts: Option<i64>,
    ssd_cache_gib: Option<i64>,
    ssd_full_layers: Option<i32>,
    ssd_preload_experts: Option<i32>,
    directional_steering_file: Option<&'a str>,
    directional_steering_ffn: Option<f32>,
    directional_steering_attn: Option<f32>,
    simulated_used_memory_gib: Option<i64>,
    expert_profile_path: Option<&'a str>,
    endpoint_port: i32,
    endpoint_enabled: bool,
    endpoint_cors: bool,
    kv_budget_gib: Option<i64>,
    kv_min_tokens: Option<i32>,
    kv_cold_max_tokens: Option<i32>,
    kv_continued_interval_tokens: Option<i32>,
}

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = projects)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct Project {
    pub id: i32,
    pub name: String,
    pub path: String,
    /// Sidebar state: hides this project's session rows.
    pub collapsed: bool,
}

#[derive(Insertable)]
#[diesel(table_name = projects)]
struct NewProject<'a> {
    name: &'a str,
    path: &'a str,
}

#[derive(Associations, Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(belongs_to(Project))]
#[diesel(table_name = sessions)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct Session {
    pub id: i32,
    pub project_id: i32,
    pub title: String,
    pub context_used: i32,
    pub context_limit: i32,
    pub last_tokens_per_second: Option<f32>,
    /// Raw column value; read it through [`Session::state`].
    state: String,
}

impl Session {
    pub fn state(&self) -> SessionState {
        SessionState::from_id(&self.state).unwrap_or_default()
    }

    #[cfg(test)]
    pub fn fixture(id: i32, project_id: i32, title: &str, state: SessionState) -> Self {
        Self {
            id,
            project_id,
            title: title.to_owned(),
            context_used: 0,
            context_limit: 0,
            last_tokens_per_second: None,
            state: state.as_id().to_owned(),
        }
    }
}

/// Where a session sits in the sidebar. These are lifecycle states, not flags:
/// a session is exactly one of them.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum SessionState {
    #[default]
    Normal,
    Pinned,
    Archived,
}

impl SessionState {
    pub fn as_id(self) -> &'static str {
        match self {
            Self::Normal => "normal",
            Self::Pinned => "pinned",
            Self::Archived => "archived",
        }
    }

    pub fn from_id(id: &str) -> Option<Self> {
        match id {
            "normal" => Some(Self::Normal),
            "pinned" => Some(Self::Pinned),
            "archived" => Some(Self::Archived),
            _ => None,
        }
    }

    /// Sidebar grouping order: pinned, then ordinary, then archived.
    pub fn rank(self) -> u8 {
        match self {
            Self::Pinned => 0,
            Self::Normal => 1,
            Self::Archived => 2,
        }
    }
}

#[derive(Insertable)]
#[diesel(table_name = sessions)]
struct NewSession<'a> {
    project_id: i32,
    title: &'a str,
}

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = messages)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct StoredMessage {
    pub id: i32,
    pub session_id: i32,
    pub user: bool,
    pub tool: bool,
    pub reasoning: Option<String>,
    pub reasoning_complete: bool,
    pub content: String,
}

#[derive(Insertable)]
#[diesel(table_name = messages)]
struct NewMessage<'a> {
    session_id: i32,
    user: bool,
    tool: bool,
    reasoning: Option<&'a str>,
    reasoning_complete: bool,
    content: &'a str,
}

#[derive(Debug)]
pub struct ProjectWithSessions {
    pub project: Project,
    pub sessions: Vec<Session>,
}

pub struct Database {
    connection: SqliteConnection,
}

impl Database {
    pub fn open(path: &Path) -> Result<Self, String> {
        let parent = path
            .parent()
            .ok_or_else(|| "database path has no parent directory".to_owned())?;
        fs::create_dir_all(parent).map_err(|error| error.to_string())?;
        let url = path
            .to_str()
            .ok_or_else(|| "database path is not valid UTF-8".to_owned())?;
        let mut connection = SqliteConnection::establish(url).map_err(|error| error.to_string())?;
        connection
            .run_pending_migrations(MIGRATIONS)
            .map_err(|error| error.to_string())?;
        Ok(Self { connection })
    }

    pub fn load_projects(&mut self) -> Result<Vec<ProjectWithSessions>, String> {
        let project_rows = projects::table
            .order(projects::id.asc())
            .select(Project::as_select())
            .load(&mut self.connection)
            .map_err(|error| error.to_string())?;
        let session_rows = sessions::table
            .order(sessions::id.asc())
            .select(Session::as_select())
            .load::<Session>(&mut self.connection)
            .map_err(|error| error.to_string())?;
        let mut sessions_by_project: HashMap<i32, Vec<Session>> = HashMap::new();
        for session in session_rows {
            sessions_by_project
                .entry(session.project_id)
                .or_default()
                .push(session);
        }
        // Pinned first, then ordinary, then archived; creation order within each
        // group. A stable sort keeps the id ordering the query established.
        for sessions in sessions_by_project.values_mut() {
            sessions.sort_by_key(|session| session.state().rank());
        }

        Ok(project_rows
            .into_iter()
            .map(|project| ProjectWithSessions {
                sessions: sessions_by_project.remove(&project.id).unwrap_or_default(),
                project,
            })
            .collect())
    }

    pub fn load_preferences(&mut self) -> Result<AppPreferences, String> {
        preferences::table
            .find(1)
            .select(AppPreferences::as_select())
            .first(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn set_sidebar_collapsed(&mut self, collapsed: bool) -> Result<(), String> {
        diesel::update(preferences::table.find(1))
            .set(preferences::sidebar_collapsed.eq(collapsed))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn set_sidebar_width(&mut self, width: i32) -> Result<(), String> {
        diesel::update(preferences::table.find(1))
            .set(preferences::sidebar_width.eq(width))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn set_last_project(&mut self, project_id: Option<i32>) -> Result<(), String> {
        diesel::update(preferences::table.find(1))
            .set(preferences::last_project_id.eq(project_id))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    #[allow(clippy::too_many_arguments)]
    pub fn update_preferences(
        &mut self,
        selected_model: &str,
        idle_timeout_minutes: i32,
        endpoint_port: i32,
        endpoint_enabled: bool,
        endpoint_cors: bool,
        generation: &GenerationPreferences,
        runtime: &RuntimePreferences,
    ) -> Result<AppPreferences, String> {
        generation.validate()?;
        let model = crate::model::ModelChoice::from_id(selected_model)
            .ok_or_else(|| format!("Unsupported model: {selected_model}"))?;
        runtime.validate(model)?;
        if !(1..=65_535).contains(&endpoint_port) {
            return Err("Endpoint port must be between 1 and 65535.".into());
        }
        let execution = &runtime.execution;
        let speculative = &runtime.speculative;
        let (ssd_cache_experts, ssd_cache_gib) = match runtime.ssd.cache {
            None => (None, None),
            Some(StreamingCacheBudget::Experts(experts)) => (Some(i64::from(experts)), None),
            Some(StreamingCacheBudget::Gib(gib)) => (
                None,
                Some(i64::try_from(gib).map_err(|_| "SSD cache budget is too large.")?),
            ),
        };
        let seed = generation.seed.map(|seed| seed.to_string());
        diesel::update(preferences::table.find(1))
            .set(PreferenceChanges {
                selected_model,
                dspark_enabled: speculative.dspark_enabled,
                idle_timeout_minutes,
                context_tokens: generation.context_tokens,
                max_generated_tokens: generation.max_generated_tokens,
                system_prompt: &generation.system_prompt,
                temperature: generation.temperature,
                top_p: generation.top_p,
                min_p: generation.min_p,
                seed: seed.as_deref(),
                reasoning_mode: generation.reasoning_mode.id(),
                cpu_threads: execution.cpu_threads.map(|value| value as i32),
                power_percent: execution.power_percent.map(i32::from),
                prefill_chunk: execution.prefill_chunk.map(|value| value as i32),
                quality: execution.quality,
                warm_weights: execution.warm_weights,
                mtp_draft_tokens: speculative.mtp_draft_tokens,
                mtp_margin: speculative.mtp_margin,
                glm_mtp: speculative.glm_mtp,
                glm_mtp_timing: speculative.glm_mtp_timing,
                dspark_confidence_threshold: speculative.dspark_confidence_threshold,
                dspark_strict: speculative.dspark_strict,
                ssd_streaming: runtime.ssd.enabled,
                ssd_streaming_cold: runtime.ssd.cold,
                ssd_cache_experts,
                ssd_cache_gib,
                ssd_full_layers: runtime.ssd.full_layers.map(|value| value as i32),
                ssd_preload_experts: runtime.ssd.preload_experts.map(|value| value as i32),
                directional_steering_file: runtime.steering.file.as_deref(),
                directional_steering_ffn: runtime.steering.ffn_scale,
                directional_steering_attn: runtime.steering.attention_scale,
                simulated_used_memory_gib: runtime
                    .diagnostics
                    .simulated_used_memory_gib
                    .map(|value| value as i64),
                expert_profile_path: runtime.diagnostics.expert_profile_path.as_deref(),
                endpoint_port,
                endpoint_enabled,
                endpoint_cors,
                kv_budget_gib: runtime
                    .kv_cache
                    .budget_gib
                    .map(i64::try_from)
                    .transpose()
                    .map_err(|_| "KV cache budget is too large.")?,
                kv_min_tokens: runtime.kv_cache.min_tokens.map(|value| value as i32),
                kv_cold_max_tokens: runtime.kv_cache.cold_max_tokens.map(|value| value as i32),
                kv_continued_interval_tokens: runtime
                    .kv_cache
                    .continued_interval_tokens
                    .map(|value| value as i32),
            })
            .returning(AppPreferences::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn create_project(&mut self, name: &str, path: &str) -> Result<Project, String> {
        diesel::insert_into(projects::table)
            .values(NewProject { name, path })
            .returning(Project::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn set_project_collapsed(
        &mut self,
        project_id: i32,
        collapsed: bool,
    ) -> Result<(), String> {
        diesel::update(projects::table.find(project_id))
            .set(projects::collapsed.eq(collapsed))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> {
        self.connection
            .transaction(|connection| {
                diesel::delete(
                    messages::table.filter(
                        messages::session_id.eq_any(
                            sessions::table
                                .filter(sessions::project_id.eq(project_id))
                                .select(sessions::id),
                        ),
                    ),
                )
                .execute(connection)?;
                diesel::delete(sessions::table.filter(sessions::project_id.eq(project_id)))
                    .execute(connection)?;
                diesel::delete(projects::table.find(project_id)).execute(connection)?;
                Ok(())
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn create_session(&mut self, project_id: i32, title: &str) -> Result<Session, String> {
        diesel::insert_into(sessions::table)
            .values(NewSession { project_id, title })
            .returning(Session::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn rename_session(&mut self, session_id: i32, title: &str) -> Result<(), String> {
        let title = title.trim();
        if title.is_empty() {
            return Err("Session title cannot be empty.".into());
        }
        diesel::update(sessions::table.find(session_id))
            .set(sessions::title.eq(title))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn set_session_state(
        &mut self,
        session_id: i32,
        state: SessionState,
    ) -> Result<(), String> {
        diesel::update(sessions::table.find(session_id))
            .set(sessions::state.eq(state.as_id()))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> {
        self.connection
            .transaction(|connection| {
                diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
                    .execute(connection)?;
                diesel::delete(sessions::table.find(session_id)).execute(connection)?;
                Ok(())
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn load_messages(&mut self, session_id: i32) -> Result<Vec<StoredMessage>, String> {
        messages::table
            .filter(messages::session_id.eq(session_id))
            .order(messages::id.asc())
            .select(StoredMessage::as_select())
            .load(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn update_session_context(
        &mut self,
        session_id: i32,
        used: u32,
        limit: u32,
        tokens_per_second: Option<f32>,
    ) -> Result<(), String> {
        let used = i32::try_from(used).map_err(|_| "Used context is too large to save")?;
        let limit = i32::try_from(limit).map_err(|_| "Context limit is too large to save")?;
        if tokens_per_second.is_some_and(|speed| !speed.is_finite() || speed < 0.0) {
            return Err("Generation speed is invalid".into());
        }
        diesel::update(sessions::table.find(session_id))
            .set((
                sessions::context_used.eq(used),
                sessions::context_limit.eq(limit),
                sessions::last_tokens_per_second.eq(tokens_per_second),
            ))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn start_chat_turn(
        &mut self,
        session_id: i32,
        prompt: &str,
        reasoning: bool,
    ) -> Result<(StoredMessage, StoredMessage), String> {
        self.connection
            .transaction(|connection| {
                let user = diesel::insert_into(messages::table)
                    .values(NewMessage {
                        session_id,
                        user: true,
                        tool: false,
                        reasoning: None,
                        reasoning_complete: true,
                        content: prompt,
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                let assistant = diesel::insert_into(messages::table)
                    .values(NewMessage {
                        session_id,
                        user: false,
                        tool: false,
                        reasoning: reasoning.then_some(""),
                        reasoning_complete: !reasoning,
                        content: "",
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                Ok((user, assistant))
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn continue_tool_turn(
        &mut self,
        session_id: i32,
        result: &str,
        reasoning: bool,
    ) -> Result<(StoredMessage, StoredMessage), String> {
        self.connection
            .transaction(|connection| {
                let tool = diesel::insert_into(messages::table)
                    .values(NewMessage {
                        session_id,
                        user: false,
                        tool: true,
                        reasoning: None,
                        reasoning_complete: true,
                        content: result,
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                let assistant = diesel::insert_into(messages::table)
                    .values(NewMessage {
                        session_id,
                        user: false,
                        tool: false,
                        reasoning: reasoning.then_some(""),
                        reasoning_complete: !reasoning,
                        content: "",
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                Ok((tool, assistant))
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn update_message(
        &mut self,
        id: i32,
        reasoning: Option<&str>,
        reasoning_complete: bool,
        content: &str,
    ) -> Result<(), String> {
        diesel::update(messages::table.find(id))
            .set((
                messages::reasoning.eq(reasoning),
                messages::reasoning_complete.eq(reasoning_complete),
                messages::content.eq(content),
            ))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::{SystemTime, UNIX_EPOCH};

    #[test]
    fn migrations_and_crud_work() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-server-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();

        let preferences = database.load_preferences().unwrap();
        assert_eq!(preferences.selected_model, "deepseek-v4-flash");
        assert!(!preferences.dspark_enabled);
        assert_eq!(preferences.idle_timeout_minutes, 10);
        assert_eq!(preferences.endpoint_port, 4000);
        assert!(preferences.endpoint_enabled);
        assert!(!preferences.endpoint_cors);
        let generation = GenerationPreferences::default();
        let runtime = RuntimePreferences::default();
        assert!(
            database
                .update_preferences(
                    "glm-5.2",
                    30,
                    4000,
                    true,
                    false,
                    &generation,
                    &RuntimePreferences {
                        speculative: SpeculativePreferences {
                            dspark_enabled: true,
                            ..runtime.speculative.clone()
                        },
                        ..runtime.clone()
                    }
                )
                .is_err()
        );
        assert!(
            database
                .update_preferences(
                    "deepseek-v4-flash",
                    0,
                    4000,
                    true,
                    false,
                    &generation,
                    &runtime,
                )
                .is_err()
        );
        let generation = GenerationPreferences {
            top_p: Some(0.7),
            seed: Some(u64::MAX),
            reasoning_mode: ReasoningMode::Max,
            ..generation
        };
        let execution = ExecutionPreferences {
            cpu_threads: Some(8),
            quality: true,
            ..runtime.execution
        };
        let speculative = SpeculativePreferences {
            glm_mtp: true,
            glm_mtp_timing: true,
            mtp_draft_tokens: 4,
            mtp_margin: 2.5,
            ..runtime.speculative
        };
        let runtime = RuntimePreferences {
            execution,
            speculative,
            ssd: SsdPreferences {
                enabled: true,
                cache: Some(StreamingCacheBudget::Gib(64)),
                full_layers: Some(0),
                preload_experts: Some(32),
                ..SsdPreferences::default()
            },
            diagnostics: DiagnosticPreferences {
                simulated_used_memory_gib: Some(8),
                expert_profile_path: Some("/tmp/experts.json".into()),
            },
            kv_cache: KvCachePreferences {
                budget_gib: Some(16),
                min_tokens: Some(256),
                cold_max_tokens: Some(0),
                continued_interval_tokens: Some(4_096),
            },
            ..RuntimePreferences::default()
        };
        database
            .update_preferences("glm-5.2", 30, 4567, false, true, &generation, &runtime)
            .unwrap();
        let preferences = database.load_preferences().unwrap();
        assert!(!preferences.endpoint_enabled);
        assert!(preferences.endpoint_cors);

        let project = database.create_project("DS4", "/tmp/ds4").unwrap();
        let first = database
            .create_session(project.id, "First session")
            .unwrap();
        database
            .create_session(project.id, "Second session")
            .unwrap();
        database.delete_session(first.id).unwrap();
        let loaded = database.load_projects().unwrap();
        assert_eq!(loaded[0].sessions[0].title, "Second session");
        assert_eq!(loaded[0].sessions[0].state(), SessionState::Normal);

        let ordinary = loaded[0].sessions[0].id;
        let pinned = database.create_session(project.id, "Pinned").unwrap();
        let archived = database.create_session(project.id, "Archived").unwrap();
        database
            .set_session_state(pinned.id, SessionState::Pinned)
            .unwrap();
        database
            .set_session_state(archived.id, SessionState::Archived)
            .unwrap();
        database.rename_session(pinned.id, "  Renamed  ").unwrap();
        assert!(database.rename_session(pinned.id, "   ").is_err());
        let loaded = database.load_projects().unwrap();
        assert_eq!(
            loaded[0]
                .sessions
                .iter()
                .map(|session| (session.id, session.state()))
                .collect::<Vec<_>>(),
            [
                (pinned.id, SessionState::Pinned),
                (ordinary, SessionState::Normal),
                (archived.id, SessionState::Archived),
            ]
        );
        assert_eq!(loaded[0].sessions[0].title, "Renamed");
        // Archiving a pinned session moves it out of the pinned group entirely.
        database
            .set_session_state(pinned.id, SessionState::Archived)
            .unwrap();
        let loaded = database.load_projects().unwrap();
        assert_eq!(loaded[0].sessions[0].id, ordinary);

        database.set_sidebar_collapsed(true).unwrap();
        database.set_last_project(Some(project.id)).unwrap();
        database.set_sidebar_width(320).unwrap();
        assert!(!loaded[0].project.collapsed);
        database.set_project_collapsed(project.id, true).unwrap();
        assert!(database.load_projects().unwrap()[0].project.collapsed);

        database.delete_project(project.id).unwrap();
        assert!(database.load_projects().unwrap().is_empty());
        drop(database);

        let mut reopened = Database::open(&path).unwrap();
        let preferences = reopened.load_preferences().unwrap();
        assert!(preferences.sidebar_collapsed);
        assert_eq!(preferences.last_project_id, Some(project.id));
        assert_eq!(preferences.sidebar_width, 320);
        assert_eq!(preferences.selected_model, "glm-5.2");
        assert_eq!(preferences.idle_timeout_minutes, 30);
        assert_eq!(preferences.endpoint_port, 4567);
        assert_eq!(preferences.generation().unwrap(), generation);
        assert_eq!(preferences.runtime().unwrap(), runtime);
        drop(reopened);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn chat_messages_survive_reopen_and_follow_session_deletion() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-chat-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();
        let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap();
        let session = database.create_session(project.id, "Chat").unwrap();
        let (_, assistant) = database
            .start_chat_turn(session.id, "Question", true)
            .unwrap();
        database
            .update_message(assistant.id, Some("Reasoning"), true, "Answer")
            .unwrap();
        database
            .continue_tool_turn(session.id, "Tool result", false)
            .unwrap();
        database
            .update_session_context(session.id, 1_234, 65_536, Some(12.5))
            .unwrap();
        drop(database);

        let mut reopened = Database::open(&path).unwrap();
        let projects = reopened.load_projects().unwrap();
        assert_eq!(projects[0].sessions[0].context_used, 1_234);
        assert_eq!(projects[0].sessions[0].context_limit, 65_536);
        assert_eq!(projects[0].sessions[0].last_tokens_per_second, Some(12.5));
        let messages = reopened.load_messages(session.id).unwrap();
        assert_eq!(messages.len(), 4);
        assert!(messages[0].user);
        assert!(!messages[0].tool);
        assert_eq!(messages[0].content, "Question");
        assert_eq!(messages[1].reasoning.as_deref(), Some("Reasoning"));
        assert!(messages[1].reasoning_complete);
        assert_eq!(messages[1].content, "Answer");
        assert!(messages[2].tool);
        assert_eq!(messages[2].content, "Tool result");
        assert!(!messages[3].user);
        assert!(!messages[3].tool);
        reopened.delete_session(session.id).unwrap();
        assert!(reopened.load_messages(session.id).unwrap().is_empty());
        drop(reopened);
        fs::remove_file(path).unwrap();
    }
}
