//! SQLite-backed persistence for research workspaces, notes, links, ghosts, and jobs.

use std::path::{Path, PathBuf};

use rusqlite::{params, Connection, OptionalExtension};

use crate::research::errors::{ResearchError, Result};
use crate::research::types::{
    AuditEvent, GhostLifecycleState, GhostNote, GhostVisibilityState, JobKind, JobStatus, NoteLink,
    NoteType, PipelineJob, ResearchNote, ResearchWorkspace, SourceRecord,
};

#[derive(Clone)]
pub struct ResearchRepository {
    db_path: PathBuf,
}

impl ResearchRepository {
    pub fn new(db_path: PathBuf) -> Result<Self> {
        if let Some(parent) = db_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let repository = Self { db_path };
        let connection = repository.open_connection()?;
        repository.initialize_schema(&connection)?;
        Ok(repository)
    }

    pub async fn create_workspace(
        &self,
        workspace: ResearchWorkspace,
    ) -> Result<ResearchWorkspace> {
        let value = workspace.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO research_workspaces (
                    id, name, primary_ticker, stage, default_view, archived, created_at, updated_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
                params![
                    value.id,
                    value.name,
                    value.primary_ticker,
                    serde_json::to_string(&value.stage)?,
                    serde_json::to_string(&value.default_view)?,
                    i64::from(value.archived),
                    value.created_at,
                    value.updated_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            Ok(value)
        })
        .await
    }

    #[allow(dead_code)]
    pub async fn save_workspace(&self, workspace: ResearchWorkspace) -> Result<ResearchWorkspace> {
        let value = workspace.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO research_workspaces (
                    id, name, primary_ticker, stage, default_view, archived, created_at, updated_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
                ON CONFLICT(id) DO UPDATE SET
                    name = excluded.name,
                    primary_ticker = excluded.primary_ticker,
                    stage = excluded.stage,
                    default_view = excluded.default_view,
                    archived = excluded.archived,
                    updated_at = excluded.updated_at,
                    entity_json = excluded.entity_json",
                params![
                    value.id,
                    value.name,
                    value.primary_ticker,
                    serde_json::to_string(&value.stage)?,
                    serde_json::to_string(&value.default_view)?,
                    i64::from(value.archived),
                    value.created_at,
                    value.updated_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            Ok(value)
        })
        .await
    }

    pub async fn list_workspaces(&self) -> Result<Vec<ResearchWorkspace>> {
        self.with_connection(|connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json
                 FROM research_workspaces
                 WHERE archived = 0
                 ORDER BY updated_at DESC",
            )?;
            let rows = statement.query_map([], |row| row.get::<_, String>(0))?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn get_workspace(&self, workspace_id: &str) -> Result<ResearchWorkspace> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let json = connection
                .query_row(
                    "SELECT entity_json FROM research_workspaces WHERE id = ?1",
                    params![workspace_id],
                    |row| row.get::<_, String>(0),
                )
                .optional()?
                .ok_or_else(|| ResearchError::WorkspaceNotFound(workspace_id.clone()))?;
            Ok(serde_json::from_str(&json)?)
        })
        .await
    }

    pub async fn create_note(&self, note: ResearchNote) -> Result<ResearchNote> {
        self.save_note(note).await
    }

    pub async fn save_note(&self, note: ResearchNote) -> Result<ResearchNote> {
        let value = note.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO research_notes (
                    id, workspace_id, note_type, ticker, source_id, archived, pinned, revision,
                    evidence_status, created_at, updated_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
                ON CONFLICT(id) DO UPDATE SET
                    workspace_id = excluded.workspace_id,
                    note_type = excluded.note_type,
                    ticker = excluded.ticker,
                    source_id = excluded.source_id,
                    archived = excluded.archived,
                    pinned = excluded.pinned,
                    revision = excluded.revision,
                    evidence_status = excluded.evidence_status,
                    updated_at = excluded.updated_at,
                    entity_json = excluded.entity_json",
                params![
                    value.id,
                    value.workspace_id,
                    serde_json::to_string(&value.note_type)?,
                    value.ticker,
                    value.source_id,
                    i64::from(value.archived),
                    i64::from(value.pinned),
                    i64::from(value.revision),
                    serde_json::to_string(&value.evidence_status)?,
                    value.created_at,
                    value.updated_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            connection.execute(
                "DELETE FROM research_fts WHERE note_id = ?1",
                params![value.id.clone()],
            )?;
            connection.execute(
                "INSERT INTO research_fts (note_id, title, cleaned_text, ai_annotation)
                 VALUES (?1, ?2, ?3, ?4)",
                params![
                    value.id,
                    value.title.clone().unwrap_or_default(),
                    value.cleaned_text,
                    value.ai_annotation.clone().unwrap_or_default(),
                ],
            )?;
            Ok(value)
        })
        .await
    }

    pub async fn get_note(&self, note_id: &str) -> Result<ResearchNote> {
        let note_id = note_id.to_string();
        self.with_connection(move |connection| {
            let json = connection
                .query_row(
                    "SELECT entity_json FROM research_notes WHERE id = ?1",
                    params![note_id],
                    |row| row.get::<_, String>(0),
                )
                .optional()?
                .ok_or_else(|| ResearchError::NoteNotFound(note_id.clone()))?;
            Ok(serde_json::from_str(&json)?)
        })
        .await
    }

    pub async fn list_notes(
        &self,
        workspace_id: &str,
        include_archived: bool,
        note_type: Option<NoteType>,
    ) -> Result<Vec<ResearchNote>> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let json_rows = if let Some(note_type) = note_type {
                let mut statement = connection.prepare(
                    "SELECT entity_json FROM research_notes
                     WHERE workspace_id = ?1 AND (?2 = 1 OR archived = 0) AND note_type = ?3
                     ORDER BY pinned DESC, updated_at DESC",
                )?;
                let rows = statement.query_map(
                    params![
                        workspace_id,
                        i64::from(include_archived),
                        serde_json::to_string(&note_type)?,
                    ],
                    |row| row.get::<_, String>(0),
                )?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            } else {
                let mut statement = connection.prepare(
                    "SELECT entity_json FROM research_notes
                     WHERE workspace_id = ?1 AND (?2 = 1 OR archived = 0)
                     ORDER BY pinned DESC, updated_at DESC",
                )?;
                let rows = statement
                    .query_map(params![workspace_id, i64::from(include_archived)], |row| {
                        row.get::<_, String>(0)
                    })?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            };

            json_rows
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn archive_note(&self, note_id: &str, archived: bool) -> Result<ResearchNote> {
        let mut note = self.get_note(note_id).await?;
        note.archived = archived;
        note.updated_at = crate::research::util::now_rfc3339();
        self.save_note(note).await
    }

    pub async fn find_source_reference_note(
        &self,
        workspace_id: &str,
        source_id: &str,
    ) -> Result<Option<ResearchNote>> {
        let workspace_id = workspace_id.to_string();
        let source_id = source_id.to_string();
        self.with_connection(move |connection| {
            let json = connection
                .query_row(
                    "SELECT entity_json
                     FROM research_notes
                     WHERE workspace_id = ?1 AND source_id = ?2 AND note_type = ?3
                     LIMIT 1",
                    params![
                        workspace_id,
                        source_id,
                        serde_json::to_string(&NoteType::SourceReference)?,
                    ],
                    |row| row.get::<_, String>(0),
                )
                .optional()?;
            json.map(|value| serde_json::from_str(&value).map_err(ResearchError::from))
                .transpose()
        })
        .await
    }

    pub async fn save_source(&self, source: SourceRecord) -> Result<SourceRecord> {
        let value = source.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO source_records (
                    id, workspace_id, ticker, kind, published_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
                ON CONFLICT(id) DO UPDATE SET
                    workspace_id = excluded.workspace_id,
                    ticker = excluded.ticker,
                    kind = excluded.kind,
                    published_at = excluded.published_at,
                    entity_json = excluded.entity_json",
                params![
                    value.id,
                    value.workspace_id,
                    value.ticker,
                    serde_json::to_string(&value.kind)?,
                    value.published_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            Ok(value)
        })
        .await
    }

    pub async fn list_sources(&self, workspace_id: &str) -> Result<Vec<SourceRecord>> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json FROM source_records WHERE workspace_id = ?1 ORDER BY published_at DESC, id DESC",
            )?;
            let rows = statement.query_map(params![workspace_id], |row| row.get::<_, String>(0))?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn list_sources_by_ids(&self, source_ids: &[String]) -> Result<Vec<SourceRecord>> {
        if source_ids.is_empty() {
            return Ok(Vec::new());
        }

        let ids = source_ids.to_vec();
        self.with_connection(move |connection| {
            let mut results = Vec::new();
            let mut statement =
                connection.prepare("SELECT entity_json FROM source_records WHERE id = ?1")?;
            for source_id in ids {
                if let Some(json) = statement
                    .query_row(params![source_id], |row| row.get::<_, String>(0))
                    .optional()?
                {
                    results.push(serde_json::from_str(&json)?);
                }
            }
            Ok(results)
        })
        .await
    }

    pub async fn find_source_by_checksum_or_accession(
        &self,
        workspace_id: &str,
        checksum: Option<&str>,
        filing_accession: Option<&str>,
    ) -> Result<Option<SourceRecord>> {
        let workspace_id = workspace_id.to_string();
        let checksum = checksum.map(ToOwned::to_owned);
        let filing_accession = filing_accession.map(ToOwned::to_owned);
        self.with_connection(move |connection| {
            let json = if let Some(checksum) = checksum {
                connection
                    .query_row(
                        "SELECT entity_json FROM source_records
                         WHERE workspace_id = ?1 AND json_extract(entity_json, '$.checksum') = ?2
                         LIMIT 1",
                        params![workspace_id, checksum],
                        |row| row.get::<_, String>(0),
                    )
                    .optional()?
            } else if let Some(filing_accession) = filing_accession {
                connection
                    .query_row(
                        "SELECT entity_json FROM source_records
                         WHERE workspace_id = ?1 AND json_extract(entity_json, '$.filingAccession') = ?2
                         LIMIT 1",
                        params![workspace_id, filing_accession],
                        |row| row.get::<_, String>(0),
                    )
                    .optional()?
            } else {
                None
            };

            json.map(|value| serde_json::from_str(&value).map_err(ResearchError::from))
                .transpose()
        })
        .await
    }

    pub async fn replace_links_for_workspace(
        &self,
        workspace_id: &str,
        links: Vec<NoteLink>,
    ) -> Result<Vec<NoteLink>> {
        let workspace_id = workspace_id.to_string();
        let values = links.clone();
        self.with_connection(move |connection| {
            let transaction = connection.transaction()?;
            transaction.execute("DELETE FROM note_links WHERE workspace_id = ?1", params![workspace_id])?;
            for link in &values {
                transaction.execute(
                    "INSERT INTO note_links (
                        id, workspace_id, from_note_id, to_note_id, link_type, stale, updated_at, entity_json
                    ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
                    params![
                        link.id,
                        link.workspace_id,
                        link.from_note_id,
                        link.to_note_id,
                        serde_json::to_string(&link.link_type)?,
                        i64::from(link.stale),
                        link.updated_at,
                        serde_json::to_string(link)?,
                    ],
                )?;
            }
            transaction.commit()?;
            Ok(values)
        })
        .await
    }

    pub async fn list_links(
        &self,
        workspace_id: &str,
        note_id: Option<&str>,
    ) -> Result<Vec<NoteLink>> {
        let workspace_id = workspace_id.to_string();
        let note_id = note_id.map(ToOwned::to_owned);
        self.with_connection(move |connection| {
            let json_rows = if let Some(note_id) = note_id {
                let mut statement = connection.prepare(
                    "SELECT entity_json
                     FROM note_links
                     WHERE workspace_id = ?1 AND (from_note_id = ?2 OR to_note_id = ?2)
                     ORDER BY updated_at DESC",
                )?;
                let rows = statement.query_map(params![workspace_id, note_id], |row| {
                    row.get::<_, String>(0)
                })?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            } else {
                let mut statement = connection.prepare(
                    "SELECT entity_json
                     FROM note_links
                     WHERE workspace_id = ?1
                     ORDER BY updated_at DESC",
                )?;
                let rows =
                    statement.query_map(params![workspace_id], |row| row.get::<_, String>(0))?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            };

            json_rows
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn replace_ghosts_for_workspace(
        &self,
        workspace_id: &str,
        ghosts: Vec<GhostNote>,
    ) -> Result<Vec<GhostNote>> {
        let workspace_id = workspace_id.to_string();
        let generated = ghosts.clone();
        self.with_connection(move |connection| {
            let transaction = connection.transaction()?;
            let existing = {
                let mut existing_statement = transaction
                    .prepare("SELECT id, entity_json FROM ghost_notes WHERE workspace_id = ?1")?;
                let existing_rows = existing_statement
                    .query_map(params![workspace_id.clone()], |row| {
                        Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
                    })?;
                existing_rows.collect::<std::result::Result<Vec<_>, _>>()?
            };
            transaction.execute(
                "DELETE FROM ghost_notes WHERE workspace_id = ?1",
                params![workspace_id],
            )?;

            for mut ghost in generated.clone() {
                if let Some((_, json)) = existing.iter().find(|(id, _)| *id == ghost.id) {
                    let prior: GhostNote = serde_json::from_str(json)?;
                    if matches!(
                        prior.state,
                        GhostLifecycleState::Accepted
                            | GhostLifecycleState::Dismissed
                            | GhostLifecycleState::Converted
                            | GhostLifecycleState::Ignored
                    ) {
                        ghost.state = prior.state;
                    }
                    if matches!(prior.visibility_state, GhostVisibilityState::Pinned) {
                        ghost.visibility_state = prior.visibility_state;
                    }
                    ghost.promoted_note_id = prior.promoted_note_id;
                }

                transaction.execute(
                    "INSERT INTO ghost_notes (
                        id, workspace_id, state, visibility_state, updated_at, entity_json
                    ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
                    params![
                        ghost.id,
                        ghost.workspace_id,
                        serde_json::to_string(&ghost.state)?,
                        serde_json::to_string(&ghost.visibility_state)?,
                        ghost.updated_at,
                        serde_json::to_string(&ghost)?,
                    ],
                )?;
            }

            transaction.commit()?;
            Ok(generated)
        })
        .await
    }

    pub async fn get_ghost(&self, ghost_note_id: &str) -> Result<GhostNote> {
        let ghost_note_id = ghost_note_id.to_string();
        self.with_connection(move |connection| {
            let json = connection
                .query_row(
                    "SELECT entity_json FROM ghost_notes WHERE id = ?1",
                    params![ghost_note_id],
                    |row| row.get::<_, String>(0),
                )
                .optional()?
                .ok_or_else(|| ResearchError::GhostNoteNotFound(ghost_note_id.clone()))?;
            Ok(serde_json::from_str(&json)?)
        })
        .await
    }

    pub async fn save_ghost(&self, ghost: GhostNote) -> Result<GhostNote> {
        let value = ghost.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO ghost_notes (
                    id, workspace_id, state, visibility_state, updated_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
                ON CONFLICT(id) DO UPDATE SET
                    state = excluded.state,
                    visibility_state = excluded.visibility_state,
                    updated_at = excluded.updated_at,
                    entity_json = excluded.entity_json",
                params![
                    value.id,
                    value.workspace_id,
                    serde_json::to_string(&value.state)?,
                    serde_json::to_string(&value.visibility_state)?,
                    value.updated_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            Ok(value)
        })
        .await
    }

    pub async fn list_ghosts(
        &self,
        workspace_id: &str,
        include_hidden: bool,
    ) -> Result<Vec<GhostNote>> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json FROM ghost_notes
                 WHERE workspace_id = ?1
                 ORDER BY updated_at DESC",
            )?;
            let rows = statement.query_map(params![workspace_id], |row| row.get::<_, String>(0))?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str::<GhostNote>(&json).map_err(ResearchError::from))
                .collect::<Result<Vec<_>>>()
                .map(|ghosts| {
                    ghosts
                        .into_iter()
                        .filter(|ghost| {
                            include_hidden
                                || !matches!(ghost.visibility_state, GhostVisibilityState::Hidden)
                        })
                        .collect()
                })
        })
        .await
    }

    pub async fn enqueue_jobs(&self, jobs: Vec<PipelineJob>) -> Result<Vec<PipelineJob>> {
        let values = jobs.clone();
        self.with_connection(move |connection| {
            let transaction = connection.transaction()?;
            for job in &values {
                transaction.execute(
                    "INSERT INTO pipeline_jobs (
                        id, workspace_id, entity_id, job_kind, status, next_attempt_at, updated_at, entity_json
                    ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
                    ON CONFLICT(id) DO UPDATE SET
                        status = excluded.status,
                        next_attempt_at = excluded.next_attempt_at,
                        updated_at = excluded.updated_at,
                        entity_json = excluded.entity_json",
                    params![
                        job.id,
                        job.workspace_id,
                        job.entity_id,
                        serde_json::to_string(&job.job_kind)?,
                        serde_json::to_string(&job.status)?,
                        job.next_attempt_at,
                        job.updated_at,
                        serde_json::to_string(job)?,
                    ],
                )?;
            }
            transaction.commit()?;
            Ok(values)
        })
        .await
    }

    pub async fn save_job(&self, job: PipelineJob) -> Result<PipelineJob> {
        let value = job.clone();
        self.enqueue_jobs(vec![value.clone()]).await?;
        Ok(value)
    }

    pub async fn list_due_jobs(&self, limit: usize) -> Result<Vec<PipelineJob>> {
        self.with_connection(move |connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json
                 FROM pipeline_jobs
                 WHERE status IN (?1, ?2)
                   AND (next_attempt_at IS NULL OR next_attempt_at <= datetime('now'))
                 ORDER BY updated_at ASC
                 LIMIT ?3",
            )?;
            let rows = statement.query_map(
                params![
                    serde_json::to_string(&JobStatus::Queued)?,
                    serde_json::to_string(&JobStatus::Failed)?,
                    i64::try_from(limit)
                        .map_err(|error| ResearchError::Validation(error.to_string()))?,
                ],
                |row| row.get::<_, String>(0),
            )?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn list_jobs(
        &self,
        workspace_id: &str,
        job_kind: Option<JobKind>,
    ) -> Result<Vec<PipelineJob>> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let json_rows = if let Some(job_kind) = job_kind {
                let mut statement = connection.prepare(
                    "SELECT entity_json FROM pipeline_jobs
                     WHERE workspace_id = ?1 AND job_kind = ?2
                     ORDER BY updated_at DESC",
                )?;
                let rows = statement.query_map(
                    params![workspace_id, serde_json::to_string(&job_kind)?],
                    |row| row.get::<_, String>(0),
                )?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            } else {
                let mut statement = connection.prepare(
                    "SELECT entity_json FROM pipeline_jobs
                     WHERE workspace_id = ?1
                     ORDER BY updated_at DESC",
                )?;
                let rows =
                    statement.query_map(params![workspace_id], |row| row.get::<_, String>(0))?;
                rows.collect::<std::result::Result<Vec<_>, _>>()?
            };

            json_rows
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn retry_failed_jobs(
        &self,
        workspace_id: &str,
        job_kind: Option<JobKind>,
    ) -> Result<Vec<PipelineJob>> {
        let mut jobs = self.list_jobs(workspace_id, job_kind).await?;
        let mut retried = Vec::new();
        for job in &mut jobs {
            if matches!(job.status, JobStatus::Failed | JobStatus::Skipped) {
                job.status = JobStatus::Queued;
                job.next_attempt_at = None;
                job.last_error = None;
                job.updated_at = crate::research::util::now_rfc3339();
                retried.push(job.clone());
            }
        }
        self.enqueue_jobs(retried.clone()).await?;
        Ok(retried)
    }

    pub async fn append_audit_event(&self, event: AuditEvent) -> Result<AuditEvent> {
        let value = event.clone();
        self.with_connection(move |connection| {
            connection.execute(
                "INSERT INTO audit_events (
                    id, workspace_id, entity_id, entity_kind, created_at, entity_json
                ) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
                params![
                    value.id,
                    value.workspace_id,
                    value.entity_id,
                    serde_json::to_string(&value.entity_kind)?,
                    value.created_at,
                    serde_json::to_string(&value)?,
                ],
            )?;
            Ok(value)
        })
        .await
    }

    pub async fn list_audit_events_for_entity(&self, entity_id: &str) -> Result<Vec<AuditEvent>> {
        let entity_id = entity_id.to_string();
        self.with_connection(move |connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json FROM audit_events WHERE entity_id = ?1 ORDER BY created_at ASC",
            )?;
            let rows = statement.query_map(params![entity_id], |row| row.get::<_, String>(0))?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    pub async fn list_audit_events_for_workspace(
        &self,
        workspace_id: &str,
    ) -> Result<Vec<AuditEvent>> {
        let workspace_id = workspace_id.to_string();
        self.with_connection(move |connection| {
            let mut statement = connection.prepare(
                "SELECT entity_json FROM audit_events WHERE workspace_id = ?1 ORDER BY created_at ASC",
            )?;
            let rows = statement.query_map(params![workspace_id], |row| row.get::<_, String>(0))?;
            rows.collect::<std::result::Result<Vec<_>, _>>()?
                .into_iter()
                .map(|json| serde_json::from_str(&json).map_err(ResearchError::from))
                .collect()
        })
        .await
    }

    fn open_connection(&self) -> Result<Connection> {
        open_connection(&self.db_path)
    }

    fn initialize_schema(&self, connection: &Connection) -> Result<()> {
        connection.execute_batch(
            "PRAGMA foreign_keys = ON;
             PRAGMA journal_mode = WAL;
             PRAGMA synchronous = NORMAL;
             CREATE TABLE IF NOT EXISTS research_workspaces (
                 id TEXT PRIMARY KEY,
                 name TEXT NOT NULL,
                 primary_ticker TEXT NOT NULL,
                 stage TEXT NOT NULL,
                 default_view TEXT NOT NULL,
                 archived INTEGER NOT NULL DEFAULT 0,
                 created_at TEXT NOT NULL,
                 updated_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE TABLE IF NOT EXISTS research_notes (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 note_type TEXT NOT NULL,
                 ticker TEXT,
                 source_id TEXT,
                 archived INTEGER NOT NULL DEFAULT 0,
                 pinned INTEGER NOT NULL DEFAULT 0,
                 revision INTEGER NOT NULL,
                 evidence_status TEXT NOT NULL,
                 created_at TEXT NOT NULL,
                 updated_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE INDEX IF NOT EXISTS research_notes_workspace_archived_updated_idx
                 ON research_notes (workspace_id, archived, updated_at DESC);
             CREATE INDEX IF NOT EXISTS research_notes_workspace_type_archived_idx
                 ON research_notes (workspace_id, note_type, archived);
             CREATE INDEX IF NOT EXISTS research_notes_workspace_ticker_updated_idx
                 ON research_notes (workspace_id, ticker, updated_at DESC);
             CREATE TABLE IF NOT EXISTS note_links (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 from_note_id TEXT NOT NULL,
                 to_note_id TEXT NOT NULL,
                 link_type TEXT NOT NULL,
                 stale INTEGER NOT NULL DEFAULT 0,
                 updated_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE INDEX IF NOT EXISTS note_links_workspace_from_type_idx
                 ON note_links (workspace_id, from_note_id, link_type);
             CREATE TABLE IF NOT EXISTS ghost_notes (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 state TEXT NOT NULL,
                 visibility_state TEXT NOT NULL,
                 updated_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE INDEX IF NOT EXISTS ghost_notes_workspace_state_visibility_idx
                 ON ghost_notes (workspace_id, state, visibility_state, updated_at DESC);
             CREATE TABLE IF NOT EXISTS source_records (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 ticker TEXT,
                 kind TEXT NOT NULL,
                 published_at TEXT,
                 entity_json TEXT NOT NULL
             );
             CREATE INDEX IF NOT EXISTS source_records_workspace_ticker_kind_published_idx
                 ON source_records (workspace_id, ticker, kind, published_at DESC);
             CREATE TABLE IF NOT EXISTS source_excerpts (
                 id TEXT PRIMARY KEY,
                 source_id TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE TABLE IF NOT EXISTS pipeline_jobs (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 entity_id TEXT NOT NULL,
                 job_kind TEXT NOT NULL,
                 status TEXT NOT NULL,
                 next_attempt_at TEXT,
                 updated_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE INDEX IF NOT EXISTS pipeline_jobs_status_next_attempt_idx
                 ON pipeline_jobs (status, next_attempt_at);
             CREATE TABLE IF NOT EXISTS workspace_view_state (
                 workspace_id TEXT PRIMARY KEY,
                 entity_json TEXT NOT NULL
             );
             CREATE TABLE IF NOT EXISTS memo_exports (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE TABLE IF NOT EXISTS audit_events (
                 id TEXT PRIMARY KEY,
                 workspace_id TEXT NOT NULL REFERENCES research_workspaces(id) ON DELETE CASCADE,
                 entity_id TEXT NOT NULL,
                 entity_kind TEXT NOT NULL,
                 created_at TEXT NOT NULL,
                 entity_json TEXT NOT NULL
             );
             CREATE VIRTUAL TABLE IF NOT EXISTS research_fts USING fts5(note_id UNINDEXED, title, cleaned_text, ai_annotation);",
        )?;
        Ok(())
    }

    async fn with_connection<F, T>(&self, task: F) -> Result<T>
    where
        F: FnOnce(&mut Connection) -> Result<T> + Send + 'static,
        T: Send + 'static,
    {
        let db_path = self.db_path.clone();
        tokio::task::spawn_blocking(move || {
            let mut connection = open_connection(&db_path)?;
            task(&mut connection)
        })
        .await
        .map_err(|error| ResearchError::Join(error.to_string()))?
    }
}

fn open_connection(path: &Path) -> Result<Connection> {
    let connection = Connection::open(path)?;
    connection.execute_batch(
        "PRAGMA foreign_keys = ON;
         PRAGMA journal_mode = WAL;
         PRAGMA synchronous = NORMAL;",
    )?;
    Ok(connection)
}

#[cfg(test)]
mod tests {
    use tempfile::tempdir;

    use super::ResearchRepository;
    use crate::research::types::{
        ResearchStage, ResearchWorkspace, WorkspaceScope, WorkspaceViewKind,
    };

    #[tokio::test]
    async fn new_should_initialize_schema() {
        let dir = tempdir().unwrap();
        let repository = ResearchRepository::new(dir.path().join("research.sqlite")).unwrap();
        let now = crate::research::util::now_rfc3339();
        repository
            .create_workspace(ResearchWorkspace {
                id: "workspace-1".to_string(),
                name: "AAPL".to_string(),
                primary_ticker: "AAPL".to_string(),
                scope: WorkspaceScope::SingleCompany,
                stage: ResearchStage::Capture,
                default_view: WorkspaceViewKind::Canvas,
                pinned_note_ids: Vec::new(),
                archived: false,
                created_at: now.clone(),
                updated_at: now,
            })
            .await
            .unwrap();

        let workspaces = repository.list_workspaces().await.unwrap();
        assert_eq!(workspaces.len(), 1);
    }
}
