use rusqlite::Connection;
use std::path::Path;
use std::sync::Mutex;

use crate::agent::AgentId;
use crate::error::SwarmResult;
use crate::message::Message;

pub struct MemoryStore {
    conn: Mutex<Connection>,
}

impl MemoryStore {
    pub fn open(path: &Path) -> SwarmResult<Self> {
        let conn = Connection::open(path)?;
        let store = Self {
            conn: Mutex::new(conn),
        };
        store.init_tables()?;
        Ok(store)
    }

    pub fn in_memory() -> SwarmResult<Self> {
        let conn = Connection::open_in_memory()?;
        let store = Self {
            conn: Mutex::new(conn),
        };
        store.init_tables()?;
        Ok(store)
    }

    fn init_tables(&self) -> SwarmResult<()> {
        let conn = self.conn.lock().unwrap();
        conn.execute_batch(
            "
            CREATE TABLE IF NOT EXISTS messages (
                id TEXT PRIMARY KEY,
                agent_id TEXT NOT NULL,
                kind TEXT NOT NULL,
                content TEXT NOT NULL,
                metadata TEXT,
                timestamp TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS agent_state (
                agent_id TEXT PRIMARY KEY,
                state TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_messages_agent
                ON messages(agent_id, timestamp);
            ",
        )?;
        Ok(())
    }

    pub fn store_message(&self, agent_id: &AgentId, msg: &Message) -> SwarmResult<()> {
        let conn = self.conn.lock().unwrap();
        conn.execute(
            "INSERT OR REPLACE INTO messages (id, agent_id, kind, content, metadata, timestamp)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            (
                msg.id.to_string(),
                &agent_id.0,
                serde_json::to_string(&msg.kind)?,
                &msg.content,
                serde_json::to_string(&msg.metadata)?,
                msg.timestamp.to_rfc3339(),
            ),
        )?;
        Ok(())
    }

    pub fn get_history(&self, agent_id: &AgentId, limit: usize) -> SwarmResult<Vec<Message>> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare(
            "SELECT id, kind, content, metadata, timestamp
             FROM messages WHERE agent_id = ?1
             ORDER BY timestamp DESC LIMIT ?2",
        )?;

        let messages = stmt
            .query_map((&agent_id.0, limit), |row| {
                let id_str: String = row.get(0)?;
                let kind_str: String = row.get(1)?;
                let content: String = row.get(2)?;
                let metadata_str: String = row.get(3)?;
                let ts_str: String = row.get(4)?;
                Ok((id_str, kind_str, content, metadata_str, ts_str))
            })?
            .filter_map(|row| {
                let (id_str, kind_str, content, metadata_str, ts_str) = row.ok()?;
                Some(Message {
                    id: id_str.parse().ok()?,
                    kind: serde_json::from_str(&kind_str).ok()?,
                    from: agent_id.clone(),
                    to: None,
                    content,
                    tool_calls: vec![],
                    tool_results: vec![],
                    metadata: serde_json::from_str(&metadata_str).unwrap_or_default(),
                    timestamp: chrono::DateTime::parse_from_rfc3339(&ts_str)
                        .ok()?
                        .with_timezone(&chrono::Utc),
                })
            })
            .collect::<Vec<_>>();

        Ok(messages)
    }

    pub fn save_agent_state(
        &self,
        agent_id: &AgentId,
        state: &serde_json::Value,
    ) -> SwarmResult<()> {
        let conn = self.conn.lock().unwrap();
        conn.execute(
            "INSERT OR REPLACE INTO agent_state (agent_id, state, updated_at)
             VALUES (?1, ?2, ?3)",
            (
                &agent_id.0,
                serde_json::to_string(state)?,
                chrono::Utc::now().to_rfc3339(),
            ),
        )?;
        Ok(())
    }

    pub fn load_agent_state(&self, agent_id: &AgentId) -> SwarmResult<Option<serde_json::Value>> {
        let conn = self.conn.lock().unwrap();
        let result = conn.query_row(
            "SELECT state FROM agent_state WHERE agent_id = ?1",
            [&agent_id.0],
            |row| {
                let state_str: String = row.get(0)?;
                Ok(state_str)
            },
        );

        match result {
            Ok(state_str) => Ok(Some(serde_json::from_str(&state_str)?)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }
}
