use crate::types::{Draft, MessageHeader, PgpStatus};
use rusqlite::{params, Connection, OptionalExtension};
use std::error::Error;

pub struct Storage {
    path: String,
}

impl Storage {
    pub fn new(path: &str) -> Result<Self, Box<dyn Error + Send + Sync>> {
        let conn = Connection::open(path)?;

        // Enable WAL mode for concurrent read/write access.
        conn.execute_batch("PRAGMA journal_mode=WAL;")?;

        // Set a 5-second busy timeout.
        conn.busy_timeout(std::time::Duration::from_secs(5))?;

        // New messages table with folder support
        conn.execute(
            "CREATE TABLE IF NOT EXISTS messages_v2 (
                folder TEXT,
                uid INTEGER,
                subject TEXT,
                sender TEXT,
                recipients TEXT,
                date INTEGER,
                is_read INTEGER,
                has_attachments INTEGER,
                pgp_status TEXT,
                flag TEXT,
                PRIMARY KEY (folder, uid)
            )",
            [],
        )?;

        // Message bodies table for full offline cache
        conn.execute(
            "CREATE TABLE IF NOT EXISTS message_bodies (
                folder TEXT,
                uid INTEGER,
                raw_content BLOB,
                PRIMARY KEY (folder, uid)
            )",
            [],
        )?;

        conn.execute(
            "CREATE TABLE IF NOT EXISTS contacts (
                email TEXT PRIMARY KEY,
                display_name TEXT,
                public_key TEXT,
                fingerprint TEXT,
                trust_level TEXT DEFAULT 'Unknown',
                first_seen INTEGER,
                last_seen INTEGER
            )",
            [],
        )?;

        // Migrations for existing databases
        let _ = conn.execute(
            "ALTER TABLE messages_v2 ADD COLUMN recipients TEXT DEFAULT '[]'",
            [],
        );
        let _ = conn.execute(
            "ALTER TABLE messages_v2 ADD COLUMN is_read INTEGER DEFAULT 0",
            [],
        );
        let _ = conn.execute(
            "ALTER TABLE messages_v2 ADD COLUMN has_attachments INTEGER DEFAULT 0",
            [],
        );
        let _ = conn.execute(
            "ALTER TABLE messages_v2 ADD COLUMN pgp_status TEXT DEFAULT 'None'",
            [],
        );
        let _ = conn.execute("ALTER TABLE contacts ADD COLUMN fingerprint TEXT", []);
        let _ = conn.execute(
            "ALTER TABLE contacts ADD COLUMN trust_level TEXT DEFAULT 'Unknown'",
            [],
        );
        let _ = conn.execute("ALTER TABLE contacts ADD COLUMN first_seen INTEGER", []);
        let _ = conn.execute("ALTER TABLE contacts ADD COLUMN last_seen INTEGER", []);
        let _ = conn.execute("ALTER TABLE messages_v2 ADD COLUMN flag TEXT", []);
        let _ = conn.execute(
            "ALTER TABLE sync_state ADD COLUMN server_total INTEGER NOT NULL DEFAULT 0",
            [],
        );

        // sync_state table
        conn.execute(
            "CREATE TABLE IF NOT EXISTS sync_state (
                account TEXT,
                folder TEXT,
                uidvalidity INTEGER,
                last_uid INTEGER,
                PRIMARY KEY (account, folder)
            )",
            [],
        )?;

        // Normalize INBOX
        let _ = conn.execute(
            "UPDATE messages_v2 SET folder = 'INBOX' WHERE LOWER(folder) = 'inbox' AND folder != 'INBOX'",
            [],
        );
        let _ = conn.execute(
            "UPDATE message_bodies SET folder = 'INBOX' WHERE LOWER(folder) = 'inbox' AND folder != 'INBOX'",
            [],
        );
        let _ = conn.execute(
            "UPDATE sync_state SET folder = 'INBOX' WHERE LOWER(folder) = 'inbox' AND folder != 'INBOX'",
            [],
        );

        conn.execute(
            "CREATE TABLE IF NOT EXISTS pgp_identities (
                fingerprint TEXT PRIMARY KEY,
                email TEXT,
                public_key TEXT,
                private_key TEXT
            )",
            [],
        )?;

        conn.execute(
            "CREATE TABLE IF NOT EXISTS drafts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                recipient TEXT,
                subject TEXT,
                body TEXT,
                last_updated INTEGER,
                last_synced_at INTEGER DEFAULT 0,
                server_uid INTEGER
            )",
            [],
        )?;

        conn.execute(
            "CREATE TABLE IF NOT EXISTS draft_attachments (
                draft_id INTEGER,
                name TEXT,
                content_type TEXT,
                content BLOB,
                FOREIGN KEY (draft_id) REFERENCES drafts(id) ON DELETE CASCADE
            )",
            [],
        )?;

        Ok(Storage {
            path: path.to_string(),
        })
    }

    fn connect(&self) -> Result<Connection, Box<dyn Error + Send + Sync>> {
        Ok(Connection::open(&self.path)?)
    }

    pub fn save_draft(&self, draft: &Draft) -> Result<i64, Box<dyn Error + Send + Sync>> {
        let mut conn = self.connect()?;
        let tx = conn.transaction()?;

        let draft_id = if let Some(id) = draft.id {
            tx.execute(
                "UPDATE drafts SET recipient = ?1, subject = ?2, body = ?3, last_updated = ?4, last_synced_at = ?5, server_uid = ?6 WHERE id = ?7",
                params![draft.to, draft.subject, draft.body, draft.last_updated, draft.last_synced_at, draft.server_uid, id],
            )?;
            id
        } else {
            tx.execute(
                "INSERT INTO drafts (recipient, subject, body, last_updated, last_synced_at, server_uid) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
                params![draft.to, draft.subject, draft.body, draft.last_updated, draft.last_synced_at, draft.server_uid],
            )?;
            tx.last_insert_rowid()
        };

        // Update attachments
        tx.execute(
            "DELETE FROM draft_attachments WHERE draft_id = ?1",
            params![draft_id],
        )?;
        for att in &draft.attachments {
            tx.execute(
                "INSERT INTO draft_attachments (draft_id, name, content_type, content) VALUES (?1, ?2, ?3, ?4)",
                params![draft_id, att.name, att.content_type, att.content],
            )?;
        }

        tx.commit()?;
        Ok(draft_id)
    }

    pub fn get_drafts(&self) -> Result<Vec<Draft>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn
            .prepare("SELECT id, recipient, subject, body, last_updated, last_synced_at, server_uid FROM drafts")?;
        let rows = stmt.query_map([], |row| {
            let id: i64 = row.get(0)?;

            // Load attachments for this draft
            let mut att_stmt = conn.prepare(
                "SELECT name, content_type, content FROM draft_attachments WHERE draft_id = ?1",
            )?;
            let att_rows = att_stmt.query_map(params![id], |att_row| {
                let content: Vec<u8> = att_row.get(2)?;
                let size = content.len();
                Ok(crate::types::Attachment {
                    name: att_row.get(0)?,
                    content_type: att_row.get(1)?,
                    size,
                    content,
                })
            })?;

            let mut attachments = Vec::new();
            for att in att_rows {
                attachments.push(att?);
            }

            Ok(Draft {
                id: Some(id),
                to: row.get(1)?,
                subject: row.get(2)?,
                body: row.get(3)?,
                attachments,
                last_updated: row.get(4).unwrap_or(0),
                last_synced_at: row.get(5).unwrap_or(0),
                server_uid: row.get(6).ok(),
            })
        })?;

        let mut drafts = Vec::new();
        for draft in rows {
            drafts.push(draft?);
        }
        Ok(drafts)
    }

    pub fn delete_draft(&self, id: i64) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute("DELETE FROM drafts WHERE id = ?1", params![id])?;
        Ok(())
    }

    pub fn save_header(
        &self,
        folder: &str,
        header: &MessageHeader,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let recipients_json =
            serde_json::to_string(&header.to).unwrap_or_else(|_| "[]".to_string());

        conn.execute(
            "INSERT OR REPLACE INTO messages_v2 (folder, uid, subject, sender, recipients, date, is_read, has_attachments, pgp_status, flag)
             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
            params![
                folder,
                header.uid,
                header.subject,
                header.from,
                recipients_json,
                header.date,
                if header.is_read { 1 } else { 0 },
                if header.has_attachments { 1 } else { 0 },
                format!("{:?}", header.pgp_status),
                header.flag.clone(),
            ],
        )?;
        Ok(())
    }

    pub fn get_uids(
        &self,
        folder: &str,
    ) -> Result<std::collections::HashSet<u32>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn.prepare("SELECT uid FROM messages_v2 WHERE folder = ?1")?;
        let rows = stmt.query_map(rusqlite::params![folder], |row| row.get::<_, u32>(0))?;

        let mut uids = std::collections::HashSet::new();
        for uid in rows {
            if let Ok(u) = uid {
                uids.insert(u);
            }
        }
        Ok(uids)
    }

    pub fn get_headers(
        &self,
        folder: &str,
    ) -> Result<Vec<MessageHeader>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn.prepare("SELECT uid, subject, sender, recipients, date, is_read, has_attachments, pgp_status, flag FROM messages_v2 WHERE folder = ?1 ORDER BY date DESC")?;
        let rows = stmt.query_map(params![folder], |row| {
            let recipients_json: String = row.get(3)?;
            let to: Vec<String> = serde_json::from_str(&recipients_json).unwrap_or_default();
            let pgp_status_str: String = row.get(7)?;

            let pgp_status = match pgp_status_str.as_str() {
                "Encrypted" => PgpStatus::Encrypted,
                "Signed" => PgpStatus::Signed,
                "SignedAndEncrypted" => PgpStatus::SignedAndEncrypted,
                "InvalidSignature" => PgpStatus::InvalidSignature,
                "EncryptedAtRest" => PgpStatus::EncryptedAtRest,
                _ => PgpStatus::None,
            };

            Ok(MessageHeader {
                uid: row.get(0)?,
                subject: row.get(1)?,
                from: row.get(2)?,
                to,
                date: row.get(4)?,
                is_read: row.get::<_, i32>(5)? == 1,
                has_attachments: row.get::<_, i32>(6)? == 1,
                pgp_status,
                flag: row.get(8).unwrap_or(None),
            })
        })?;

        let mut headers = Vec::new();
        for header in rows {
            headers.push(header?);
        }
        Ok(headers)
    }

    pub fn delete_header(
        &self,
        folder: &str,
        uid: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "DELETE FROM messages_v2 WHERE folder = ?1 AND uid = ?2",
            params![folder, uid],
        )?;
        conn.execute(
            "DELETE FROM message_bodies WHERE folder = ?1 AND uid = ?2",
            params![folder, uid],
        )?;
        Ok(())
    }

    pub fn delete_headers(
        &self,
        folder: &str,
        uids: &[u32],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        if uids.is_empty() {
            return Ok(());
        }
        let conn = self.connect()?;
        let placeholders = (0..uids.len())
            .map(|i| format!("?{}", i + 2))
            .collect::<Vec<_>>()
            .join(",");
        let sql_messages = format!(
            "DELETE FROM messages_v2 WHERE folder = ?1 AND uid IN ({})",
            placeholders
        );
        let sql_bodies = format!(
            "DELETE FROM message_bodies WHERE folder = ?1 AND uid IN ({})",
            placeholders
        );
        {
            let mut stmt = conn.prepare(&sql_messages)?;
            let mut p: Vec<Box<dyn rusqlite::ToSql>> = Vec::with_capacity(uids.len() + 1);
            p.push(Box::new(folder.to_string()));
            for &uid in uids {
                p.push(Box::new(uid));
            }
            let refs: Vec<&dyn rusqlite::ToSql> = p.iter().map(|v| v.as_ref()).collect();
            stmt.execute(refs.as_slice())?;
        }
        {
            let mut stmt = conn.prepare(&sql_bodies)?;
            let mut p: Vec<Box<dyn rusqlite::ToSql>> = Vec::with_capacity(uids.len() + 1);
            p.push(Box::new(folder.to_string()));
            for &uid in uids {
                p.push(Box::new(uid));
            }
            let refs: Vec<&dyn rusqlite::ToSql> = p.iter().map(|v| v.as_ref()).collect();
            stmt.execute(refs.as_slice())?;
        }
        Ok(())
    }

    pub fn wipe_folder(&self, folder: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute("DELETE FROM messages_v2 WHERE folder = ?1", params![folder])?;
        conn.execute(
            "DELETE FROM message_bodies WHERE folder = ?1",
            params![folder],
        )?;
        Ok(())
    }

    pub fn update_read_status(
        &self,
        folder: &str,
        uid: u32,
        is_read: bool,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "UPDATE messages_v2 SET is_read = ?1 WHERE folder = ?2 AND uid = ?3",
            params![if is_read { 1 } else { 0 }, folder, uid],
        )?;
        Ok(())
    }

    pub fn mark_all_as_read(&self, folder: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "UPDATE messages_v2 SET is_read = 1 WHERE folder = ?1",
            params![folder],
        )?;
        Ok(())
    }

    pub fn update_message_flag(
        &self,
        folder: &str,
        uid: u32,
        flag: Option<String>,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "UPDATE messages_v2 SET flag = ?1 WHERE folder = ?2 AND uid = ?3",
            params![flag, folder, uid],
        )?;
        Ok(())
    }

    pub fn save_message_raw(
        &self,
        folder: &str,
        uid: u32,
        content: &[u8],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "INSERT OR REPLACE INTO message_bodies (folder, uid, raw_content) VALUES (?1, ?2, ?3)",
            params![folder, uid, content],
        )?;
        Ok(())
    }

    pub fn get_message_raw(
        &self,
        folder: &str,
        uid: u32,
    ) -> Result<Option<Vec<u8>>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt =
            conn.prepare("SELECT raw_content FROM message_bodies WHERE folder = ?1 AND uid = ?2")?;
        let result = stmt
            .query_row(params![folder, uid], |row| row.get(0))
            .optional()?;
        Ok(result)
    }

    pub fn save_contact(
        &self,
        contact: &crate::types::Contact,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "INSERT OR IGNORE INTO contacts (email, display_name, public_key, fingerprint, trust_level, first_seen, last_seen) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
            params![
                contact.email.to_lowercase(),
                contact.display_name,
                contact.public_key,
                contact.fingerprint,
                format!("{:?}", contact.trust_level),
                contact.first_seen,
                contact.last_seen
            ],
        )?;
        Ok(())
    }

    pub fn get_contacts(&self) -> Result<Vec<crate::types::Contact>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn.prepare(
            "SELECT email, display_name, public_key, fingerprint, trust_level, first_seen, last_seen FROM contacts",
        )?;
        let rows = stmt.query_map([], |row| {
            let trust_level_str: String = row.get(4)?;
            let trust_level = match trust_level_str.as_str() {
                "Trusted" => crate::types::TrustLevel::Trusted,
                "Established" => crate::types::TrustLevel::Established,
                "Warning" => crate::types::TrustLevel::Warning,
                _ => crate::types::TrustLevel::Unknown,
            };

            Ok(crate::types::Contact {
                email: row.get(0)?,
                display_name: row.get(1)?,
                public_key: row.get(2)?,
                fingerprint: row.get(3)?,
                trust_level,
                first_seen: row.get(5).unwrap_or(0),
                last_seen: row.get(6).unwrap_or(0),
            })
        })?;

        let mut contacts = Vec::new();
        for contact in rows {
            contacts.push(contact?);
        }
        Ok(contacts)
    }

    pub fn get_contact_by_email(
        &self,
        email: &str,
    ) -> Result<Option<crate::types::Contact>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt =
            conn.prepare("SELECT email, display_name, public_key, fingerprint, trust_level, first_seen, last_seen FROM contacts WHERE LOWER(email) = LOWER(?1)")?;
        let contact = stmt.query_row(params![email.to_lowercase()], |row| {
            let trust_level_str: String = row.get(4)?;
            let trust_level = match trust_level_str.as_str() {
                "Trusted" => crate::types::TrustLevel::Trusted,
                "Established" => crate::types::TrustLevel::Established,
                "Warning" => crate::types::TrustLevel::Warning,
                _ => crate::types::TrustLevel::Unknown,
            };

            Ok(crate::types::Contact {
                email: row.get(0)?,
                display_name: row.get(1)?,
                public_key: row.get(2)?,
                fingerprint: row.get(3)?,
                trust_level,
                first_seen: row.get(5).unwrap_or(0),
                last_seen: row.get(6).unwrap_or(0),
            })
        });

        match contact {
            Ok(c) => Ok(Some(c)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    pub fn delete_contact(&self, email: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "DELETE FROM contacts WHERE LOWER(email) = LOWER(?1)",
            params![email.to_lowercase()],
        )?;
        Ok(())
    }

    pub fn save_identity(
        &self,
        email: &str,
        keypair: &crate::types::PgpKeyPair,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "INSERT OR REPLACE INTO pgp_identities (fingerprint, email, public_key, private_key) VALUES (?1, ?2, ?3, ?4)",
            params![keypair.fingerprint, email.to_lowercase(), keypair.public_key, keypair.private_key],
        )?;
        Ok(())
    }

    pub fn list_identities(
        &self,
        email: &str,
    ) -> Result<Vec<crate::types::PgpKeyPair>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn.prepare(
            "SELECT public_key, private_key, fingerprint FROM pgp_identities WHERE LOWER(email) = LOWER(?1)",
        )?;
        let rows = stmt.query_map(params![email.to_lowercase()], |row| {
            let pub_key: String = row.get(0)?;
            let priv_key: String = row.get(1)?;
            let fingerprint: String = row.get(2)?;

            Ok(crate::types::PgpKeyPair {
                public_key: pub_key,
                private_key: priv_key,
                key_info: format!("Key: {}", fingerprint),
                fingerprint,
                user_ids: Vec::new(),
            })
        })?;

        let mut identities = Vec::new();
        for id in rows {
            identities.push(id?);
        }
        Ok(identities)
    }

    pub fn get_identity_by_fingerprint(
        &self,
        fingerprint: &str,
    ) -> Result<Option<crate::types::PgpKeyPair>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt =
            conn.prepare("SELECT public_key, private_key, fingerprint FROM pgp_identities WHERE fingerprint = ?1")?;
        let keypair = stmt.query_row(params![fingerprint], |row| {
            let pub_key: String = row.get(0)?;
            let priv_key: String = row.get(1)?;
            let fingerprint: String = row.get(2)?;

            Ok(crate::types::PgpKeyPair {
                public_key: pub_key,
                private_key: priv_key,
                key_info: format!("Key: {}", fingerprint),
                fingerprint,
                user_ids: Vec::new(),
            })
        });

        match keypair {
            Ok(k) => Ok(Some(k)),
            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    pub fn delete_identity(&self, fingerprint: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "DELETE FROM pgp_identities WHERE fingerprint = ?1",
            params![fingerprint],
        )?;
        Ok(())
    }

    pub fn update_sync_state(
        &self,
        account: &str,
        folder: &str,
        uidvalidity: u32,
        last_uid: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "INSERT OR REPLACE INTO sync_state (account, folder, uidvalidity, last_uid) VALUES (?1, ?2, ?3, ?4)",
            params![account, folder, uidvalidity, last_uid],
        )?;
        Ok(())
    }

    pub fn get_sync_state(
        &self,
        account: &str,
        folder: &str,
    ) -> Result<Option<(u32, u32)>, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let mut stmt = conn.prepare(
            "SELECT uidvalidity, last_uid FROM sync_state WHERE account = ?1 AND folder = ?2",
        )?;

        let mut rows = stmt.query(params![account, folder])?;

        if let Some(row) = rows.next()? {
            Ok(Some((row.get(0)?, row.get(1)?)))
        } else {
            Ok(None)
        }
    }

    pub fn get_unread_count(&self, folder: &str) -> Result<i32, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let count: i32 = conn.query_row(
            "SELECT COUNT(*) FROM messages_v2 WHERE folder = ?1 AND is_read = 0",
            params![folder],
            |row| row.get(0),
        )?;
        Ok(count)
    }

    pub fn update_server_total(
        &self,
        account: &str,
        folder: &str,
        total: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        conn.execute(
            "INSERT OR IGNORE INTO sync_state (account, folder, uidvalidity, last_uid, server_total) VALUES (?1, ?2, 0, 0, 0)",
            params![account, folder],
        )?;
        conn.execute(
            "UPDATE sync_state SET server_total = ?3 WHERE account = ?1 AND folder = ?2",
            params![account, folder, total],
        )?;
        Ok(())
    }

    pub fn get_folder_stats(
        &self,
        account: &str,
        folder: &str,
    ) -> Result<(i32, i32), Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let unread: i32 = conn.query_row(
            "SELECT COUNT(*) FROM messages_v2 WHERE folder = ?1 AND is_read = 0",
            params![folder],
            |row| row.get(0),
        )?;
        let server_total: i32 = conn
            .query_row(
                "SELECT COALESCE(server_total, 0) FROM sync_state WHERE account = ?1 AND folder = ?2",
                params![account, folder],
                |row| row.get(0),
            )
            .unwrap_or(0);
        Ok((unread, server_total))
    }

    pub fn get_total_unread_count(&self) -> Result<i32, Box<dyn Error + Send + Sync>> {
        let conn = self.connect()?;
        let count: i32 = conn.query_row(
            "SELECT COUNT(*) FROM messages_v2 WHERE is_read = 0 AND LOWER(folder) = 'inbox'",
            [],
            |row| row.get(0),
        )?;
        Ok(count)
    }
}
