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, projects, sessions};

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

#[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,
    pub compacted_summary: Option<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(),
            compacted_summary: None,
        }
    }
}

/// 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,
    pub system: bool,
    pub compaction: bool,
    pub compaction_tail_start: Option<i32>,
}

#[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,
    system: bool,
    compaction: bool,
    compaction_tail_start: Option<i32>,
}

#[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 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,
        system_messages: &[String],
        reasoning: bool,
    ) -> Result<Vec<StoredMessage>, String> {
        self.connection
            .transaction(|connection| {
                let mut stored = Vec::with_capacity(system_messages.len() + 2);
                for content in system_messages {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage {
                                session_id,
                                user: false,
                                tool: false,
                                reasoning: None,
                                reasoning_complete: true,
                                content,
                                system: true,
                                compaction: false,
                                compaction_tail_start: None,
                            })
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                let user = diesel::insert_into(messages::table)
                    .values(NewMessage {
                        session_id,
                        user: true,
                        tool: false,
                        reasoning: None,
                        reasoning_complete: true,
                        content: prompt,
                        system: false,
                        compaction: false,
                        compaction_tail_start: None,
                    })
                    .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: "",
                        system: false,
                        compaction: false,
                        compaction_tail_start: None,
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                stored.extend([user, assistant]);
                Ok(stored)
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn continue_tool_turn(
        &mut self,
        session_id: i32,
        result: &str,
        queued_user: Option<&str>,
        system_messages: &[String],
        reasoning: bool,
    ) -> Result<Vec<StoredMessage>, String> {
        self.connection
            .transaction(|connection| {
                let mut stored = Vec::with_capacity(system_messages.len() + 3);
                stored.push(
                    diesel::insert_into(messages::table)
                        .values(NewMessage {
                            session_id,
                            user: false,
                            tool: true,
                            reasoning: None,
                            reasoning_complete: true,
                            content: result,
                            system: false,
                            compaction: false,
                            compaction_tail_start: None,
                        })
                        .returning(StoredMessage::as_returning())
                        .get_result(connection)?,
                );
                if let Some(content) = queued_user {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage {
                                session_id,
                                user: true,
                                tool: false,
                                reasoning: None,
                                reasoning_complete: true,
                                content,
                                system: false,
                                compaction: false,
                                compaction_tail_start: None,
                            })
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                for content in system_messages {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage {
                                session_id,
                                user: false,
                                tool: false,
                                reasoning: None,
                                reasoning_complete: true,
                                content,
                                system: true,
                                compaction: false,
                                compaction_tail_start: None,
                            })
                            .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: "",
                        system: false,
                        compaction: false,
                        compaction_tail_start: None,
                    })
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                stored.push(assistant);
                Ok(stored)
            })
            .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())
    }

    pub fn record_compaction(
        &mut self,
        session_id: i32,
        summary: &str,
        tail_start: Option<i32>,
        running_jobs: Option<&str>,
        context_used: u32,
        context_limit: u32,
    ) -> Result<Vec<StoredMessage>, String> {
        let context_used = i32::try_from(context_used)
            .map_err(|_| "Compacted context is too large to save".to_owned())?;
        let context_limit = i32::try_from(context_limit)
            .map_err(|_| "Context limit is too large to save".to_owned())?;
        self.connection
            .transaction(|connection| {
                diesel::update(sessions::table.find(session_id))
                    .set((
                        sessions::compacted_summary.eq(Some(summary)),
                        sessions::context_used.eq(context_used),
                        sessions::context_limit.eq(context_limit),
                        sessions::last_tokens_per_second.eq(None::<f32>),
                    ))
                    .execute(connection)?;
                let mut stored = Vec::with_capacity(2);
                stored.push(
                    diesel::insert_into(messages::table)
                        .values(NewMessage {
                            session_id,
                            user: false,
                            tool: false,
                            reasoning: None,
                            reasoning_complete: true,
                            content: summary,
                            system: true,
                            compaction: true,
                            compaction_tail_start: tail_start,
                        })
                        .returning(StoredMessage::as_returning())
                        .get_result(connection)?,
                );
                if let Some(content) = running_jobs {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage {
                                session_id,
                                user: false,
                                tool: true,
                                reasoning: None,
                                reasoning_complete: true,
                                content,
                                system: false,
                                compaction: false,
                                compaction_tail_start: None,
                            })
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                Ok(stored)
            })
            .map_err(|error: diesel::result::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 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);

        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();
        assert!(reopened.load_projects().unwrap().is_empty());
        drop(reopened);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn chat_and_compaction_history_survive_reopen_and_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 mut opening = database
            .start_chat_turn(session.id, "Question", &["Date context".into()], true)
            .unwrap();
        let assistant = opening.pop().unwrap();
        database
            .update_message(assistant.id, Some("Reasoning"), true, "Answer")
            .unwrap();
        database
            .continue_tool_turn(
                session.id,
                "Tool result",
                Some("Queued correction"),
                &["Tool reminder".into()],
                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(), 7);
        assert!(messages[0].system);
        assert_eq!(messages[1].content, "Question");
        assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning"));
        assert!(messages[2].reasoning_complete);
        assert_eq!(messages[2].content, "Answer");
        assert!(messages[3].tool);
        assert_eq!(messages[3].content, "Tool result");
        assert!(messages[4].user);
        assert_eq!(messages[4].content, "Queued correction");
        assert!(messages[5].system);
        assert_eq!(messages[5].content, "Tool reminder");
        assert!(!messages[6].user);
        assert!(!messages[6].tool);
        let first = reopened
            .record_compaction(
                session.id,
                "First durable state.",
                Some(messages[4].id),
                Some("bash job=1 status=running"),
                321,
                32_768,
            )
            .unwrap();
        assert_eq!(first.len(), 2);
        assert!(first[0].compaction);
        assert_eq!(first[0].compaction_tail_start, Some(messages[4].id));
        drop(reopened);
        let mut reopened = Database::open(&path).unwrap();
        assert_eq!(
            reopened.load_projects().unwrap()[0].sessions[0]
                .compacted_summary
                .as_deref(),
            Some("First durable state.")
        );
        assert_eq!(
            reopened.load_projects().unwrap()[0].sessions[0].context_used,
            321
        );
        let messages = reopened.load_messages(session.id).unwrap();
        assert_eq!(messages.len(), 9);
        assert_eq!(messages[1].content, "Question");
        assert!(messages[7].compaction);
        assert_eq!(messages[7].content, "First durable state.");
        assert_eq!(messages[8].content, "bash job=1 status=running");
        reopened
            .continue_tool_turn(session.id, "Reloaded tool result", None, &[], false)
            .unwrap();
        let continued = reopened.load_messages(session.id).unwrap();
        let second_tail = continued[9].id;
        reopened
            .record_compaction(
                session.id,
                "Second durable state.",
                Some(second_tail),
                None,
                222,
                32_768,
            )
            .unwrap();
        reopened
            .continue_tool_turn(session.id, "Final tool result", None, &[], false)
            .unwrap();
        let continued = reopened.load_messages(session.id).unwrap();
        let third_tail = continued[12].id;
        reopened
            .record_compaction(
                session.id,
                "Third durable state.",
                Some(third_tail),
                None,
                111,
                32_768,
            )
            .unwrap();
        reopened
            .start_chat_turn(session.id, "After third compaction", &[], false)
            .unwrap();
        drop(reopened);

        let mut reopened = Database::open(&path).unwrap();
        let history = reopened.load_messages(session.id).unwrap();
        assert_eq!(history.len(), 17);
        assert_eq!(history[1].content, "Question");
        let markers = history
            .iter()
            .filter(|message| message.compaction)
            .collect::<Vec<_>>();
        assert_eq!(markers.len(), 3);
        assert_eq!(markers[0].content, "First durable state.");
        assert_eq!(markers[1].content, "Second durable state.");
        assert_eq!(markers[2].content, "Third durable state.");
        assert_eq!(markers[2].compaction_tail_start, Some(third_tail));
        assert_eq!(history[15].content, "After third compaction");
        reopened.delete_session(session.id).unwrap();
        assert!(reopened.load_messages(session.id).unwrap().is_empty());
        drop(reopened);
        fs::remove_file(path).unwrap();
    }
}
