//! Durable bounded semantic source tasks. Only completed source groups may
//! supply vectors; tokenizer/inference work stays outside SQLite transactions.
use super::*;
use searchi_core::IndexedPassage;
use store::{content_hash, SourceTask};

pub(super) struct BuildWork {
    pub report: EmbeddingReport,
    pub keys: Vec<String>,
    pub media: Vec<searchi_media::MediaIdentity>,
    documents: HashMap<String, searchi_media::MediaIdentity>,
}

impl BuildWork {
    pub fn track(&mut self, task: &SourceTask) -> Result<()> {
        ensure!(
            self.keys.len() < MAX_SCAN_CHUNKS,
            "Semantic source task budget exceeded"
        );
        self.keys.push(task.key.clone());
        Ok(())
    }
}

pub(super) fn task_key(parts: &[&str]) -> String {
    let mut input = String::new();
    for part in parts {
        input.push_str(&format!("{}:", part.len()));
        input.push_str(part);
    }
    content_hash(&input)
}

pub(super) fn text_task(passage: &IndexedPassage, policy: &str) -> Result<SourceTask> {
    let key = task_key(&[
        "text",
        &passage.path,
        &passage.source.kind,
        &passage.source.number.to_string(),
    ]);
    Ok(SourceTask {
        key,
        path: passage.path.clone(),
        source_type: "text".into(),
        source_hash: content_hash(&passage.text),
        title_hash: content_hash(bounded_title(&passage.filename)),
        policy: policy.into(),
        passage_id: passage.id,
    })
}

impl HybridSearch {
    pub(super) fn rebuild_generation(
        &mut self,
        fingerprint: String,
        build: impl FnOnce(&mut Self, i64, &mut BuildWork) -> Result<()>,
    ) -> Result<EmbeddingReport> {
        let lock = self.store.build_lock()?;
        let (generation, resumed) = self.store.start_or_resume(&fingerprint, &lock)?;
        // This revision covers the entire discovery/build/finalization pass.
        // A concurrent lexical update aborts publication; next invocation
        // reconciles and resumes all still-compatible completed groups.
        let revision = self.store.revision()?;
        let previous_fingerprint = std::mem::replace(&mut self.generation_fingerprint, fingerprint);
        let mut work = BuildWork {
            report: EmbeddingReport {
                media_files: 0,
                media_failures: 0,
                failures: Vec::new(),
                passages: 0,
                chunks: 0,
                generation,
                model_fingerprint: self.generation_fingerprint.clone(),
                reused_sources: 0,
                reused_chunks: 0,
                encoded_sources: 0,
                document_inference_calls: 0,
                image_inference_calls: 0,
                audio_inference_calls: 0,
                resumed,
            },
            keys: Vec::new(),
            media: Vec::new(),
            documents: HashMap::new(),
        };
        let result = build(self, generation, &mut work).and_then(|()| {
            for identity in work.media.iter().chain(work.documents.values()) {
                // Full bytes, not just stat metadata: in-place changes can
                // preserve size and timestamps, including on mounted volumes.
                let current = searchi_media::MediaIdentity::read(&identity.path, identity.bytes)?;
                ensure!(
                    current.sha256 == identity.sha256
                        && current.modified_ns == identity.modified_ns,
                    "Source changed during semantic indexing; retry required"
                );
            }
            self.store
                .reconcile_tasks(generation, &work.keys, revision)?;
            self.store.activate_checked(generation, revision)?;
            Ok(work.report)
        });
        if result.is_err() {
            self.generation_fingerprint = previous_fingerprint;
        }
        result
    }

    pub(super) fn build(&mut self, generation: i64, work: &mut BuildWork) -> Result<()> {
        let mut after_id = 0;
        loop {
            let passages = self.lexical.passage_batch(after_id, 64)?;
            if passages.is_empty() {
                break;
            }
            for passage in passages {
                after_id = passage.id;
                if passage.text.trim().is_empty() {
                    continue;
                }
                if !work.documents.contains_key(&passage.path) {
                    ensure!(
                        work.documents.len() < MAX_SCAN_CHUNKS,
                        "Document source budget exceeded"
                    );
                    let identity = searchi_media::MediaIdentity::read(
                        Path::new(&passage.path),
                        32 * 1024 * 1024,
                    )?;
                    work.documents.insert(passage.path.clone(), identity);
                }
                let task = text_task(&passage, &self.generation_fingerprint)?;
                work.track(&task)?;
                self.store.stage_task(generation, &task)?;
                let result = self.build_text_source(generation, &task, &passage, work);
                if let Err(error) = result {
                    if let Err(cleanup) =
                        self.store
                            .fail_source(generation, &task, &format!("{error:#}"))
                    {
                        return Err(error.context(format!(
                            "Failed to discard incomplete text source: {cleanup:#}"
                        )));
                    }
                    return Err(error);
                }
            }
        }
        Ok(())
    }

    fn build_text_source(
        &mut self,
        generation: i64,
        task: &SourceTask,
        passage: &IndexedPassage,
        work: &mut BuildWork,
    ) -> Result<()> {
        if let Some(count) = self.store.reuse_text_source(generation, task, passage)? {
            ensure!(
                work.report.chunks + count <= MAX_SCAN_CHUNKS,
                "Corpus exceeds ANN chunk budget"
            );
            work.report.chunks += count;
            work.report.passages += 1;
            work.report.reused_chunks += count;
            work.report.reused_sources += 1;
            return Ok(());
        }
        let title = bounded_title(&passage.filename);
        let windows = passage_windows(&mut *self.embedder, &passage.text, title)?;
        ensure!(
            work.report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
            "Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
        );
        for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
            let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
            work.report.document_inference_calls += 1;
            let vectors = self.embedder.embed_documents(&texts, title)?;
            self.store.insert(
                generation,
                passage,
                batch_index * MAX_BATCH,
                batch,
                &vectors,
            )?;
        }
        self.store
            .complete_source(generation, task, windows.len())?;
        work.report.chunks += windows.len();
        work.report.passages += 1;
        work.report.encoded_sources += 1;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Fault-only provider: never returns a vector or claims semantic quality.
    struct RejectInference;
    impl EmbeddingProvider for RejectInference {
        fn fingerprint(&self) -> &str {
            "fault-injection-only"
        }
        fn dimensions(&self) -> usize {
            GEMMA_DIMENSIONS
        }
        fn split_windows(&mut self, text: &str, _: &str) -> Result<Vec<TextWindow>> {
            Ok(vec![TextWindow {
                text: text.into(),
                start_byte: 0,
                end_byte: text.len(),
            }])
        }
        fn embed_documents(&mut self, _: &[&str], _: &str) -> Result<Vec<Vec<f32>>> {
            anyhow::bail!("injected primary inference failure")
        }
        fn embed_query(&mut self, _: &str) -> Result<Vec<f32>> {
            anyhow::bail!("fault provider never supplies query vectors")
        }
    }

    #[test]
    fn cleanup_failure_keeps_primary_error_reader_policy_and_staging() -> Result<()> {
        let temp = tempfile::tempdir()?;
        let root = temp.path().join("corpus");
        std::fs::create_dir(&root)?;
        std::fs::write(
            root.join("notes.txt"),
            "A real source whose inference is deliberately rejected.",
        )?;
        let db = temp.path().join("index.db");
        SearchIndex::open(&db)?.index_folder(&root)?;
        let mut service = HybridSearch::with_provider(&db, Box::new(RejectInference))?;
        let prior_policy = service.generation_fingerprint.clone();
        let connection = rusqlite::Connection::open(&db)?;
        connection.execute_batch("CREATE TRIGGER reject_task_cleanup BEFORE UPDATE OF state ON semantic_tasks WHEN new.state='failed' BEGIN SELECT RAISE(ABORT,'injected task cleanup failure'); END;")?;
        let error = service.index_embeddings().unwrap_err();
        let message = format!("{error:#}");
        assert!(
            message.contains("injected primary inference failure"),
            "{message}"
        );
        assert!(
            message.contains("injected task cleanup failure"),
            "{message}"
        );
        assert_eq!(service.generation_fingerprint, prior_policy);
        let staging: usize =
            connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
                r.get(0)
            })?;
        assert_eq!(staging, 1);
        let state: String =
            connection.query_row("SELECT state FROM semantic_tasks", [], |r| r.get(0))?;
        assert_eq!(
            state, "running",
            "Failed cleanup transaction must roll back"
        );
        connection.execute_batch("DROP TRIGGER reject_task_cleanup;")?;
        let error = service.index_embeddings().unwrap_err();
        assert!(format!("{error:#}").contains("injected primary inference failure"));
        assert_eq!(service.generation_fingerprint, prior_policy);
        let staging: usize =
            connection.query_row("SELECT count(*) FROM semantic_generations", [], |r| {
                r.get(0)
            })?;
        assert_eq!(staging, 1, "Resume must retain the same staging generation");
        Ok(())
    }
}