use crate::traits::{Memory, MemoryCategory, MemoryEntry};
use async_trait::async_trait;
use chrono::Utc;
use rusqlite::{params, Connection};
use std::path::{Path, PathBuf};

pub struct SqliteMemory {
    db_path: PathBuf,
}

impl SqliteMemory {
    pub fn new(workspace_dir: &Path) -> anyhow::Result<Self> {
        let db_path = workspace_dir.join("memory").join("brain.db");
        if let Some(parent) = db_path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let this = Self { db_path };
        this.init_schema()?;
        Ok(this)
    }

    fn open(&self) -> anyhow::Result<Connection> {
        Ok(Connection::open(&self.db_path)?)
    }

    fn init_schema(&self) -> anyhow::Result<()> {
        let conn = self.open()?;
        conn.execute_batch(
            "
            CREATE TABLE IF NOT EXISTS memory_entries (
                id TEXT PRIMARY KEY,
                key TEXT NOT NULL UNIQUE,
                content TEXT NOT NULL,
                category TEXT NOT NULL,
                timestamp TEXT NOT NULL,
                session_id TEXT
            );
            CREATE INDEX IF NOT EXISTS idx_memory_key ON memory_entries(key);
            CREATE INDEX IF NOT EXISTS idx_memory_category ON memory_entries(category);
            CREATE INDEX IF NOT EXISTS idx_memory_session_id ON memory_entries(session_id);
            ",
        )?;
        Ok(())
    }

    fn to_category(raw: &str) -> MemoryCategory {
        match raw {
            "core" => MemoryCategory::Core,
            "daily" => MemoryCategory::Daily,
            "conversation" => MemoryCategory::Conversation,
            other => MemoryCategory::Custom(other.to_string()),
        }
    }
}

#[async_trait]
impl Memory for SqliteMemory {
    fn name(&self) -> &str {
        "sqlite"
    }

    async fn store(
        &self,
        key: &str,
        content: &str,
        category: MemoryCategory,
        session_id: Option<&str>,
    ) -> anyhow::Result<()> {
        let db_path = self.db_path.clone();
        let key = key.to_string();
        let content = content.to_string();
        let category = category.to_string();
        let session_id = session_id.map(|s| s.to_string());

        tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
            let conn = Connection::open(db_path)?;
            conn.execute(
                "
                INSERT INTO memory_entries (id, key, content, category, timestamp, session_id)
                VALUES (?1, ?2, ?3, ?4, ?5, ?6)
                ON CONFLICT(key)
                DO UPDATE SET
                    content=excluded.content,
                    category=excluded.category,
                    timestamp=excluded.timestamp,
                    session_id=excluded.session_id
                ",
                params![
                    uuid::Uuid::new_v4().to_string(),
                    key,
                    content,
                    category,
                    Utc::now().to_rfc3339(),
                    session_id,
                ],
            )?;
            Ok(())
        })
        .await??;

        Ok(())
    }

    async fn recall(
        &self,
        query: &str,
        limit: usize,
        session_id: Option<&str>,
    ) -> anyhow::Result<Vec<MemoryEntry>> {
        let db_path = self.db_path.clone();
        let query = query.to_string();
        let session = session_id.map(|s| s.to_string());

        tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<MemoryEntry>> {
            let conn = Connection::open(db_path)?;
            let mut out = Vec::new();
            let like = format!("%{}%", query);

            if let Some(session_id) = session {
                let mut stmt = conn.prepare(
                    "
                    SELECT id, key, content, category, timestamp, session_id
                    FROM memory_entries
                    WHERE (key LIKE ?1 OR content LIKE ?1)
                      AND session_id = ?2
                    ORDER BY timestamp DESC
                    LIMIT ?3
                    ",
                )?;
                let rows = stmt.query_map(params![like, session_id, limit as i64], |row| {
                    let category_raw: String = row.get(3)?;
                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        key: row.get(1)?,
                        content: row.get(2)?,
                        category: Self::to_category(&category_raw),
                        timestamp: row.get(4)?,
                        session_id: row.get(5)?,
                        score: None,
                    })
                })?;

                for row in rows {
                    out.push(row?);
                }
            } else {
                let mut stmt = conn.prepare(
                    "
                    SELECT id, key, content, category, timestamp, session_id
                    FROM memory_entries
                    WHERE key LIKE ?1 OR content LIKE ?1
                    ORDER BY timestamp DESC
                    LIMIT ?2
                    ",
                )?;
                let rows = stmt.query_map(params![like, limit as i64], |row| {
                    let category_raw: String = row.get(3)?;
                    Ok(MemoryEntry {
                        id: row.get(0)?,
                        key: row.get(1)?,
                        content: row.get(2)?,
                        category: Self::to_category(&category_raw),
                        timestamp: row.get(4)?,
                        session_id: row.get(5)?,
                        score: None,
                    })
                })?;

                for row in rows {
                    out.push(row?);
                }
            }

            Ok(out)
        })
        .await?
    }

    async fn get(&self, key: &str) -> anyhow::Result<Option<MemoryEntry>> {
        let db_path = self.db_path.clone();
        let key = key.to_string();

        tokio::task::spawn_blocking(move || -> anyhow::Result<Option<MemoryEntry>> {
            let conn = Connection::open(db_path)?;
            let mut stmt = conn.prepare(
                "SELECT id, key, content, category, timestamp, session_id FROM memory_entries WHERE key = ?1",
            )?;
            let mut rows = stmt.query(params![key])?;
            if let Some(row) = rows.next()? {
                let category_raw: String = row.get(3)?;
                return Ok(Some(MemoryEntry {
                    id: row.get(0)?,
                    key: row.get(1)?,
                    content: row.get(2)?,
                    category: Self::to_category(&category_raw),
                    timestamp: row.get(4)?,
                    session_id: row.get(5)?,
                    score: None,
                }));
            }
            Ok(None)
        })
        .await?
    }

    async fn list(
        &self,
        category: Option<&MemoryCategory>,
        session_id: Option<&str>,
    ) -> anyhow::Result<Vec<MemoryEntry>> {
        let db_path = self.db_path.clone();
        let category = category.map(|c| c.to_string());
        let session = session_id.map(|s| s.to_string());

        tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<MemoryEntry>> {
            let conn = Connection::open(db_path)?;
            let mut out = Vec::new();

            let sql = match (category.is_some(), session.is_some()) {
                (false, false) => "SELECT id, key, content, category, timestamp, session_id FROM memory_entries ORDER BY timestamp DESC",
                (true, false) => "SELECT id, key, content, category, timestamp, session_id FROM memory_entries WHERE category = ?1 ORDER BY timestamp DESC",
                (false, true) => "SELECT id, key, content, category, timestamp, session_id FROM memory_entries WHERE session_id = ?1 ORDER BY timestamp DESC",
                (true, true) => "SELECT id, key, content, category, timestamp, session_id FROM memory_entries WHERE category = ?1 AND session_id = ?2 ORDER BY timestamp DESC",
            };

            let mut stmt = conn.prepare(sql)?;
            let mut rows = match (category, session) {
                (None, None) => stmt.query(params![] )?,
                (Some(c), None) => stmt.query(params![c])?,
                (None, Some(s)) => stmt.query(params![s])?,
                (Some(c), Some(s)) => stmt.query(params![c, s])?,
            };

            while let Some(row) = rows.next()? {
                let category_raw: String = row.get(3)?;
                out.push(MemoryEntry {
                    id: row.get(0)?,
                    key: row.get(1)?,
                    content: row.get(2)?,
                    category: Self::to_category(&category_raw),
                    timestamp: row.get(4)?,
                    session_id: row.get(5)?,
                    score: None,
                });
            }

            Ok(out)
        })
        .await?
    }

    async fn forget(&self, key: &str) -> anyhow::Result<bool> {
        let db_path = self.db_path.clone();
        let key = key.to_string();
        tokio::task::spawn_blocking(move || -> anyhow::Result<bool> {
            let conn = Connection::open(db_path)?;
            let affected =
                conn.execute("DELETE FROM memory_entries WHERE key = ?1", params![key])?;
            Ok(affected > 0)
        })
        .await?
    }

    async fn count(&self) -> anyhow::Result<usize> {
        let db_path = self.db_path.clone();
        tokio::task::spawn_blocking(move || -> anyhow::Result<usize> {
            let conn = Connection::open(db_path)?;
            let count: i64 =
                conn.query_row("SELECT COUNT(*) FROM memory_entries", params![], |r| {
                    r.get(0)
                })?;
            Ok(count.max(0) as usize)
        })
        .await?
    }

    async fn health_check(&self) -> bool {
        self.open().is_ok()
    }
}
