use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::error::Error;
use std::fmt;

/// Stable model identifier for the built-in local text embedding provider.
pub const LOCAL_TEXT_EMBEDDING_MODEL_ID: &str = "mentisdb-local-text";
/// Stable version label for the built-in local text embedding provider.
pub const LOCAL_TEXT_EMBEDDING_VERSION: &str = "v1";
/// Fixed dimension for the built-in local text embedding provider.
pub const LOCAL_TEXT_EMBEDDING_DIMENSION: usize = 256;

/// Stable metadata describing one embedding space.
///
/// A vector index should only contain vectors generated by the same metadata
/// tuple so similarity scores remain comparable.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub struct EmbeddingMetadata {
    /// Stable provider or model identifier.
    pub model_id: String,
    /// Fixed vector dimension for this embedding space.
    pub dimension: usize,
    /// Version label for model or preprocessing changes.
    pub embedding_version: String,
}

impl EmbeddingMetadata {
    /// Create embedding metadata with model id, vector dimension, and version.
    pub fn new(
        model_id: impl Into<String>,
        dimension: usize,
        embedding_version: impl Into<String>,
    ) -> Self {
        Self {
            model_id: model_id.into(),
            dimension: dimension.max(1),
            embedding_version: embedding_version.into(),
        }
    }
}

/// Input payload for one provider embedding request.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct EmbeddingInput {
    /// Caller-defined stable id that should map back to the source document.
    pub input_id: String,
    /// Text to embed.
    pub text: String,
}

impl EmbeddingInput {
    /// Create one embedding input with a stable id and source text.
    pub fn new(input_id: impl Into<String>, text: impl Into<String>) -> Self {
        Self {
            input_id: input_id.into(),
            text: text.into(),
        }
    }
}

/// One embedding vector returned by a provider.
#[derive(Debug, Clone, PartialEq)]
pub struct EmbeddingVector {
    /// Dense numeric representation in the provider's embedding space.
    pub values: Vec<f32>,
}

impl EmbeddingVector {
    /// Wrap an embedding vector.
    pub fn new(values: Vec<f32>) -> Self {
        Self { values }
    }
}

/// Provider-agnostic batch embedding interface.
pub trait EmbeddingProvider {
    /// Provider-specific embedding error.
    type Error: Error + Send + Sync + 'static;

    /// Return metadata for vectors produced by this provider.
    fn metadata(&self) -> &EmbeddingMetadata;

    /// Embed a batch of inputs in the provider's declared vector space.
    fn embed_batch(&self, inputs: &[EmbeddingInput]) -> Result<Vec<EmbeddingVector>, Self::Error>;
}

/// Non-failing error type for the built-in local text embedding provider.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LocalTextEmbeddingError;

impl fmt::Display for LocalTextEmbeddingError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("local text embedding provider does not fail")
    }
}

impl Error for LocalTextEmbeddingError {}

/// Deterministic in-process embedding provider used by `mentisdbd`.
///
/// This provider hashes normalized tokens and trigrams into a fixed dense
/// vector. It is lightweight, fully local, and requires no network calls or
/// model downloads, which makes it suitable as the daemon's default sidecar
/// provider.
#[derive(Debug, Clone)]
pub struct LocalTextEmbeddingProvider {
    metadata: EmbeddingMetadata,
}

impl Default for LocalTextEmbeddingProvider {
    fn default() -> Self {
        Self::new()
    }
}

impl LocalTextEmbeddingProvider {
    /// Create the built-in local text embedding provider.
    pub fn new() -> Self {
        Self {
            metadata: EmbeddingMetadata::new(
                LOCAL_TEXT_EMBEDDING_MODEL_ID,
                LOCAL_TEXT_EMBEDDING_DIMENSION,
                LOCAL_TEXT_EMBEDDING_VERSION,
            ),
        }
    }

    fn vector_for_text(&self, text: &str) -> Vec<f32> {
        let mut vector = vec![0.0; self.metadata.dimension];
        let normalized = text.to_ascii_lowercase();
        let trimmed = normalized.trim();
        let mut token_count = 0usize;

        for token in normalized
            .split(|character: char| !character.is_ascii_alphanumeric())
            .filter(|token| !token.is_empty())
        {
            token_count += 1;
            accumulate_hashed_feature(&mut vector, token.as_bytes(), 1.0);
        }

        let compact: Vec<u8> = normalized
            .bytes()
            .filter(|byte| byte.is_ascii_alphanumeric())
            .collect();
        for trigram in compact.windows(3) {
            accumulate_hashed_feature(&mut vector, trigram, 0.35);
        }

        if token_count == 0 && !trimmed.is_empty() {
            accumulate_hashed_feature(&mut vector, trimmed.as_bytes(), 1.0);
        }
        if vector.iter().all(|value| *value == 0.0) {
            vector[0] = 1.0;
            return vector;
        }

        let magnitude = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
        if magnitude > 0.0 {
            for value in &mut vector {
                *value /= magnitude;
            }
        }
        vector
    }
}

impl EmbeddingProvider for LocalTextEmbeddingProvider {
    type Error = LocalTextEmbeddingError;

    fn metadata(&self) -> &EmbeddingMetadata {
        &self.metadata
    }

    fn embed_batch(&self, inputs: &[EmbeddingInput]) -> Result<Vec<EmbeddingVector>, Self::Error> {
        Ok(inputs
            .iter()
            .map(|input| EmbeddingVector::new(self.vector_for_text(&input.text)))
            .collect())
    }
}

/// Failure while converting provider outputs into vector documents.
#[derive(Debug, PartialEq, Eq)]
pub enum EmbeddingBuildError<E> {
    /// Provider returned an error.
    Provider(E),
    /// Provider returned a different number of vectors than requested inputs.
    OutputCountMismatch {
        /// Number of input rows requested.
        expected: usize,
        /// Number of vectors returned by the provider.
        actual: usize,
    },
    /// Returned vector did not match provider metadata dimension.
    DimensionMismatch {
        /// Expected vector dimension from provider metadata.
        expected: usize,
        /// Actual vector length returned by provider.
        actual: usize,
        /// Position of the offending input/vector pair.
        input_index: usize,
    },
    /// Returned vector contained `NaN` or `inf`.
    NonFiniteValue {
        /// Position of the offending input/vector pair.
        input_index: usize,
        /// Position of the non-finite value inside the vector.
        value_index: usize,
    },
}

impl<E: fmt::Display> fmt::Display for EmbeddingBuildError<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Provider(error) => write!(f, "embedding provider error: {error}"),
            Self::OutputCountMismatch { expected, actual } => write!(
                f,
                "embedding provider returned {actual} vectors for {expected} inputs"
            ),
            Self::DimensionMismatch {
                expected,
                actual,
                input_index,
            } => write!(
                f,
                "embedding vector dimension mismatch at input index {input_index}: expected {expected}, got {actual}"
            ),
            Self::NonFiniteValue {
                input_index,
                value_index,
            } => write!(
                f,
                "embedding vector contains non-finite value at input index {input_index}, value index {value_index}"
            ),
        }
    }
}

impl<E: Error + 'static> Error for EmbeddingBuildError<E> {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Provider(error) => Some(error),
            _ => None,
        }
    }
}

/// Convert provider embeddings into deterministic vector documents.
pub fn embed_batch_to_documents<P: EmbeddingProvider>(
    provider: &P,
    inputs: &[EmbeddingInput],
) -> Result<Vec<VectorDocument>, EmbeddingBuildError<P::Error>> {
    let outputs = provider
        .embed_batch(inputs)
        .map_err(EmbeddingBuildError::Provider)?;
    if outputs.len() != inputs.len() {
        return Err(EmbeddingBuildError::OutputCountMismatch {
            expected: inputs.len(),
            actual: outputs.len(),
        });
    }

    let dimension = provider.metadata().dimension;
    inputs
        .iter()
        .zip(outputs)
        .enumerate()
        .map(|(input_index, (input, output))| {
            if output.values.len() != dimension {
                return Err(EmbeddingBuildError::DimensionMismatch {
                    expected: dimension,
                    actual: output.values.len(),
                    input_index,
                });
            }
            for (value_index, value) in output.values.iter().enumerate() {
                if !value.is_finite() {
                    return Err(EmbeddingBuildError::NonFiniteValue {
                        input_index,
                        value_index,
                    });
                }
            }
            Ok(VectorDocument::new(input.input_id.clone(), output.values))
        })
        .collect()
}

/// One vectorized document tracked by the index.
#[derive(Debug, Clone, PartialEq)]
pub struct VectorDocument {
    /// Stable caller-defined id for deterministic tie-breaking and lookup.
    pub document_id: String,
    /// Dense vector in a shared embedding space.
    pub vector: Vec<f32>,
}

impl VectorDocument {
    /// Create a vector document with id and vector values.
    pub fn new(document_id: impl Into<String>, vector: Vec<f32>) -> Self {
        Self {
            document_id: document_id.into(),
            vector,
        }
    }
}

/// Query payload for cosine-ranked vector retrieval.
#[derive(Debug, Clone, PartialEq)]
pub struct VectorQuery {
    /// Query vector in the same embedding space as the index.
    pub vector: Vec<f32>,
    /// Maximum number of hits to return.
    pub limit: usize,
}

impl VectorQuery {
    /// Create a query with default limit of 10.
    pub fn new(vector: Vec<f32>) -> Self {
        Self { vector, limit: 10 }
    }

    /// Set maximum number of hits to return.
    pub fn with_limit(mut self, limit: usize) -> Self {
        self.limit = limit.max(1);
        self
    }
}

/// Ranked vector-search hit.
#[derive(Debug, Clone, PartialEq)]
pub struct VectorSearchHit {
    /// Stable document id.
    pub document_id: String,
    /// Cosine similarity score in `[-1.0, 1.0]`.
    pub score: f32,
}

/// In-memory deterministic cosine index for vector documents.
#[derive(Debug, Clone)]
pub struct VectorIndex {
    metadata: EmbeddingMetadata,
    documents: BTreeMap<String, Vec<f32>>,
}

impl VectorIndex {
    /// Create an empty index for one embedding space.
    pub fn new(metadata: EmbeddingMetadata) -> Self {
        Self {
            metadata,
            documents: BTreeMap::new(),
        }
    }

    /// Build an index from existing vector documents.
    pub fn from_documents(
        metadata: EmbeddingMetadata,
        documents: Vec<VectorDocument>,
    ) -> Result<Self, VectorIndexError> {
        let mut index = Self::new(metadata);
        for document in documents {
            index.upsert_document(document)?;
        }
        Ok(index)
    }

    /// Return immutable embedding metadata for this index.
    pub fn metadata(&self) -> &EmbeddingMetadata {
        &self.metadata
    }

    /// Return number of stored vector documents.
    pub fn document_count(&self) -> usize {
        self.documents.len()
    }

    /// Insert or replace one document vector by id.
    pub fn upsert_document(&mut self, document: VectorDocument) -> Result<(), VectorIndexError> {
        validate_vector(
            &document.vector,
            self.metadata.dimension,
            Some(document.document_id.as_str()),
            "document",
        )?;
        self.documents.insert(document.document_id, document.vector);
        Ok(())
    }

    /// Remove one document by id, returning `true` when it existed.
    pub fn remove_document(&mut self, document_id: &str) -> bool {
        self.documents.remove(document_id).is_some()
    }

    /// Rank indexed vectors by cosine similarity to the query vector.
    pub fn search(&self, query: &VectorQuery) -> Result<Vec<VectorSearchHit>, VectorIndexError> {
        validate_vector(&query.vector, self.metadata.dimension, None, "query")?;

        let mut hits: Vec<VectorSearchHit> = self
            .documents
            .iter()
            .map(|(document_id, vector)| VectorSearchHit {
                document_id: document_id.clone(),
                score: cosine_similarity(&query.vector, vector).unwrap_or(0.0),
            })
            .collect();
        hits.sort_by(|left, right| {
            right
                .score
                .total_cmp(&left.score)
                .then_with(|| left.document_id.cmp(&right.document_id))
        });
        if hits.len() > query.limit {
            hits.truncate(query.limit);
        }
        Ok(hits)
    }
}

/// Failure while validating vectors or executing search.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VectorIndexError {
    /// Vector length does not match embedding metadata dimension.
    DimensionMismatch {
        /// Expected dimension from index metadata.
        expected: usize,
        /// Actual vector length.
        actual: usize,
        /// Context string: `"document"` or `"query"`.
        context: &'static str,
        /// Optional document id for document validation failures.
        document_id: Option<String>,
    },
    /// Vector contains `NaN` or `inf`.
    NonFiniteValue {
        /// Context string: `"document"` or `"query"`.
        context: &'static str,
        /// Optional document id for document validation failures.
        document_id: Option<String>,
        /// Index of the non-finite value.
        value_index: usize,
    },
}

impl fmt::Display for VectorIndexError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::DimensionMismatch {
                expected,
                actual,
                context,
                document_id,
            } => {
                if let Some(document_id) = document_id {
                    write!(
                        f,
                        "{context} vector for '{document_id}' has dimension {actual}, expected {expected}"
                    )
                } else {
                    write!(
                        f,
                        "{context} vector has dimension {actual}, expected {expected}"
                    )
                }
            }
            Self::NonFiniteValue {
                context,
                document_id,
                value_index,
            } => {
                if let Some(document_id) = document_id {
                    write!(
                        f,
                        "{context} vector for '{document_id}' contains non-finite value at index {value_index}"
                    )
                } else {
                    write!(
                        f,
                        "{context} vector contains non-finite value at index {value_index}"
                    )
                }
            }
        }
    }
}

impl Error for VectorIndexError {}

/// Compute cosine similarity for two vectors in the same space.
///
/// Returns `None` when dimensions differ or either vector has zero magnitude.
pub fn cosine_similarity(left: &[f32], right: &[f32]) -> Option<f32> {
    if left.len() != right.len() || left.is_empty() {
        return None;
    }

    let mut dot = 0.0_f32;
    let mut left_norm_sq = 0.0_f32;
    let mut right_norm_sq = 0.0_f32;
    for (left_value, right_value) in left.iter().zip(right.iter()) {
        dot += left_value * right_value;
        left_norm_sq += left_value * left_value;
        right_norm_sq += right_value * right_value;
    }

    if left_norm_sq == 0.0 || right_norm_sq == 0.0 {
        return None;
    }
    Some(dot / (left_norm_sq.sqrt() * right_norm_sq.sqrt()))
}

fn validate_vector(
    vector: &[f32],
    expected_dimension: usize,
    document_id: Option<&str>,
    context: &'static str,
) -> Result<(), VectorIndexError> {
    if vector.len() != expected_dimension {
        return Err(VectorIndexError::DimensionMismatch {
            expected: expected_dimension,
            actual: vector.len(),
            context,
            document_id: document_id.map(ToOwned::to_owned),
        });
    }
    for (value_index, value) in vector.iter().enumerate() {
        if !value.is_finite() {
            return Err(VectorIndexError::NonFiniteValue {
                context,
                document_id: document_id.map(ToOwned::to_owned),
                value_index,
            });
        }
    }
    Ok(())
}

fn accumulate_hashed_feature(vector: &mut [f32], feature: &[u8], weight: f32) {
    if vector.is_empty() {
        return;
    }
    let digest = Sha256::digest(feature);
    let primary =
        u32::from_le_bytes([digest[0], digest[1], digest[2], digest[3]]) as usize % vector.len();
    let secondary =
        u32::from_le_bytes([digest[4], digest[5], digest[6], digest[7]]) as usize % vector.len();
    let sign = if digest[8] & 1 == 0 { 1.0 } else { -1.0 };
    vector[primary] += weight * sign;
    vector[secondary] += (weight * 0.5) * -sign;
}
