use engram_embeddings::cache::EmbeddingCache;
use engram_embeddings::embedder::Embedder;
use engram_embeddings::error::EmbeddingError;
use engram_embeddings::hyde;
use engram_llm_client::error::ApiError;
use engram_llm_client::provider::{EmbeddingProvider, TextGenerator};

// ── Mock providers ─────────────────────────────────────────────────────

struct MockEmbeddingProvider {
    dimension: usize,
}

impl EmbeddingProvider for MockEmbeddingProvider {
    fn embed(&self, text: &str, input_type: Option<&str>) -> Result<Vec<f32>, ApiError> {
        let mut embedding = vec![0.0_f32; self.dimension];
        for (index, byte) in text.bytes().enumerate() {
            embedding[index % self.dimension] += byte as f32 / 255.0;
        }
        // Differentiate document vs query so cache-collision tests can
        // distinguish what the embedder actually requested. Offsets are
        // small constants — vector remains deterministic per (text, kind).
        if let Some(kind) = input_type {
            let offset = match kind {
                "document" => 0.01,
                "query" => 0.02,
                _ => 0.0,
            };
            for value in &mut embedding {
                *value += offset;
            }
        }
        Ok(embedding)
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn model_name(&self) -> &str {
        "mock-embedding"
    }
}

struct FailingEmbeddingProvider;

impl EmbeddingProvider for FailingEmbeddingProvider {
    fn embed(&self, _text: &str, _input_type: Option<&str>) -> Result<Vec<f32>, ApiError> {
        Err(ApiError::EmbeddingApiUnavailable("mock failure".into()))
    }

    fn dimension(&self) -> usize {
        128
    }

    fn model_name(&self) -> &str {
        "failing-embedding"
    }
}

struct MockTextGenerator {
    response: String,
}

impl TextGenerator for MockTextGenerator {
    fn generate(&self, _prompt: &str) -> Result<String, ApiError> {
        Ok(self.response.clone())
    }

    fn model_name(&self) -> &str {
        "mock-generator"
    }
}

struct FailingTextGenerator;

impl TextGenerator for FailingTextGenerator {
    fn generate(&self, _prompt: &str) -> Result<String, ApiError> {
        Err(ApiError::HyDeGenerationFailed(
            "mock generation failure".into(),
        ))
    }

    fn model_name(&self) -> &str {
        "failing-generator"
    }
}

// ── embed_fields ───────────────────────────────────────────────────────

#[test]
fn embed_fields_returns_three_embeddings() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);

    let result = embedder
        .embed_fields(
            "user opened the editor and started typing a document",
            "typed a paragraph about rust programming language features",
            "the document was saved successfully to the local filesystem",
            &provider,
            None,
        )
        .unwrap();

    assert_eq!(result.context.len(), 4);
    assert_eq!(result.action.len(), 4);
    assert_eq!(result.result.len(), 4);
}

#[test]
fn embed_fields_caches_identical_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);

    let same_text = "user opened the editor and started typing a long document about rust";

    embedder
        .embed_fields(same_text, same_text, same_text, &provider, None)
        .unwrap();

    assert_eq!(embedder.cache().len(), 1);
}

#[test]
fn embed_fields_different_texts_produce_different_embeddings() {
    let provider = MockEmbeddingProvider { dimension: 8 };
    let embedder = Embedder::new(0);

    let result = embedder
        .embed_fields(
            "user opened the editor window and began writing a new document from scratch",
            "typed several paragraphs about the rust programming language and its features",
            "the document was saved successfully to the local filesystem without any errors",
            &provider,
            None,
        )
        .unwrap();

    assert_ne!(result.context, result.action);
    assert_ne!(result.action, result.result);
    assert_ne!(result.context, result.result);
}

// ── embed_query ────────────────────────────────────────────────────────

#[test]
fn embed_query_returns_single_embedding() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);

    let long_query =
        "search for all documents that were created in the last week about rust programming";
    let result = embedder.embed_query(long_query, &provider, None).unwrap();

    assert_eq!(result.len(), 4);
}

#[test]
fn embed_query_uses_cache() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);

    let query = "search for documents about rust programming language features and memory safety";

    let first = embedder.embed_query(query, &provider, None).unwrap();
    let second = embedder.embed_query(query, &provider, None).unwrap();

    assert_eq!(first, second);
    assert_eq!(embedder.cache().len(), 1);
}

// ── HyDE ───────────────────────────────────────────────────────────────

#[test]
fn hyde_triggered_for_short_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let generator = MockTextGenerator {
        response:
            "this is a hypothetical document generated by the text generator for short queries"
                .into(),
    };
    let embedder = Embedder::new(20);

    let short_query = "rust memory safety";

    embedder
        .embed_query(short_query, &provider, Some(&generator))
        .unwrap();

    // ADR 2026-05-05: cache key is now ORIGINAL query, not hypothesis
    assert!(embedder.cache().get(short_query, Some("query")).is_some());
    assert!(
        embedder
            .cache()
            .get(&generator.response, Some("query"))
            .is_none()
    );
}

#[test]
fn hyde_not_triggered_for_long_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let generator = MockTextGenerator {
        response:
            "hypothesis that should not be used because the input text is already long enough"
                .into(),
    };
    let embedder = Embedder::new(20);

    let long_text = "one two three four five six seven eight nine ten \
                     eleven twelve thirteen fourteen fifteen sixteen \
                     seventeen eighteen nineteen twenty";

    embedder
        .embed_query(long_text, &provider, Some(&generator))
        .unwrap();

    // Cache should contain the raw text, not the hypothesis
    assert!(embedder.cache().get(long_text, Some("query")).is_some());
    assert!(
        embedder
            .cache()
            .get(&generator.response, Some("query"))
            .is_none()
    );
}

#[test]
fn hyde_failure_falls_back_to_raw_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let generator = FailingTextGenerator;
    let embedder = Embedder::new(20);

    let short_query = "rust memory";

    let result = embedder.embed_query(short_query, &provider, Some(&generator));

    assert!(result.is_ok());
    assert!(embedder.cache().get(short_query, Some("query")).is_some());
}

#[test]
fn hyde_without_text_generator_embeds_raw() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(20);

    let short_query = "rust memory";

    let result = embedder.embed_query(short_query, &provider, None);

    assert!(result.is_ok());
    assert!(embedder.cache().get(short_query, Some("query")).is_some());
}

// ── Error propagation ──────────────────────────────────────────────────

#[test]
fn provider_error_propagates() {
    let provider = FailingEmbeddingProvider;
    let embedder = Embedder::new(0);

    let long_text =
        "a sufficiently long text that will not trigger hyde generation at all for testing";
    let result = embedder.embed_query(long_text, &provider, None);

    assert!(result.is_err());
    let error = result.unwrap_err();
    assert!(
        matches!(error, EmbeddingError::ProviderError(_)),
        "expected ProviderError, got: {error}"
    );
}

// ── Cache ──────────────────────────────────────────────────────────────

#[test]
fn cache_clear_works() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);

    let query = "search for documents about rust programming and its unique memory safety features";

    embedder.embed_query(query, &provider, None).unwrap();
    assert_eq!(embedder.cache().len(), 1);

    embedder.clear_cache();
    assert!(embedder.cache().is_empty());

    // Re-embedding should work after cache clear
    embedder.embed_query(query, &provider, None).unwrap();
    assert_eq!(embedder.cache().len(), 1);
}

#[test]
fn cache_starts_empty() {
    let cache = EmbeddingCache::new();
    assert!(cache.is_empty());
    assert_eq!(cache.len(), 0);
}

// ── hyde module boundary tests ─────────────────────────────────────────

#[test]
fn should_use_hyde_boundary() {
    // Exactly 20 words -> false
    let twenty_words = "one two three four five six seven eight nine ten \
                        eleven twelve thirteen fourteen fifteen sixteen \
                        seventeen eighteen nineteen twenty";
    assert!(!hyde::should_use_hyde(twenty_words, 20));
}

#[test]
fn should_use_hyde_under_threshold() {
    // 19 words -> true
    let nineteen_words = "one two three four five six seven eight nine ten \
                          eleven twelve thirteen fourteen fifteen sixteen \
                          seventeen eighteen nineteen";
    assert!(hyde::should_use_hyde(nineteen_words, 20));
}

// ── EmbeddingError ─────────────────────────────────────────────────────

#[test]
fn embedding_error_display_format() {
    let provider_error =
        EmbeddingError::ProviderError(ApiError::EmbeddingApiUnavailable("timeout".into()));
    assert_eq!(
        provider_error.to_string(),
        "embedding provider error: [2001] embedding api unavailable: timeout"
    );
}

#[test]
fn embedding_error_from_api_error() {
    let api_error = ApiError::EmbeddingApiUnavailable("connection lost".into());
    let embedding_error: EmbeddingError = api_error.into();

    assert!(matches!(embedding_error, EmbeddingError::ProviderError(_)));
    assert_eq!(
        embedding_error.to_string(),
        "embedding provider error: [2001] embedding api unavailable: connection lost"
    );
}

// ── embed_fields with failing text generator ──────────────────────────

#[test]
fn embed_fields_with_failing_generator_falls_back_to_raw_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let generator = FailingTextGenerator;
    let embedder = Embedder::new(0);

    let result = embedder.embed_fields(
        "rust memory",
        "open editor",
        "saved file",
        &provider,
        Some(&generator),
    );

    assert!(result.is_ok());
    let embedding = result.unwrap();
    assert_eq!(embedding.context.len(), 4);
    assert_eq!(embedding.action.len(), 4);
    assert_eq!(embedding.result.len(), 4);
}

// ── Error::source() ───────────────────────────────────────────────────

#[test]
fn embedding_error_source_returns_underlying_api_error() {
    use std::error::Error;

    let api_error = ApiError::EmbeddingApiUnavailable("connection reset".into());
    let embedding_error = EmbeddingError::ProviderError(api_error);

    let source = embedding_error.source();
    assert!(source.is_some());
    assert!(source.unwrap().to_string().contains("connection reset"));
}

// ── Default trait impls ───────────────────────────────────────────────

#[test]
fn embedder_default_has_empty_cache() {
    let embedder = Embedder::default();
    assert!(embedder.cache().is_empty());
}

#[test]
fn embedding_cache_default_is_empty() {
    let cache = EmbeddingCache::default();
    assert!(cache.is_empty());
    assert_eq!(cache.len(), 0);
}

// ── embed_fields with HyDE ───────────────────────────────────────────

#[test]
fn embed_fields_with_hyde_uses_hypothesis_for_short_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let generator = MockTextGenerator {
        response:
            "this is a hypothetical document generated by the text generator for short fields"
                .into(),
    };
    let embedder = Embedder::new(20);

    let short_context = "rust memory";
    let short_action = "open editor";
    let short_result = "saved file";

    let result = embedder.embed_fields(
        short_context,
        short_action,
        short_result,
        &provider,
        Some(&generator),
    );

    assert!(result.is_ok());
    // All short fields should trigger HyDE and resolve to the same hypothesis.
    // embed_fields stores under input_type="document" (key includes input_type).
    assert!(
        embedder
            .cache()
            .get(&generator.response, Some("document"))
            .is_some()
    );
    assert!(
        embedder
            .cache()
            .get(short_context, Some("document"))
            .is_none()
    );
    assert!(
        embedder
            .cache()
            .get(short_action, Some("document"))
            .is_none()
    );
    assert!(
        embedder
            .cache()
            .get(short_result, Some("document"))
            .is_none()
    );
}

// ── ADR 2026-05-05 cache-by-original + HyDE opt-in ──────────────────

struct CountingTextGenerator {
    response: String,
    calls: std::sync::atomic::AtomicUsize,
}

impl CountingTextGenerator {
    fn new(response: impl Into<String>) -> Self {
        Self {
            response: response.into(),
            calls: std::sync::atomic::AtomicUsize::new(0),
        }
    }

    fn calls(&self) -> usize {
        self.calls.load(std::sync::atomic::Ordering::SeqCst)
    }
}

impl TextGenerator for CountingTextGenerator {
    fn generate(&self, _prompt: &str) -> Result<String, ApiError> {
        self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        Ok(self.response.clone())
    }

    fn model_name(&self) -> &str {
        "counting-generator"
    }
}

struct CountingEmbeddingProvider {
    dimension: usize,
    calls: std::sync::atomic::AtomicUsize,
}

impl CountingEmbeddingProvider {
    fn new(dimension: usize) -> Self {
        Self {
            dimension,
            calls: std::sync::atomic::AtomicUsize::new(0),
        }
    }

    fn calls(&self) -> usize {
        self.calls.load(std::sync::atomic::Ordering::SeqCst)
    }
}

impl EmbeddingProvider for CountingEmbeddingProvider {
    fn embed(&self, text: &str, _input_type: Option<&str>) -> Result<Vec<f32>, ApiError> {
        self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let mut embedding = vec![0.0_f32; self.dimension];
        for (index, byte) in text.bytes().enumerate() {
            embedding[index % self.dimension] += byte as f32 / 255.0;
        }
        Ok(embedding)
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn model_name(&self) -> &str {
        "counting-embedding"
    }
}

#[test]
fn embed_query_caches_by_original_query() {
    let provider = CountingEmbeddingProvider::new(4);
    let generator = CountingTextGenerator::new("hypothetical document for short query");
    let embedder = Embedder::new(20);
    let query = "rust";

    embedder
        .embed_query(query, &provider, Some(&generator))
        .unwrap();
    embedder
        .embed_query(query, &provider, Some(&generator))
        .unwrap();

    assert!(embedder.cache().get(query, Some("query")).is_some());
    assert!(
        embedder
            .cache()
            .get(&generator.response, Some("query"))
            .is_none()
    );
    assert_eq!(provider.calls(), 1);
}

#[test]
fn embed_query_cache_hit_skips_hyde() {
    let provider = CountingEmbeddingProvider::new(4);
    let generator = CountingTextGenerator::new("hypothesis");
    let embedder = Embedder::new(20);

    embedder
        .embed_query("rust", &provider, Some(&generator))
        .unwrap();
    embedder
        .embed_query("rust", &provider, Some(&generator))
        .unwrap();

    assert_eq!(generator.calls(), 1);
}

#[test]
fn should_use_hyde_threshold_zero_disables() {
    assert!(!hyde::should_use_hyde("any short text here", 0));
    assert!(!hyde::should_use_hyde("", 0));
    assert!(!hyde::should_use_hyde("rust", 0));
}

#[test]
fn should_use_hyde_threshold_strict_less_than() {
    assert!(!hyde::should_use_hyde("five word text count check", 5));
    assert!(hyde::should_use_hyde("four word text count", 5));
}

#[test]
fn should_use_hyde_empty_text_returns_false() {
    assert!(!hyde::should_use_hyde("", 5));
    assert!(!hyde::should_use_hyde("   ", 5));
    assert!(!hyde::should_use_hyde("\t\n", 20));
}

#[test]
fn should_use_hyde_threshold_one_excludes_single_word() {
    assert!(!hyde::should_use_hyde("hello", 1));
}

#[test]
fn should_use_hyde_threshold_max_active_for_short_text() {
    assert!(hyde::should_use_hyde("rust", usize::MAX));
}

#[test]
fn embedder_threshold_zero_skips_hyde_even_with_generator() {
    let provider = CountingEmbeddingProvider::new(4);
    let generator = CountingTextGenerator::new("never invoked");
    let embedder = Embedder::new(0);

    embedder
        .embed_query("rust", &provider, Some(&generator))
        .unwrap();

    assert_eq!(generator.calls(), 0);
    assert!(embedder.cache().get("rust", Some("query")).is_some());
}

#[test]
fn embed_fields_caches_under_document_input_type() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);
    let context = "user opened the editor and started typing a new document file";
    let action = "typed several paragraphs about rust programming language and its features";
    let result = "the document was saved successfully to the local filesystem without errors";

    embedder
        .embed_fields(context, action, result, &provider, None)
        .unwrap();

    assert!(embedder.cache().get(context, Some("document")).is_some());
    assert!(embedder.cache().get(action, Some("document")).is_some());
    assert!(embedder.cache().get(result, Some("document")).is_some());
    // Query-keyed lookup of the same text MUST miss — different cache slot.
    assert!(embedder.cache().get(context, Some("query")).is_none());
}

#[test]
fn embed_query_caches_under_query_input_type() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);
    let query = "search for documents about rust programming language and its memory model";

    embedder.embed_query(query, &provider, None).unwrap();

    assert!(embedder.cache().get(query, Some("query")).is_some());
    // Document-keyed lookup MUST miss for the same text.
    assert!(embedder.cache().get(query, Some("document")).is_none());
}

#[test]
fn cache_distinguishes_document_and_query_for_same_text() {
    let provider = MockEmbeddingProvider { dimension: 4 };
    let embedder = Embedder::new(0);
    let text = "rust ownership and borrowing rules govern memory access at compile time";

    // First write the document embedding via embed_fields.
    embedder
        .embed_fields(text, text, text, &provider, None)
        .unwrap();
    // Then write the query embedding for the same text.
    embedder.embed_query(text, &provider, None).unwrap();

    let document_embedding = embedder.cache().get(text, Some("document")).unwrap();
    let query_embedding = embedder.cache().get(text, Some("query")).unwrap();

    // MockEmbeddingProvider yields different offsets for document vs query —
    // proves the two entries are independently stored, not collapsed.
    assert_ne!(document_embedding, query_embedding);
}

#[test]
fn embedder_threshold_twenty_invokes_hyde_for_short_text() {
    let provider = CountingEmbeddingProvider::new(4);
    let generator = CountingTextGenerator::new("hypothetical document");
    let embedder = Embedder::new(20);

    embedder
        .embed_query("rust", &provider, Some(&generator))
        .unwrap();

    assert_eq!(generator.calls(), 1);
}
