use std::path::Path;
use std::sync::Arc;
use anyhow::{Context, Result};
use arrow_array::{
    Float32Array, Int64Array, RecordBatch, RecordBatchIterator, StringArray,
    FixedSizeListArray, ArrayRef,
};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use futures::TryStreamExt;
use lancedb::query::{ExecutableQuery, QueryBase};
use tokio::sync::RwLock;

const TABLE_NAME: &str = "vectors";
const RECIPE_TABLE: &str = "recipes";
const EMBEDDING_DIM: i32 = 768;

/// LanceDB vector store for Tier 1 search cache.
/// Uses the embedded LanceDB engine for persistent vector storage.
pub struct LanceDb {
    conn: lancedb::Connection,
    table: RwLock<Option<lancedb::Table>>,
    recipe_table: RwLock<Option<lancedb::Table>>,
}

impl LanceDb {
    pub async fn new(path: &Path) -> Result<Self> {
        std::fs::create_dir_all(path)?;

        let uri = path.to_string_lossy().to_string();
        let conn = lancedb::connect(&uri)
            .execute()
            .await
            .context("Failed to connect to LanceDB")?;

        let db = Self {
            conn,
            table: RwLock::new(None),
            recipe_table: RwLock::new(None),
        };

        db.ensure_table().await?;
        db.ensure_recipe_table().await?;
        Ok(db)
    }

    fn schema() -> SchemaRef {
        Arc::new(Schema::new(vec![
            Field::new("file_path", DataType::Utf8, false),
            Field::new("content", DataType::Utf8, true),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    EMBEDDING_DIM,
                ),
                true,
            ),
            Field::new("indexed_at", DataType::Int64, false),
        ]))
    }

    async fn ensure_table(&self) -> Result<()> {
        let mut table_guard = self.table.write().await;

        // Try to open existing table first
        match self.conn.open_table(TABLE_NAME).execute().await {
            Ok(t) => {
                *table_guard = Some(t);
                return Ok(());
            }
            Err(_) => {
                // Table doesn't exist, create it with an empty batch
                let schema = Self::schema();
                let table = self.conn
                    .create_empty_table(TABLE_NAME, schema)
                    .execute()
                    .await
                    .context("Failed to create vectors table")?;
                *table_guard = Some(table);
            }
        }
        Ok(())
    }

    async fn get_table(&self) -> Result<lancedb::Table> {
        let table_guard = self.table.read().await;
        table_guard.clone().context("LanceDB table not initialized")
    }

    /// Build a FixedSizeListArray from a flat Float32Array.
    fn make_vector_array(data: Vec<f32>) -> Result<FixedSizeListArray> {
        let float_arr = Float32Array::from(data);
        let field = Arc::new(Field::new("item", DataType::Float32, true));
        FixedSizeListArray::try_new(field, EMBEDDING_DIM, Arc::new(float_arr) as ArrayRef, None)
            .context("Failed to create vector array")
    }

    /// Search for similar vectors using cosine similarity.
    pub async fn search(&self, query_embedding: &[f32], limit: usize) -> Result<Vec<VectorSearchResult>> {
        let table = self.get_table().await?;

        let count = table.count_rows(None).await.unwrap_or(0);
        if count == 0 {
            return Ok(Vec::new());
        }

        let query_vec: Vec<f32> = if query_embedding.len() == EMBEDDING_DIM as usize {
            query_embedding.to_vec()
        } else {
            // Pad or truncate to match dimension
            let mut v = vec![0.0f32; EMBEDDING_DIM as usize];
            let copy_len = query_embedding.len().min(EMBEDDING_DIM as usize);
            v[..copy_len].copy_from_slice(&query_embedding[..copy_len]);
            v
        };

        let results: Vec<RecordBatch> = table
            .vector_search(query_vec.as_slice())?
            .limit(limit)
            .execute()
            .await?
            .try_collect()
            .await?;

        let mut search_results = Vec::new();
        for batch in &results {
            let file_paths = batch
                .column_by_name("file_path")
                .and_then(|c| c.as_any().downcast_ref::<StringArray>());
            let contents = batch
                .column_by_name("content")
                .and_then(|c| c.as_any().downcast_ref::<StringArray>());
            let distances = batch
                .column_by_name("_distance")
                .and_then(|c| c.as_any().downcast_ref::<Float32Array>());

            if let (Some(fps), Some(cts)) = (file_paths, contents) {
                for i in 0..batch.num_rows() {
                    let file_path = fps.value(i).to_string();
                    let content = cts.value(i).to_string();
                    // LanceDB returns distance (lower = more similar for cosine)
                    // Convert to similarity score: 1.0 - distance
                    let score = distances
                        .map(|d| 1.0 - d.value(i) as f64)
                        .unwrap_or(0.0);

                    search_results.push(VectorSearchResult {
                        file_path,
                        content,
                        score,
                    });
                }
            }
        }

        // Sort by score descending
        search_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
        search_results.truncate(limit);
        Ok(search_results)
    }

    /// Upsert a document with its embedding.
    pub async fn upsert(&self, file_path: &str, content: &str, embedding: Vec<f32>) -> Result<()> {
        let table = self.get_table().await?;

        // Delete existing row if present
        let predicate = format!("file_path = '{}'", file_path.replace('\'', "''"));
        if let Err(e) = table.delete(&predicate).await {
            tracing::warn!(file_path, error = %e, "Failed to delete existing row before upsert");
        }

        // Prepare embedding vector (pad/truncate to EMBEDDING_DIM)
        let mut vec_data = vec![0.0f32; EMBEDDING_DIM as usize];
        let copy_len = embedding.len().min(EMBEDDING_DIM as usize);
        vec_data[..copy_len].copy_from_slice(&embedding[..copy_len]);

        let indexed_at = chrono::Utc::now().timestamp();

        // Build Arrow arrays
        let file_path_arr = StringArray::from(vec![file_path]);
        let content_arr = StringArray::from(vec![content]);
        let indexed_at_arr = Int64Array::from(vec![indexed_at]);
        let vector_arr = Self::make_vector_array(vec_data)?;

        let batch = RecordBatch::try_new(
            Self::schema(),
            vec![
                Arc::new(file_path_arr),
                Arc::new(content_arr),
                Arc::new(vector_arr),
                Arc::new(indexed_at_arr),
            ],
        )
        .context("Failed to create RecordBatch")?;

        let batches = RecordBatchIterator::new(
            vec![Ok(batch)],
            Self::schema(),
        );

        table.add(batches)
            .execute()
            .await
            .context("Failed to add data to LanceDB")?;

        Ok(())
    }

    /// Get the indexed timestamp for a file.
    pub async fn get_indexed_time(&self, file_path: &str) -> Result<i64> {
        let table = self.get_table().await?;

        let predicate = format!("file_path = '{}'", file_path.replace('\'', "''"));
        let results: Vec<RecordBatch> = table
            .query()
            .only_if(predicate)
            .select(lancedb::query::Select::columns(&["indexed_at"]))
            .limit(1)
            .execute()
            .await?
            .try_collect()
            .await?;

        for batch in &results {
            if batch.num_rows() > 0 {
                if let Some(col) = batch.column_by_name("indexed_at") {
                    if let Some(arr) = col.as_any().downcast_ref::<Int64Array>() {
                        return Ok(arr.value(0));
                    }
                }
            }
        }

        Ok(0)
    }

    /// Remove a document from the store.
    pub async fn remove(&self, file_path: &str) -> Result<()> {
        let table = self.get_table().await?;
        let predicate = format!("file_path = '{}'", file_path.replace('\'', "''"));
        if let Err(e) = table.delete(&predicate).await {
            tracing::warn!(file_path, error = %e, "Failed to delete row from LanceDB");
        }
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Recipe table
    // -----------------------------------------------------------------------

    fn recipe_schema() -> SchemaRef {
        Arc::new(Schema::new(vec![
            Field::new("name", DataType::Utf8, false),
            Field::new("content", DataType::Utf8, true),
            Field::new(
                "vector",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    EMBEDDING_DIM,
                ),
                true,
            ),
            Field::new("indexed_at", DataType::Int64, false),
        ]))
    }

    async fn ensure_recipe_table(&self) -> Result<()> {
        let mut table_guard = self.recipe_table.write().await;
        match self.conn.open_table(RECIPE_TABLE).execute().await {
            Ok(t) => {
                *table_guard = Some(t);
            }
            Err(_) => {
                let schema = Self::recipe_schema();
                let table = self.conn
                    .create_empty_table(RECIPE_TABLE, schema)
                    .execute()
                    .await
                    .context("Failed to create recipes table")?;
                *table_guard = Some(table);
            }
        }
        Ok(())
    }

    async fn get_recipe_table(&self) -> Result<lancedb::Table> {
        let guard = self.recipe_table.read().await;
        guard.clone().context("LanceDB recipe table not initialized")
    }

    /// Upsert a recipe embedding.
    pub async fn upsert_recipe(&self, name: &str, content: &str, embedding: Vec<f32>) -> Result<()> {
        let table = self.get_recipe_table().await?;

        // Delete existing row if present
        let predicate = format!("name = '{}'", name.replace('\'', "''"));
        if let Err(e) = table.delete(&predicate).await {
            tracing::warn!(name = %name, error = %e, "Failed to delete existing recipe before upsert");
        }

        let mut vec_data = vec![0.0f32; EMBEDDING_DIM as usize];
        let copy_len = embedding.len().min(EMBEDDING_DIM as usize);
        vec_data[..copy_len].copy_from_slice(&embedding[..copy_len]);

        let indexed_at = chrono::Utc::now().timestamp();

        let name_arr = StringArray::from(vec![name]);
        let content_arr = StringArray::from(vec![content]);
        let indexed_at_arr = Int64Array::from(vec![indexed_at]);
        let vector_arr = Self::make_vector_array(vec_data)?;

        let batch = RecordBatch::try_new(
            Self::recipe_schema(),
            vec![
                Arc::new(name_arr),
                Arc::new(content_arr),
                Arc::new(vector_arr),
                Arc::new(indexed_at_arr),
            ],
        )
        .context("Failed to create recipe RecordBatch")?;

        let batches = RecordBatchIterator::new(
            vec![Ok(batch)],
            Self::recipe_schema(),
        );

        table.add(batches)
            .execute()
            .await
            .context("Failed to add recipe to LanceDB")?;

        Ok(())
    }

    /// Search recipes by semantic similarity.
    pub async fn search_recipes(&self, query_embedding: &[f32], limit: usize) -> Result<Vec<VectorSearchResult>> {
        let table = self.get_recipe_table().await?;

        let count = table.count_rows(None).await.unwrap_or(0);
        if count == 0 {
            return Ok(Vec::new());
        }

        let query_vec: Vec<f32> = if query_embedding.len() == EMBEDDING_DIM as usize {
            query_embedding.to_vec()
        } else {
            let mut v = vec![0.0f32; EMBEDDING_DIM as usize];
            let copy_len = query_embedding.len().min(EMBEDDING_DIM as usize);
            v[..copy_len].copy_from_slice(&query_embedding[..copy_len]);
            v
        };

        let results: Vec<RecordBatch> = table
            .vector_search(query_vec.as_slice())?
            .limit(limit)
            .execute()
            .await?
            .try_collect()
            .await?;

        let mut search_results = Vec::new();
        for batch in &results {
            let names = batch
                .column_by_name("name")
                .and_then(|c| c.as_any().downcast_ref::<StringArray>());
            let contents = batch
                .column_by_name("content")
                .and_then(|c| c.as_any().downcast_ref::<StringArray>());
            let distances = batch
                .column_by_name("_distance")
                .and_then(|c| c.as_any().downcast_ref::<Float32Array>());

            if let (Some(ns), Some(cts)) = (names, contents) {
                for i in 0..batch.num_rows() {
                    let name = ns.value(i).to_string();
                    let content = cts.value(i).to_string();
                    let score = distances
                        .map(|d| 1.0 - d.value(i) as f64)
                        .unwrap_or(0.0);

                    search_results.push(VectorSearchResult {
                        file_path: name,
                        content,
                        score,
                    });
                }
            }
        }

        search_results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
        search_results.truncate(limit);
        Ok(search_results)
    }

    /// Remove a recipe from the vector store.
    pub async fn remove_recipe(&self, name: &str) -> Result<()> {
        let table = self.get_recipe_table().await?;
        let predicate = format!("name = '{}'", name.replace('\'', "''"));
        if let Err(e) = table.delete(&predicate).await {
            tracing::warn!(name = %name, error = %e, "Failed to delete recipe from LanceDB");
        }
        Ok(())
    }

    /// Get document count.
    pub async fn count(&self) -> usize {
        if let Ok(table) = self.get_table().await {
            table.count_rows(None).await.unwrap_or(0)
        } else {
            0
        }
    }
}

#[derive(Debug, Clone)]
pub struct VectorSearchResult {
    pub file_path: String,
    pub content: String,
    pub score: f64,
}
