use crate::types::{FolderInfo, FolderType, MailServerConfig};
use async_imap::extensions::idle::IdleResponse;
use async_imap::Client;
use futures::stream::TryStreamExt;
use mail_parser::MimeHeaders;
use std::error::Error;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio_rustls::client::TlsStream;
use tokio_rustls::TlsConnector;

/// Returned by `ImapClient::sync_mailbox`. Contains data needed for post-sync operations
/// that must run outside the shared session lock.
pub struct SyncResult {
    /// UIDs of messages already in local cache that need their FLAGS refreshed from the server.
    /// The caller should run `refresh_flags` on a dedicated connection after releasing the
    /// shared session mutex.
    pub uids_for_flag_refresh: Vec<u32>,
    /// Newly-fetched message headers (same as the previous `Vec<MessageHeader>` return value).
    pub new_headers: Vec<crate::types::MessageHeader>,
}

pub struct ImapClient {
    session: Option<async_imap::Session<TlsStream<TcpStream>>>,
    config: MailServerConfig,
    email: String,
}

impl ImapClient {
    pub fn new(config: MailServerConfig, email: String) -> Self {
        Self {
            session: None,
            config,
            email,
        }
    }

    pub async fn connect(&mut self, password: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        tracing::debug!(
            "ImapClient: Starting connection to {}:{}",
            self.config.imap_host,
            self.config.imap_port
        );

        let mut root_cert_store = rustls::RootCertStore::empty();
        root_cert_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());

        tracing::debug!("ImapClient: Loading native certs...");
        // Load native certs
        match rustls_native_certs::load_native_certs() {
            Ok(certs) => {
                for cert in certs {
                    let _ = root_cert_store.add(cert);
                }
            }
            Err(e) => {
                tracing::warn!("ImapClient: Failed to load native certs: {}", e);
            }
        }

        tracing::debug!("ImapClient: Building TLS config...");
        let config = rustls::ClientConfig::builder_with_provider(Arc::new(
            rustls::crypto::ring::default_provider(),
        ))
        .with_safe_default_protocol_versions()
        .map_err(|e| format!("Failed to set TLS protocol versions: {}", e))?
        .with_root_certificates(root_cert_store)
        .with_no_client_auth();

        let connector = TlsConnector::from(Arc::new(config));

        let dns_name = self.config.imap_host.clone();
        let port = self.config.imap_port;

        let addr = format!("{}:{}", dns_name, port);
        tracing::debug!("ImapClient: Connecting to TCP stream at {}...", addr);
        let stream = tokio::time::timeout(Duration::from_secs(15), TcpStream::connect(&addr))
            .await
            .map_err(|_| format!("Connection to {} timed out after 15 seconds", addr))?
            .map_err(|e| format!("TCP connection to {} failed: {}", addr, e))?;

        // Establish TLS
        let server_name = rustls_pki_types::ServerName::try_from(dns_name.clone())
            .map_err(|_| format!("Invalid DNS name: {}", dns_name))?
            .to_owned();

        tracing::debug!("ImapClient: Performing TLS handshake for {}...", dns_name);
        let tls_stream = connector.connect(server_name, stream).await?;

        // Create IMAP client
        tracing::debug!("ImapClient: Creating async_imap client...");
        let client = Client::new(tls_stream);

        // Login
        tracing::debug!("ImapClient: Attempting login for {}...", self.email);
        let session = client
            .login(&self.email, password)
            .await
            .map_err(|(e, _client)| {
                tracing::error!("ImapClient: Login failed: {}", e);
                e.to_string()
            })?;

        tracing::debug!("ImapClient: Successfully logged in as {}", self.email);
        self.session = Some(session);
        Ok(())
    }

    pub async fn reconnect(&mut self, password: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        // Force close existing session if any
        if let Some(mut session) = self.session.take() {
            let _ = session.logout().await;
        }
        self.connect(password).await
    }

    pub async fn logout(&mut self) -> Result<(), Box<dyn Error + Send + Sync>> {
        if let Some(mut session) = self.session.take() {
            session.logout().await.map_err(|e| e.to_string())?;
        }
        Ok(())
    }

    /// Checks connection health by sending a NOOP command.
    /// Returns true if healthy, false if broken.
    pub async fn keep_alive(&mut self) -> bool {
        if let Some(session) = self.session.as_mut() {
            match session.noop().await {
                Ok(_) => true,
                Err(e) => {
                    tracing::warn!("Session check failed for {}: {}", self.email, e);
                    false
                }
            }
        } else {
            false
        }
    }

    /// Use this before any operation to ensure we have a valid session.
    pub async fn ensure_connected(
        &mut self,
        password: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        if !self.keep_alive().await {
            tracing::debug!("Reconnecting session for {}", self.email);
            self.reconnect(password).await?;
        }
        Ok(())
    }

    pub async fn list_folders(
        &mut self,
        show_all: bool,
    ) -> Result<Vec<FolderInfo>, Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;

        let mut mailboxes = session
            .list(Some(""), Some("*"))
            .await
            .map_err(|e| e.to_string())?;

        let mut all_data = Vec::new();
        while let Some(mailbox) = mailboxes.try_next().await.map_err(|e| e.to_string())? {
            let mut is_selectable = true;
            let mut special_use_type: Option<FolderType> = None;

            // Check for special-use attributes (RFC 6154) which are reliable indicators
            for attr in mailbox.attributes() {
                match attr {
                    async_imap::types::NameAttribute::NoSelect => is_selectable = false,
                    async_imap::types::NameAttribute::All => {
                        special_use_type = Some(FolderType::Archive)
                    }
                    async_imap::types::NameAttribute::Drafts => {
                        special_use_type = Some(FolderType::Drafts)
                    }
                    async_imap::types::NameAttribute::Flagged => {} // Not a standard folder type for us
                    async_imap::types::NameAttribute::Junk => {
                        special_use_type = Some(FolderType::Junk)
                    }
                    async_imap::types::NameAttribute::Sent => {
                        special_use_type = Some(FolderType::Sent)
                    }
                    async_imap::types::NameAttribute::Trash => {
                        special_use_type = Some(FolderType::Trash)
                    }
                    async_imap::types::NameAttribute::Archive => {
                        special_use_type = Some(FolderType::Archive)
                    }
                    _ => {}
                }
            }

            // Keep the raw IMAP name for SELECT/EXAMINE commands.
            let imap_name = mailbox.name().to_string();
            // Decode RFC 3501 modified UTF-7 folder names (for Cyrillic, Arabic, CJK, etc.)
            let folder_name = super::utf7::decode_modified_utf7(mailbox.name());
            all_data.push((imap_name, folder_name, is_selectable, special_use_type));
        }

        // 1. Map to FolderInfo and determine type
        let mut folders = Vec::new();
        for (imap_name, name, is_selectable, special_use_type) in all_data {
            let upper = name.to_uppercase();
            let parts: Vec<&str> = upper.split('/').collect();
            let last_part = parts.last().cloned().unwrap_or("");

            let (folder_type, canonical_name, canonical_imap_name) = if let Some(sut) =
                special_use_type
            {
                // RFC 6154 special-use attribute takes priority
                (sut, name, imap_name)
            } else if upper == "INBOX" {
                (FolderType::Inbox, "INBOX".to_string(), "INBOX".to_string())
            } else if last_part == "SENT"
                || last_part == "SENT MAIL"
                || last_part == "SENT MESSAGES"
                || last_part == "SENT ITEMS"
            {
                (FolderType::Sent, name, imap_name)
            } else if last_part == "DRAFTS" || last_part == "DRAFT" {
                (FolderType::Drafts, name, imap_name)
            } else if last_part == "TRASH"
                || last_part == "BIN"
                || last_part == "DELETED"
                || last_part == "DELETED ITEMS"
                || last_part == "DELETED MESSAGES"
            {
                (FolderType::Trash, name, imap_name)
            } else if last_part == "JUNK" || last_part == "SPAM" || last_part == "JUNK EMAIL" {
                (FolderType::Junk, name, imap_name)
            } else if last_part == "ARCHIVE" || last_part == "ARCHIVES" || last_part == "ALL MAIL" {
                (FolderType::Archive, name, imap_name)
            } else {
                (FolderType::Custom, name, imap_name)
            };

            // Always include all folders; deduplication logic below handles standard folder duplicates
            folders.push(FolderInfo {
                name: canonical_name,
                imap_name: canonical_imap_name,
                folder_type,
                is_selectable,
                unread_count: 0,
            });
        }

        // 2. Multi-standard-folder deduplication (e.g. Gmail's [Gmail]/Sent vs Sent)
        // If multiple folders have the same non-Custom type, only keep the "best" one
        let mut best_typed_folders: std::collections::HashMap<FolderType, usize> =
            std::collections::HashMap::new();

        for (i, folder) in folders.iter().enumerate() {
            if folder.folder_type == FolderType::Custom {
                continue;
            }

            let entry = best_typed_folders
                .entry(folder.folder_type.clone())
                .or_insert(i);
            let current_best_name = &folders[*entry].name;

            // Heuristic: prefer shorter names or non-prefixed names for the "primary" type
            if folder.name.len() < current_best_name.len() {
                *entry = i;
            }
        }

        // Build final list: all custom folders + only the chosen primary ones
        let mut final_folders = Vec::new();
        let best_indices: std::collections::HashSet<usize> =
            best_typed_folders.values().cloned().collect();

        for (i, folder) in folders.into_iter().enumerate() {
            if folder.folder_type == FolderType::Custom {
                final_folders.push(folder);
            } else if best_indices.contains(&i) {
                final_folders.push(folder);
            } else if show_all {
                // Keep the duplicate as a Custom folder if show_all is on
                final_folders.push(FolderInfo {
                    name: folder.name,
                    imap_name: folder.imap_name,
                    folder_type: FolderType::Custom,
                    is_selectable: folder.is_selectable,
                    unread_count: 0,
                });
            }
        }

        Ok(final_folders)
    }

    pub async fn select_mailbox(
        &mut self,
        mailbox: &str,
    ) -> Result<(u32, u32), Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;
        let mb = session.select(mailbox).await.map_err(|e| e.to_string())?;
        Ok((mb.uid_validity.unwrap_or(0), mb.exists))
    }

    /// Search for UIDs strictly below `below_uid`. Returns results sorted descending (newest first).
    pub async fn search_uids_below(
        &mut self,
        below_uid: u32,
    ) -> Result<Vec<u32>, Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;
        let query = format!("UID 1:{}", below_uid - 1);
        let mut uids: Vec<u32> = session
            .uid_search(&query)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .collect();
        uids.sort_unstable_by(|a, b| b.cmp(a));
        Ok(uids)
    }

    pub async fn fetch_headers(
        &mut self,
        range: &str,
    ) -> Result<Vec<crate::types::MessageHeader>, Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;

        // Use BODY.PEEK[HEADER] instead of full BODY.PEEK[] to avoid massive memory/bandwidth usage
        let mut fetches = session
            .uid_fetch(range, "(UID FLAGS BODY.PEEK[HEADER])")
            .await
            .map_err(|e| e.to_string())?;

        let mut headers = Vec::new();
        while let Ok(Some(fetch)) = fetches.try_next().await {
            let uid = match fetch.uid {
                Some(u) => u,
                None => {
                    tracing::warn!("Fetch response missing UID, skipping message");
                    continue;
                }
            };

            // Use .header() if available, otherwise .body() as fallback
            let full_message_bytes = fetch
                .header()
                .unwrap_or_else(|| fetch.body().unwrap_or(&[]));
            if full_message_bytes.is_empty() {
                continue;
            }

            let preprocessed = preprocess_header_bytes(full_message_bytes);
            let parsed_opt = mail_parser::MessageParser::default().parse(&preprocessed);

            let mut from = "Unknown".to_string();
            let mut subject = "No Subject".to_string();
            let mut date = 0;
            let mut to = Vec::new();

            let mut has_attachments = false;
            let mut pgp_status = crate::types::PgpStatus::None;

            if let Some(message) = parsed_opt {
                from = message
                    .from()
                    .and_then(|fl| match fl {
                        mail_parser::Address::List(l) => l.first().cloned(),
                        mail_parser::Address::Group(g) => g
                            .iter()
                            .flat_map(|group| group.addresses.iter())
                            .next()
                            .cloned(),
                    })
                    .map(|addr| {
                        let name = addr.name().unwrap_or("").trim();
                        let address = addr.address().unwrap_or("").trim();
                        if name.is_empty() {
                            address.to_string()
                        } else {
                            format!("{} <{}>", name, address)
                        }
                    })
                    .unwrap_or_else(|| "Unknown".to_string());

                subject = message.subject().unwrap_or("No Subject").to_string();
                date = message.date().map(|dt| dt.to_timestamp()).unwrap_or(0);

                if let Some(to_addr) = message.to() {
                    match to_addr {
                        mail_parser::Address::List(l) => {
                            for addr in l {
                                let name = addr.name().unwrap_or("").trim();
                                let address = addr.address().unwrap_or("").trim();
                                if name.is_empty() {
                                    to.push(address.to_string());
                                } else {
                                    to.push(format!("{} <{}>", name, address));
                                }
                            }
                        }
                        mail_parser::Address::Group(groups) => {
                            for group in groups {
                                for addr in &group.addresses {
                                    let name = addr.name().unwrap_or("").trim();
                                    let address = addr.address().unwrap_or("").trim();
                                    if name.is_empty() {
                                        to.push(address.to_string());
                                    } else {
                                        to.push(format!("{} <{}>", name, address));
                                    }
                                }
                            }
                        }
                    }
                }

                if let Some(ct) = message.content_type() {
                    let ctype = ct.ctype();
                    let subtype = ct.subtype().unwrap_or("");
                    if ctype == "multipart" && subtype == "signed" {
                        pgp_status = crate::types::PgpStatus::Signed;
                    } else if ctype == "multipart" && subtype == "encrypted" {
                        pgp_status = crate::types::PgpStatus::Encrypted;
                    } else if ctype == "multipart" && subtype == "mixed" {
                        has_attachments = true;
                    }
                }
            }

            let mut message_flag: Option<String> = None;
            for f in fetch.flags() {
                match f {
                    async_imap::types::Flag::Flagged => {
                        if message_flag.is_none() {
                            message_flag = Some("red".to_string());
                        }
                    }
                    async_imap::types::Flag::Custom(c) => {
                        let c_lower = c.to_lowercase();
                        if c_lower == "korax-flag-red" {
                            message_flag = Some("red".to_string());
                        } else if c_lower == "korax-flag-orange" {
                            message_flag = Some("orange".to_string());
                        } else if c_lower == "korax-flag-yellow" {
                            message_flag = Some("yellow".to_string());
                        } else if c_lower == "korax-flag-green" {
                            message_flag = Some("green".to_string());
                        } else if c_lower == "korax-flag-blue" {
                            message_flag = Some("blue".to_string());
                        }
                    }
                    _ => {}
                }
            }

            headers.push(crate::types::MessageHeader {
                uid,
                subject,
                from,
                to,
                date,
                is_read: fetch
                    .flags()
                    .any(|f| matches!(f, async_imap::types::Flag::Seen)),
                has_attachments,
                pgp_status,
                flag: message_flag,
            });
        }

        Ok(headers)
    }

    pub async fn fetch_body_raw(
        &mut self,
        uid: u32,
    ) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;

        let mut fetches = session
            .uid_fetch(uid.to_string(), "(BODY.PEEK[])")
            .await
            .map_err(|e| e.to_string())?;

        while let Some(fetch) = fetches.try_next().await.map_err(|e| e.to_string())? {
            // Ensure this is the message we want (though UID fetch acts as filter)
            if let Some(fetched_uid) = fetch.uid {
                if fetched_uid != uid {
                    continue;
                }
            }

            let body = fetch.body().ok_or("No body")?;
            return Ok(body.to_vec());
        }

        Err(format!("Message with UID {} not found on server", uid).into())
    }

    pub async fn fetch_body(
        &mut self,
        uid: u32,
    ) -> Result<crate::types::MessageBody, Box<dyn Error + Send + Sync>> {
        let raw_body = self.fetch_body_raw(uid).await?;

        // Parse with mail-parser
        let message = mail_parser::MessageParser::default()
            .parse(&raw_body)
            .ok_or("Failed to parse MIME")?;

        let mut attachments = Vec::new();
        for part in message.attachments() {
            let content_type = part
                .content_type()
                .map(|ct| {
                    let ctype = ct.ctype();
                    let subtype = ct.subtype();
                    format!("{}/{}", ctype, subtype.unwrap_or(""))
                })
                .unwrap_or_else(|| "application/octet-stream".to_string());

            attachments.push(crate::types::Attachment {
                name: part.attachment_name().unwrap_or("unnamed").to_string(),
                content_type,
                size: part.contents().len(),
                content: part.contents().to_vec(),
            });
        }

        let from_addr = message.from().and_then(|f| {
            f.iter()
                .next()
                .map(|a| a.address().unwrap_or("").to_string())
        });

        Ok(crate::types::MessageBody {
            uid,
            from: from_addr,
            text_plain: message.body_text(0).map(|s| s.to_string()),
            text_html: message.body_html(0).map(|s| s.to_string()),
            attachments,
            pgp_status: crate::types::PgpStatus::None,
            signature_fingerprint: None,
        })
    }

    pub async fn sync_mailbox(
        &mut self,
        mailbox: &str,
        storage: &crate::storage::Storage,
        cache_days: i64,
    ) -> Result<SyncResult, Box<dyn Error + Send + Sync>> {
        let (current_validity, mailbox_exists) = self.select_mailbox(mailbox).await?;

        // 1. Get stored state & known UIDs
        let state = storage.get_sync_state(&self.email, mailbox)?;

        let last_known_uid = if let Some((stored_validity, stored_last_uid)) = state {
            if stored_validity == current_validity {
                stored_last_uid
            } else {
                // UIDVALIDITY changed: existing UIDs now refer to different messages.
                // Wipe the local cache before re-syncing from scratch.
                storage.wipe_folder(mailbox)?;
                0
            }
        } else {
            0
        };

        let local_uids = storage.get_uids(mailbox).unwrap_or_default();

        // 2. Fetch server UIDs
        let session = self.session.as_mut().ok_or("Not connected")?;

        let mut server_uids: std::collections::HashSet<u32> = std::collections::HashSet::new();
        let unseen_uids: std::collections::HashSet<u32>;
        let mut min_since_uid = 0;

        if cache_days > 0 {
            // Priority 1: Recent messages (e.g. last 30 days)
            let since_date = chrono::Utc::now() - chrono::Duration::days(cache_days);
            let since_query = format!("SINCE {}", since_date.format("%d-%b-%Y"));
            let since_uids: std::collections::HashSet<u32> = session
                .uid_search(&since_query)
                .await
                .map_err(|e| e.to_string())?
                .into_iter()
                .collect();

            min_since_uid = since_uids.iter().cloned().min().unwrap_or(0);
            server_uids.extend(since_uids.clone());

            // Priority 2: All unread messages (regardless of age)
            unseen_uids = session
                .uid_search("UNSEEN")
                .await
                .map_err(|e| e.to_string())?
                .into_iter()
                .collect();
            server_uids.extend(unseen_uids.clone());
        } else {
            // Fetch everything
            let all_uids: std::collections::HashSet<u32> = session
                .uid_search("ALL")
                .await
                .map_err(|e| e.to_string())?
                .into_iter()
                .collect();
            server_uids.extend(all_uids);

            // Still fetch unseen_uids to maintain prioritization in the sync chunk
            unseen_uids = session
                .uid_search("UNSEEN")
                .await
                .map_err(|e| e.to_string())?
                .into_iter()
                .collect();
        };

        // 3. Find missing UIDs
        let mut missing_unseen: Vec<u32> = server_uids
            .iter()
            .cloned()
            .filter(|&uid| !local_uids.contains(&uid) && unseen_uids.contains(&uid))
            .collect();
        missing_unseen.sort_unstable();

        let mut missing_since_read: Vec<u32> = server_uids
            .iter()
            .cloned()
            .filter(|&uid| !local_uids.contains(&uid) && !unseen_uids.contains(&uid))
            .collect();
        missing_since_read.sort_unstable();

        // 4. Chunk missing UIDs (LIMIT to 500 total, but prioritize UNSEEN)
        // This ensures the unread count is corrected immediately even if there's a huge backlog.
        let max_chunk_size = 500;
        let mut uids_to_fetch = Vec::new();

        // First, add all missing unseen messages (up to the limit)
        for uid in missing_unseen {
            if uids_to_fetch.len() >= max_chunk_size {
                break;
            }
            uids_to_fetch.push(uid);
        }

        // Then fill remaining space with the NEWEST missing "Since" (read) messages.
        // We take from the end of missing_since_read because higher UIDs are newer.
        while uids_to_fetch.len() < max_chunk_size && !missing_since_read.is_empty() {
            if let Some(uid) = missing_since_read.pop() {
                uids_to_fetch.push(uid);
            }
        }

        let mut new_headers = Vec::new();

        if !uids_to_fetch.is_empty() {
            // Sort for optimal IMAP fetch range building (ascending)
            uids_to_fetch.sort_unstable();
            let search_range = uids_to_fetch
                .iter()
                .map(|id| id.to_string())
                .collect::<Vec<String>>()
                .join(",");

            tracing::debug!("Syncing {} with range: {}", mailbox, search_range);
            let headers = self.fetch_headers(&search_range).await?;

            // Save & Update State for new headers
            let mut max_uid = last_known_uid;
            for header in headers {
                if header.uid > max_uid {
                    max_uid = header.uid;
                }
                if header.uid > last_known_uid {
                    new_headers.push(header.clone());
                }
                storage.save_header(mailbox, &header)?;
            }

            // Update last UID if we found new ones
            if max_uid > last_known_uid {
                storage.update_sync_state(&self.email, mailbox, current_validity, max_uid)?;
            }
        }

        // 4.5 (FLAG refresh) has been removed from this method. The caller is responsible for
        // running `refresh_flags` on a dedicated connection after releasing the shared session
        // mutex. The UIDs to refresh are returned in `SyncResult::uids_for_flag_refresh`.
        let uids_for_flag_refresh: Vec<u32> =
            local_uids.intersection(&server_uids).cloned().collect();

        // 4.6: Minimum fill — if the windowed sync leaves fewer than MIN_FILL messages
        // in the local cache, backfill with the most recent server messages we don't
        // have yet. This prevents a near-empty view when most of the inbox predates
        // cache_days (e.g. an account with 3 000+ old read messages and only 2 recent ones).
        // Fill headers are saved to storage but NOT added to new_headers so the polling
        // service doesn't treat them as newly arrived.
        const MIN_FILL: usize = 50;
        if cache_days > 0 && (local_uids.len() + new_headers.len()) < MIN_FILL {
            let needed = MIN_FILL - (local_uids.len() + new_headers.len());
            let already_have: std::collections::HashSet<u32> = local_uids
                .iter()
                .cloned()
                .chain(new_headers.iter().map(|h| h.uid))
                .collect();

            let mut all_server_uids: Vec<u32> = {
                let s = self.session.as_mut().ok_or("Not connected")?;
                s.uid_search("ALL")
                    .await
                    .map_err(|e| e.to_string())?
                    .into_iter()
                    .collect()
            };
            all_server_uids.sort_unstable_by(|a, b| b.cmp(a)); // newest first

            let mut to_fill: Vec<u32> = all_server_uids
                .into_iter()
                .filter(|uid| !already_have.contains(uid))
                .take(needed)
                .collect();

            if !to_fill.is_empty() {
                to_fill.sort_unstable(); // ascending for optimal IMAP range building
                let range = to_fill
                    .iter()
                    .map(|u| u.to_string())
                    .collect::<Vec<_>>()
                    .join(",");
                tracing::debug!(
                    "Minimum fill: fetching {} headers for {}",
                    to_fill.len(),
                    mailbox
                );
                let fill_headers = self.fetch_headers(&range).await?;
                for header in fill_headers {
                    storage.save_header(mailbox, &header)?;
                }
            }
        }

        // 5. Pruning: Remove local messages within the search range that no longer exist on server
        tracing::debug!("Pruning {}...", mailbox);
        let local_headers = storage.get_headers(mailbox).map_err(|e| e.to_string())?;

        for local in local_headers {
            // Safety: only prune if we are sure we searched the range where this message resides
            let should_be_present = if cache_days > 0 {
                // Important: It must either be within the SINCE window (local.uid >= min_since_uid)
                // OR it was previously locally unread (so we expected to find it in the UNSEEN search).
                // If it meets either criteria but it is NOT in server_uids, it has been permanently deleted
                // or marked as read outside the sync window.
                (local.uid >= min_since_uid && min_since_uid > 0) || !local.is_read
            } else {
                // ALL search, so for every local UID it MUST be in the server set or it's gone
                true
            };

            if should_be_present && !server_uids.contains(&local.uid) {
                tracing::debug!("Pruning UID {} from {}", local.uid, mailbox);
                let _ = storage.delete_header(mailbox, local.uid);
            }
        }

        // Store server total so the UI can show accurate counts
        let _ = storage.update_server_total(&self.email, mailbox, mailbox_exists);

        Ok(SyncResult {
            uids_for_flag_refresh,
            new_headers,
        })
    }

    /// Fetches FLAGS for a set of already-cached UIDs and updates local read-status in storage.
    ///
    /// This is the extracted step 4.5 from `sync_mailbox`. It must be called on a **dedicated**
    /// `ImapClient` (not the shared session) so it does not block other callers while issuing
    /// potentially many sequential `UID FETCH … FLAGS` round-trips.
    pub async fn refresh_flags(
        &mut self,
        mailbox: &str,
        uids: &[u32],
        storage: &crate::storage::Storage,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        if uids.is_empty() {
            return Ok(());
        }

        self.select_mailbox(mailbox).await?;

        for chunk in uids.chunks(200) {
            let range = chunk
                .iter()
                .map(|id| id.to_string())
                .collect::<Vec<String>>()
                .join(",");

            let session = self.session.as_mut().ok_or("Not connected")?;
            let mut fetches = session
                .uid_fetch(&range, "FLAGS")
                .await
                .map_err(|e| e.to_string())?;

            loop {
                match fetches.try_next().await {
                    Ok(Some(fetch)) => {
                        if let Some(uid) = fetch.uid {
                            let is_read = fetch
                                .flags()
                                .any(|f| matches!(f, async_imap::types::Flag::Seen));
                            let _ = storage.update_read_status(mailbox, uid, is_read);
                        }
                    }
                    Ok(None) => break,
                    Err(e) => {
                        tracing::warn!("FLAG refresh fetch error: {}", e);
                        break;
                    }
                }
            }
        }

        Ok(())
    }

    pub async fn create_folder(
        &mut self,
        folder_name: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        eprintln!("[ImapClient::create_folder] folder_name={folder_name}");
        let session = self.session.as_mut().ok_or("Not connected")?;
        session.create(folder_name).await.map_err(|e| {
            let msg = e.to_string();
            eprintln!("[ImapClient::create_folder] ERROR: {msg}");
            msg
        })?;
        eprintln!("[ImapClient::create_folder] session.create returned");
        // Best-effort subscribe so folder appears on servers that use LSUB
        let _ = session.subscribe(folder_name).await;
        Ok(())
    }

    pub async fn delete_folder(
        &mut self,
        folder_name: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;
        let _ = session.unsubscribe(folder_name).await; // best-effort
        session
            .delete(folder_name)
            .await
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    pub async fn move_message(
        &mut self,
        uid: u32,
        from_mailbox: &str,
        to_mailbox: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(from_mailbox).await?;

        let session = self.session.as_mut().ok_or("Not connected")?;

        // Move the message
        session
            .uid_mv(uid.to_string(), to_mailbox)
            .await
            .map_err(|e| e.to_string())?;

        // Expunge to finalize deletion from source
        let _ = session
            .expunge()
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }

    pub async fn move_messages(
        &mut self,
        uids: &[u32],
        from_mailbox: &str,
        to_mailbox: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        if uids.is_empty() {
            return Ok(());
        }
        self.select_mailbox(from_mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;
        let uid_set: String = uids
            .iter()
            .map(|u| u.to_string())
            .collect::<Vec<_>>()
            .join(",");
        session
            .uid_mv(&uid_set, to_mailbox)
            .await
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    pub async fn mark_as_read(
        &mut self,
        mailbox: &str,
        uid: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        // Add \Seen flag
        session
            .uid_store(uid.to_string(), "+FLAGS (\\Seen)")
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }

    pub async fn mark_as_unread(
        &mut self,
        mailbox: &str,
        uid: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        // Remove \Seen flag
        session
            .uid_store(uid.to_string(), "-FLAGS (\\Seen)")
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }

    pub async fn mark_all_as_read(
        &mut self,
        mailbox: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        session
            .uid_store("1:*", "+FLAGS (\\Seen)")
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }

    pub async fn set_message_flag(
        &mut self,
        mailbox: &str,
        uid: u32,
        flag: Option<String>,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        let to_remove = "-FLAGS (\\Flagged korax-flag-red korax-flag-orange korax-flag-yellow korax-flag-green korax-flag-blue)";
        let _ = session
            .uid_store(uid.to_string(), to_remove)
            .await?
            .try_collect::<Vec<_>>()
            .await?;

        if let Some(f) = flag {
            let to_add = match f.as_str() {
                "red" => "+FLAGS (\\Flagged korax-flag-red)",
                "orange" => "+FLAGS (\\Flagged korax-flag-orange)",
                "yellow" => "+FLAGS (\\Flagged korax-flag-yellow)",
                "green" => "+FLAGS (\\Flagged korax-flag-green)",
                "blue" => "+FLAGS (\\Flagged korax-flag-blue)",
                _ => return Err("Invalid flag".into()),
            };
            let _ = session
                .uid_store(uid.to_string(), to_add)
                .await?
                .try_collect::<Vec<_>>()
                .await?;
        }
        Ok(())
    }

    /// Waits for an update from the server using IDLE.
    /// Returns true if new data is available, false on timeout.
    pub async fn wait_for_update(
        &mut self,
        timeout_secs: u64,
    ) -> Result<bool, Box<dyn Error + Send + Sync>> {
        let session = self.session.take().ok_or("Not connected")?;

        // Transition to IDLE state
        let mut handle = session.idle();
        handle.init().await.map_err(|e| e.to_string())?;

        // Wait for update or timeout
        // IMAP IDLE usually recommends a 29-minute timeout, but we allow shorter for polling synergy
        let (wait_fut, _stop_source) = handle.wait_with_timeout(Duration::from_secs(timeout_secs));
        let result = wait_fut.await;

        // Transition back to normal session
        self.session = Some(handle.done().await.map_err(|e| e.to_string())?);

        match result {
            Ok(IdleResponse::NewData(_)) => Ok(true),
            Ok(IdleResponse::Timeout) | Ok(IdleResponse::ManualInterrupt) => {
                tracing::debug!("IDLE timed out or interrupted for {}", self.email);
                Ok(false)
            }
            Err(e) => Err(format!("IDLE error: {}", e).into()),
        }
    }
    pub async fn append_message(
        &mut self,
        mailbox: &str,
        content: &[u8],
        _flags: Option<String>,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let session = self.session.as_mut().ok_or("Not connected")?;
        session
            .append(mailbox, content)
            .await
            .map_err(|e| e.to_string())?;
        Ok(())
    }

    pub async fn delete_all_messages(
        &mut self,
        mailbox: &str,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox).await?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        // Mark every message as deleted
        session
            .uid_store("1:*", "+FLAGS (\\Deleted)")
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        // Expunge to permanently remove
        session
            .expunge()
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }

    pub async fn delete_message(
        &mut self,
        mailbox: &str,
        uid: u32,
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        self.select_mailbox(mailbox)
            .await
            .map_err(|e| e.to_string())?;
        let session = self.session.as_mut().ok_or("Not connected")?;

        // 1. Mark as deleted
        let _ = session
            .uid_store(uid.to_string(), "+FLAGS (\\Deleted)")
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        // 2. Expunge to permanently remove
        let _ = session
            .expunge()
            .await
            .map_err(|e| e.to_string())?
            .try_collect::<Vec<_>>()
            .await
            .map_err(|e| e.to_string())?;

        Ok(())
    }
}

pub fn preprocess_header_bytes(bytes: &[u8]) -> Vec<u8> {
    // Convert to string; if not UTF-8 leave untouched.
    let s = match std::str::from_utf8(bytes) {
        Ok(s) => s,
        Err(_) => return bytes.to_vec(),
    };

    // Join adjacent RFC 2047 encoded-words that share the same charset and
    // transfer-encoding. Many mailers split multi-byte UTF-8 sequences across
    // word boundaries, which confuses mail-parser.
    //
    // A word has the form:  =?CHARSET?ENCODING?PAYLOAD?=
    // Two adjacent words (possibly separated by whitespace) can be merged into:
    //   =?CHARSET?ENCODING?PAYLOAD1PAYLOAD2?=
    let result = join_adjacent_encoded_words(s);
    result.into_bytes()
}

/// Parse one RFC 2047 encoded word starting at `s`.
/// Returns `(charset, encoding, payload, remaining)` on success.
fn parse_encoded_word(s: &str) -> Option<(&str, &str, &str, &str)> {
    let s = s.strip_prefix("=?")?;
    let (charset, rest) = s.split_once('?')?;
    let (encoding, rest) = rest.split_once('?')?;
    // encoding must be a single B or Q (case-insensitive)
    if encoding.len() != 1 {
        return None;
    }
    let enc_char = encoding.chars().next()?.to_ascii_uppercase();
    if enc_char != 'B' && enc_char != 'Q' {
        return None;
    }
    let end = rest.find("?=")?;
    let payload = &rest[..end];
    let remaining = &rest[end + 2..];
    Some((charset, encoding, payload, remaining))
}

fn join_adjacent_encoded_words(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut pos = 0;

    while pos < s.len() {
        let remaining = &s[pos..];

        if !remaining.starts_with("=?") {
            // Not an encoded word – emit one char and advance.
            let ch = remaining.chars().next().unwrap();
            out.push(ch);
            pos += ch.len_utf8();
            continue;
        }

        // Try to parse an encoded word.
        match parse_encoded_word(remaining) {
            None => {
                // Looks like "=?" but is not a valid encoded word.
                out.push('=');
                pos += 1;
            }
            Some((charset, encoding, payload, after_first)) => {
                // Accumulate adjacent words with matching charset+encoding.
                let mut combined_payload = payload.to_string();
                let mut cursor = after_first;

                loop {
                    // Skip optional linear whitespace between adjacent words, including header folding (CRLF).
                    let ws_stripped = cursor.trim_start_matches(|c: char| c.is_whitespace());

                    if !ws_stripped.starts_with("=?") {
                        break;
                    }

                    match parse_encoded_word(ws_stripped) {
                        None => break,
                        Some((cs2, enc2, payload2, after_next)) => {
                            // Only merge if charset AND encoding match (case-insensitive).
                            if cs2.eq_ignore_ascii_case(charset)
                                && enc2.eq_ignore_ascii_case(encoding)
                            {
                                combined_payload.push_str(payload2);
                                cursor = after_next;
                            } else {
                                break;
                            }
                        }
                    }
                }

                // Emit merged word.
                out.push_str("=?");
                out.push_str(charset);
                out.push('?');
                out.push_str(encoding);
                out.push('?');
                out.push_str(&combined_payload);
                out.push_str("?=");

                // Advance pos past everything we consumed.
                pos += remaining.len() - cursor.len();
            }
        }
    }

    out
}
