use super::{StorageBackend, StorageError, SnapshotQuery, FlushResult};
use crate::models::{DecisionSnapshot, Snapshot, SnapshotType};
use rusqlite::{Connection, params, OptionalExtension};
use std::path::Path;
use std::sync::{Arc, Mutex};
use tokio::task;
use chrono::{DateTime, Utc};
use flate2::{Compression, write::GzEncoder, read::GzDecoder};
use std::io::{Write, Read};

#[derive(Debug, Clone, PartialEq)]
pub enum CompressionType {
    None,
    Gzip,
    Zstd,
}

pub struct SqliteBackend {
    conn: Arc<Mutex<Connection>>,
    compression: CompressionType,
}

impl SqliteBackend {
    /// Create or open a SQLite database at the given path
    pub fn new(path: impl AsRef<Path>) -> Result<Self, StorageError> {
        let conn = Connection::open(path)
            .map_err(|e| StorageError::ConnectionError(format!("Failed to open database: {}", e)))?;

        let backend = Self {
            conn: Arc::new(Mutex::new(conn)),
            compression: CompressionType::Gzip,
        };

        // Run migrations
        {
            let conn_guard = backend.conn.lock().unwrap();
            Self::run_migrations(&conn_guard)?;
        }

        Ok(backend)
    }

    /// Create an in-memory database (for testing)
    pub fn in_memory() -> Result<Self, StorageError> {
        let conn = Connection::open(":memory:")
            .map_err(|e| StorageError::ConnectionError(format!("Failed to create in-memory database: {}", e)))?;

        let backend = Self {
            conn: Arc::new(Mutex::new(conn)),
            compression: CompressionType::Gzip,
        };

        // Run migrations
        {
            let conn_guard = backend.conn.lock().unwrap();
            Self::run_migrations(&conn_guard)?;
        }

        Ok(backend)
    }

    /// Create with custom compression
    pub fn with_compression(path: impl AsRef<Path>, compression: CompressionType) -> Result<Self, StorageError> {
        let mut backend = Self::new(path)?;
        backend.compression = compression;
        Ok(backend)
    }

    /// Run database migrations
    fn run_migrations(conn: &Connection) -> Result<(), StorageError> {
        // Enable WAL mode for better concurrent access
        conn.execute("PRAGMA journal_mode=WAL", [])
            .map_err(|e| StorageError::ConnectionError(format!("Failed to set WAL mode: {}", e)))?;

        // Enable foreign keys
        conn.execute("PRAGMA foreign_keys=ON", [])
            .map_err(|e| StorageError::ConnectionError(format!("Failed to enable foreign keys: {}", e)))?;

        // Create snapshots table
        conn.execute(
            r#"
            CREATE TABLE IF NOT EXISTS snapshots (
                id TEXT PRIMARY KEY,
                snapshot_type TEXT NOT NULL,
                metadata_json BLOB NOT NULL,
                data_json BLOB NOT NULL,
                created_at DATETIME NOT NULL,
                created_by TEXT,
                checksum TEXT,
                INDEX(created_at),
                INDEX(snapshot_type),
                INDEX(created_by)
            )
            "#,
            [],
        ).map_err(|e| StorageError::ConnectionError(format!("Failed to create snapshots table: {}", e)))?;

        // Create decisions table
        conn.execute(
            r#"
            CREATE TABLE IF NOT EXISTS decisions (
                id TEXT PRIMARY KEY,
                snapshot_id TEXT,
                function_name TEXT NOT NULL,
                module_name TEXT,
                data_json BLOB NOT NULL,
                created_at DATETIME NOT NULL,
                execution_time_ms REAL,
                error TEXT,
                error_type TEXT,
                INDEX(snapshot_id),
                INDEX(function_name),
                INDEX(module_name),
                INDEX(created_at),
                FOREIGN KEY(snapshot_id) REFERENCES snapshots(id) ON DELETE CASCADE
            )
            "#,
            [],
        ).map_err(|e| StorageError::ConnectionError(format!("Failed to create decisions table: {}", e)))?;

        // Create tags table for decision tags
        conn.execute(
            r#"
            CREATE TABLE IF NOT EXISTS tags (
                decision_id TEXT NOT NULL,
                key TEXT NOT NULL,
                value TEXT NOT NULL,
                PRIMARY KEY(decision_id, key),
                FOREIGN KEY(decision_id) REFERENCES decisions(id) ON DELETE CASCADE,
                INDEX(key),
                INDEX(value)
            )
            "#,
            [],
        ).map_err(|e| StorageError::ConnectionError(format!("Failed to create tags table: {}", e)))?;

        // Create model_parameters table
        conn.execute(
            r#"
            CREATE TABLE IF NOT EXISTS model_parameters (
                decision_id TEXT NOT NULL,
                model_name TEXT NOT NULL,
                model_version TEXT,
                provider TEXT,
                PRIMARY KEY(decision_id),
                FOREIGN KEY(decision_id) REFERENCES decisions(id) ON DELETE CASCADE,
                INDEX(model_name),
                INDEX(provider)
            )
            "#,
            [],
        ).map_err(|e| StorageError::ConnectionError(format!("Failed to create model_parameters table: {}", e)))?;

        Ok(())
    }

    fn compress(&self, data: &[u8]) -> Vec<u8> {
        match self.compression {
            CompressionType::None => data.to_vec(),
            CompressionType::Gzip => {
                let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
                encoder.write_all(data).unwrap();
                encoder.finish().unwrap()
            }
            CompressionType::Zstd => {
                zstd::bulk::compress(data, 3).unwrap_or_else(|_| data.to_vec())
            }
        }
    }

    fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, StorageError> {
        match self.compression {
            CompressionType::None => Ok(data.to_vec()),
            CompressionType::Gzip => {
                let mut decoder = GzDecoder::new(data);
                let mut result = Vec::new();
                decoder.read_to_end(&mut result)
                    .map_err(|e| StorageError::SerializationError(format!("Failed to decompress: {}", e)))?;
                Ok(result)
            }
            CompressionType::Zstd => {
                zstd::bulk::decompress(data, 10 * 1024 * 1024) // 10MB limit
                    .map_err(|e| StorageError::SerializationError(format!("Failed to decompress: {}", e)))
            }
        }
    }

    fn serialize_and_compress<T: serde::Serialize>(&self, data: &T) -> Result<Vec<u8>, StorageError> {
        let json_bytes = serde_json::to_vec(data)?;
        Ok(self.compress(&json_bytes))
    }

    fn decompress_and_deserialize<T: serde::de::DeserializeOwned>(&self, data: &[u8]) -> Result<T, StorageError> {
        let json_bytes = self.decompress(data)?;
        let result = serde_json::from_slice(&json_bytes)?;
        Ok(result)
    }
}

#[async_trait::async_trait]
impl StorageBackend for SqliteBackend {
    async fn save(&self, snapshot: &Snapshot) -> Result<String, StorageError> {
        let snapshot_id = snapshot.metadata.snapshot_id.to_string();
        let snapshot_clone = snapshot.clone();
        let compression = self.compression.clone();
        let conn = Arc::clone(&self.conn);

        task::spawn_blocking(move || {
            let backend = SqliteBackend { conn: conn.clone(), compression };
            let conn_guard = conn.lock().unwrap();

            // Serialize snapshot
            let metadata_blob = backend.serialize_and_compress(&snapshot_clone.metadata)?;
            let data_blob = backend.serialize_and_compress(&snapshot_clone)?;

            // Insert snapshot
            conn_guard.execute(
                r#"
                INSERT OR REPLACE INTO snapshots (
                    id, snapshot_type, metadata_json, data_json,
                    created_at, created_by, checksum
                ) VALUES (?, ?, ?, ?, ?, ?, ?)
                "#,
                params![
                    snapshot_id,
                    format!("{:?}", snapshot_clone.snapshot_type),
                    metadata_blob,
                    data_blob,
                    snapshot_clone.metadata.timestamp.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
                    snapshot_clone.metadata.created_by,
                    snapshot_clone.metadata.checksum,
                ],
            ).map_err(|e| StorageError::ConnectionError(format!("Failed to insert snapshot: {}", e)))?;

            // Save decisions
            for decision in &snapshot_clone.decisions {
                backend.save_decision_internal(&conn_guard, decision)?;
            }

            Ok(snapshot_id)
        }).await.unwrap()
    }

    async fn save_decision(&self, decision: &DecisionSnapshot) -> Result<String, StorageError> {
        let decision_id = decision.metadata.snapshot_id.to_string();
        let decision_clone = decision.clone();

        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();
            self.save_decision_internal(&conn_guard, &decision_clone)?;
            Ok(decision_id)
        }).await.unwrap()
    }

    async fn load(&self, snapshot_id: &str) -> Result<Snapshot, StorageError> {
        let id = snapshot_id.to_string();

        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            let row: Option<(Vec<u8>,)> = conn_guard.query_row(
                "SELECT data_json FROM snapshots WHERE id = ?",
                params![id],
                |row| Ok((row.get(0)?,))
            ).optional()
            .map_err(|e| StorageError::ConnectionError(format!("Failed to query snapshot: {}", e)))?;

            match row {
                Some((data_blob,)) => {
                    let snapshot: Snapshot = self.decompress_and_deserialize(&data_blob)?;
                    Ok(snapshot)
                }
                None => Err(StorageError::NotFound(format!("Snapshot {} not found", id)))
            }
        }).await.unwrap()
    }

    async fn load_decision(&self, decision_id: &str) -> Result<DecisionSnapshot, StorageError> {
        let id = decision_id.to_string();

        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            let row: Option<(Vec<u8>,)> = conn_guard.query_row(
                "SELECT data_json FROM decisions WHERE id = ?",
                params![id],
                |row| Ok((row.get(0)?,))
            ).optional()
            .map_err(|e| StorageError::ConnectionError(format!("Failed to query decision: {}", e)))?;

            match row {
                Some((data_blob,)) => {
                    let decision: DecisionSnapshot = self.decompress_and_deserialize(&data_blob)?;
                    Ok(decision)
                }
                None => Err(StorageError::NotFound(format!("Decision {} not found", id)))
            }
        }).await.unwrap()
    }

    async fn query(&self, query: SnapshotQuery) -> Result<Vec<Snapshot>, StorageError> {
        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            let mut sql = "SELECT data_json FROM snapshots WHERE 1=1".to_string();
            let mut params_vec: Vec<Box<dyn rusqlite::ToSql>> = Vec::new();

            // Build WHERE clause
            if let Some(start_time) = query.start_time {
                sql.push_str(" AND created_at >= ?");
                params_vec.push(Box::new(start_time.format("%Y-%m-%d %H:%M:%S%.3f").to_string()));
            }

            if let Some(end_time) = query.end_time {
                sql.push_str(" AND created_at <= ?");
                params_vec.push(Box::new(end_time.format("%Y-%m-%d %H:%M:%S%.3f").to_string()));
            }

            // Add function name filter by joining with decisions
            if query.function_name.is_some() || query.module_name.is_some() || query.model_name.is_some() || query.tags.is_some() {
                sql = format!(
                    "SELECT DISTINCT s.data_json FROM snapshots s
                     JOIN decisions d ON s.id = d.snapshot_id
                     LEFT JOIN model_parameters mp ON d.id = mp.decision_id
                     LEFT JOIN tags t ON d.id = t.decision_id
                     WHERE 1=1"
                );

                if let Some(function_name) = query.function_name {
                    sql.push_str(" AND d.function_name = ?");
                    params_vec.push(Box::new(function_name));
                }

                if let Some(module_name) = query.module_name {
                    sql.push_str(" AND d.module_name = ?");
                    params_vec.push(Box::new(module_name));
                }

                if let Some(model_name) = query.model_name {
                    sql.push_str(" AND mp.model_name = ?");
                    params_vec.push(Box::new(model_name));
                }

                if let Some(tags) = query.tags {
                    for (key, value) in tags {
                        sql.push_str(" AND EXISTS (SELECT 1 FROM tags t2 WHERE t2.decision_id = d.id AND t2.key = ? AND t2.value = ?)");
                        params_vec.push(Box::new(key));
                        params_vec.push(Box::new(value));
                    }
                }
            }

            // Add ordering and pagination
            sql.push_str(" ORDER BY created_at DESC");

            if let Some(limit) = query.limit {
                sql.push_str(" LIMIT ?");
                params_vec.push(Box::new(limit as i64));
            }

            if let Some(offset) = query.offset {
                sql.push_str(" OFFSET ?");
                params_vec.push(Box::new(offset as i64));
            }

            // Execute query
            let mut stmt = conn_guard.prepare(&sql)
                .map_err(|e| StorageError::InvalidQuery(format!("Invalid query: {}", e)))?;

            let param_refs: Vec<&dyn rusqlite::ToSql> = params_vec.iter().map(|p| p.as_ref()).collect();
            let rows = stmt.query_map(param_refs.as_slice(), |row| {
                Ok(row.get::<_, Vec<u8>>(0)?)
            }).map_err(|e| StorageError::ConnectionError(format!("Query failed: {}", e)))?;

            let mut snapshots = Vec::new();
            for row in rows {
                let data_blob = row.map_err(|e| StorageError::ConnectionError(format!("Row error: {}", e)))?;
                let snapshot: Snapshot = self.decompress_and_deserialize(&data_blob)?;
                snapshots.push(snapshot);
            }

            Ok(snapshots)
        }).await.unwrap()
    }

    async fn delete(&self, snapshot_id: &str) -> Result<bool, StorageError> {
        let id = snapshot_id.to_string();

        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            let rows_affected = conn_guard.execute(
                "DELETE FROM snapshots WHERE id = ?",
                params![id],
            ).map_err(|e| StorageError::ConnectionError(format!("Failed to delete snapshot: {}", e)))?;

            Ok(rows_affected > 0)
        }).await.unwrap()
    }

    async fn flush(&self) -> Result<FlushResult, StorageError> {
        // SQLite doesn't need explicit flushing, but we can optimize
        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            // Force WAL checkpoint
            conn_guard.execute("PRAGMA wal_checkpoint(TRUNCATE)", [])
                .map_err(|e| StorageError::ConnectionError(format!("Failed to checkpoint WAL: {}", e)))?;

            // Get stats
            let snapshot_count: i64 = conn_guard.query_row(
                "SELECT COUNT(*) FROM snapshots",
                [],
                |row| row.get(0)
            ).unwrap_or(0);

            Ok(FlushResult {
                snapshots_written: snapshot_count as usize,
                bytes_written: 0, // SQLite doesn't easily report this
                checkpoint_id: None,
            })
        }).await.unwrap()
    }

    async fn health_check(&self) -> Result<bool, StorageError> {
        task::spawn_blocking(move || {
            let conn_guard = self.conn.lock().unwrap();

            // Simple query to check connection
            let _: i64 = conn_guard.query_row("SELECT 1", [], |row| row.get(0))
                .map_err(|e| StorageError::ConnectionError(format!("Health check failed: {}", e)))?;

            Ok(true)
        }).await.unwrap()
    }
}

impl SqliteBackend {
    fn save_decision_internal(&self, conn: &Connection, decision: &DecisionSnapshot) -> Result<(), StorageError> {
        let decision_id = decision.metadata.snapshot_id.to_string();
        let data_blob = self.serialize_and_compress(decision)?;

        // Insert decision
        conn.execute(
            r#"
            INSERT OR REPLACE INTO decisions (
                id, function_name, module_name, data_json,
                created_at, execution_time_ms, error, error_type
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            "#,
            params![
                decision_id,
                decision.function_name,
                decision.module_name,
                data_blob,
                decision.metadata.timestamp.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
                decision.execution_time_ms,
                decision.error,
                decision.error_type,
            ],
        ).map_err(|e| StorageError::ConnectionError(format!("Failed to insert decision: {}", e)))?;

        // Save tags
        for (key, value) in &decision.tags {
            conn.execute(
                "INSERT OR REPLACE INTO tags (decision_id, key, value) VALUES (?, ?, ?)",
                params![decision_id, key, value],
            ).map_err(|e| StorageError::ConnectionError(format!("Failed to insert tag: {}", e)))?;
        }

        // Save model parameters
        if let Some(model_params) = &decision.model_parameters {
            conn.execute(
                r#"
                INSERT OR REPLACE INTO model_parameters (
                    decision_id, model_name, model_version, provider
                ) VALUES (?, ?, ?, ?)
                "#,
                params![
                    decision_id,
                    model_params.model_name,
                    model_params.model_version,
                    model_params.provider,
                ],
            ).map_err(|e| StorageError::ConnectionError(format!("Failed to insert model parameters: {}", e)))?;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::*;
    use tempfile::tempdir;
    use serde_json::json;

    async fn create_test_snapshot() -> Snapshot {
        let input = Input::new("test_input", json!("value"), "string");
        let output = Output::new("test_output", json!("result"), "string");
        let model_params = ModelParameters::new("gpt-4");

        let decision = DecisionSnapshot::new("test_function")
            .with_module("test_module")
            .add_input(input)
            .add_output(output)
            .with_model_parameters(model_params)
            .add_tag("env", "test");

        let mut snapshot = Snapshot::new(SnapshotType::Session);
        snapshot.add_decision(decision);
        snapshot
    }

    #[tokio::test]
    async fn test_sqlite_in_memory() {
        let backend = SqliteBackend::in_memory().unwrap();
        assert!(backend.health_check().await.unwrap());
    }

    #[tokio::test]
    async fn test_save_and_load_snapshot() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;

        let snapshot_id = backend.save(&snapshot).await.unwrap();
        let loaded_snapshot = backend.load(&snapshot_id).await.unwrap();

        assert_eq!(snapshot.decisions.len(), loaded_snapshot.decisions.len());
        assert_eq!(snapshot.snapshot_type, loaded_snapshot.snapshot_type);
    }

    #[tokio::test]
    async fn test_save_and_load_decision() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;
        let decision = snapshot.decisions[0].clone();

        let decision_id = backend.save_decision(&decision).await.unwrap();
        let loaded_decision = backend.load_decision(&decision_id).await.unwrap();

        assert_eq!(decision.function_name, loaded_decision.function_name);
        assert_eq!(decision.tags, loaded_decision.tags);
    }

    #[tokio::test]
    async fn test_query_by_function_name() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;
        backend.save(&snapshot).await.unwrap();

        let query = SnapshotQuery::new().with_function_name("test_function");
        let results = backend.query(query).await.unwrap();

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].decisions[0].function_name, "test_function");
    }

    #[tokio::test]
    async fn test_query_by_tags() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;
        backend.save(&snapshot).await.unwrap();

        let query = SnapshotQuery::new().with_tag("env", "test");
        let results = backend.query(query).await.unwrap();

        assert_eq!(results.len(), 1);
    }

    #[tokio::test]
    async fn test_delete_snapshot() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;

        let snapshot_id = backend.save(&snapshot).await.unwrap();
        assert!(backend.delete(&snapshot_id).await.unwrap());

        let result = backend.load(&snapshot_id).await;
        assert!(matches!(result, Err(StorageError::NotFound(_))));
    }

    #[tokio::test]
    async fn test_query_with_limit_offset() {
        let backend = SqliteBackend::in_memory().unwrap();

        // Save multiple snapshots
        for i in 0..5 {
            let mut snapshot = create_test_snapshot().await;
            snapshot.decisions[0].function_name = format!("test_function_{}", i);
            backend.save(&snapshot).await.unwrap();
        }

        let query = SnapshotQuery::new().with_limit(2).with_offset(1);
        let results = backend.query(query).await.unwrap();

        assert_eq!(results.len(), 2);
    }

    #[tokio::test]
    async fn test_compression() {
        let backend = SqliteBackend::in_memory().unwrap();
        let snapshot = create_test_snapshot().await;

        // Test that data is actually compressed
        let uncompressed = serde_json::to_vec(&snapshot).unwrap();
        let compressed = backend.compress(&uncompressed);

        // Gzip should reduce size for JSON data
        assert!(compressed.len() < uncompressed.len());

        // Test decompression
        let decompressed = backend.decompress(&compressed).unwrap();
        assert_eq!(uncompressed, decompressed);
    }

    #[tokio::test]
    async fn test_file_backend() {
        let temp_dir = tempdir().unwrap();
        let db_path = temp_dir.path().join("test.db");

        let backend = SqliteBackend::new(db_path).unwrap();
        let snapshot = create_test_snapshot().await;

        let snapshot_id = backend.save(&snapshot).await.unwrap();
        let loaded_snapshot = backend.load(&snapshot_id).await.unwrap();

        assert_eq!(snapshot.decisions.len(), loaded_snapshot.decisions.len());
    }
}