elektromail 0.1.1

use std::{
    io,
    sync::{Arc, Mutex},
};

use base64::{Engine as _, engine::general_purpose};
use tokio::{
    io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
    net::{TcpListener, TcpStream},
    sync::broadcast,
};

use crate::auth::SharedUserStore;
use crate::store::{FlagOp, FlagSet, Store, current_internal_date, uid_to_seq};
use crate::tls::{ImapStream, build_tls_acceptor};
use crate::{MailboxEvent, MailboxNotifier};

type ImapReader = BufReader<ImapStream>;

pub(crate) mod fetch;
pub(crate) mod mime;
mod search;
mod util;

const UIDVALIDITY: u32 = 1;

use fetch::{
    FetchRequest, format_flags, parse_fetch_request, send_fetch_seq, send_fetch_uid,
    write_fetch_response,
};
use search::{parse_search_args, send_search};
use util::{
    list_match, parse_append_mailbox, parse_copy_target, parse_flag_list_from_line,
    parse_imap_args, parse_internal_date_from_line, parse_literal_size, parse_mailbox_name,
    parse_move_target, parse_seq_range, parse_sequence_set, parse_uid_range,
};

async fn write_raw(writer: &mut ImapReader, data: &[u8]) -> io::Result<()> {
    writer.get_mut().write_all(data).await
}

async fn write_tagged(writer: &mut ImapReader, tag: &str, msg: &str) -> io::Result<()> {
    write_raw(writer, format!("{} {}\r\n", tag, msg).as_bytes()).await
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum SessionState {
    NotAuthenticated,
    Authenticated,
    Selected,
    SelectedReadOnly,
}

pub(crate) async fn run_imap(
    listener: TcpListener,
    store: Arc<Mutex<Store>>,
    auth: SharedUserStore,
    mut shutdown_rx: broadcast::Receiver<()>,
    mailbox_notifier: MailboxNotifier,
) {
    loop {
        tokio::select! {
            accept = listener.accept() => {
                let Ok((stream, _)) = accept else { break };
                let store = store.clone();
                let auth = auth.clone();
                let notifier = mailbox_notifier.clone();
                tokio::spawn(async move {
                    let _ = handle_imap(stream, store, auth, notifier).await;
                });
            }
            _ = shutdown_rx.recv() => {
                break;
            }
        }
    }
}

async fn handle_imap(
    stream: TcpStream,
    store: Arc<Mutex<Store>>,
    auth: SharedUserStore,
    mailbox_notifier: MailboxNotifier,
) -> io::Result<()> {
    let mut reader: ImapReader = BufReader::new(ImapStream::Plain(stream));
    let mut tls_active = false;
    let capability_line = |tls_enabled: bool| {
        if tls_enabled {
            "IMAP4rev1 IMAP4rev2 UIDPLUS LITERAL+ IDLE MOVE AUTH=PLAIN QUOTA"
        } else {
            "IMAP4rev1 IMAP4rev2 UIDPLUS LITERAL+ IDLE MOVE AUTH=PLAIN QUOTA STARTTLS"
        }
    };

    write_raw(
        &mut reader,
        format!(
            "* OK [CAPABILITY {}] elektromail ready\r\n",
            capability_line(false)
        )
        .as_bytes(),
    )
    .await?;

    let mut line = String::new();
    let mut current_user = "user".to_string();
    let mut current_mailbox = "INBOX".to_string();
    let mut state = SessionState::NotAuthenticated;

    loop {
        line.clear();
        let bytes = reader.read_line(&mut line).await?;
        if bytes == 0 {
            break;
        }

        let trimmed = line.trim_end_matches(&['\r', '\n'][..]).to_string();
        let mut parts = trimmed.split_whitespace();
        let Some(tag) = parts.next() else { continue };
        let Some(cmd) = parts.next() else {
            write_tagged(&mut reader, tag, "BAD Missing command").await?;
            continue;
        };

        let upper_cmd = cmd.to_ascii_uppercase();
        match upper_cmd.as_str() {
            "CAPABILITY" => {
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD CAPABILITY takes no arguments").await?;
                } else {
                    let line = format!("* CAPABILITY {}\r\n", capability_line(tls_active));
                    write_raw(&mut reader, line.as_bytes()).await?;
                    write_tagged(&mut reader, tag, "OK CAPABILITY completed").await?;
                }
            }
            "STARTTLS" => {
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD STARTTLS takes no arguments").await?;
                    continue;
                }
                if state != SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD STARTTLS not permitted").await?;
                    continue;
                }
                if tls_active {
                    write_tagged(&mut reader, tag, "BAD TLS already active").await?;
                    continue;
                }
                write_tagged(&mut reader, tag, "OK Begin TLS negotiation").await?;
                let acceptor = build_tls_acceptor()?;
                let inner = reader.into_inner();
                let ImapStream::Plain(stream) = inner else {
                    return Err(io::Error::other("STARTTLS requires plaintext stream"));
                };
                let tls_stream = acceptor.accept(stream).await.map_err(io::Error::other)?;
                reader = BufReader::new(ImapStream::Tls(tls_stream));
                tls_active = true;
                state = SessionState::NotAuthenticated;
                current_user = "user".to_string();
                current_mailbox = "INBOX".to_string();
            }
            "AUTHENTICATE" => {
                if state != SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Already authenticated").await?;
                    continue;
                }
                let args = parse_imap_args(&trimmed);
                if args.is_empty() {
                    write_tagged(&mut reader, tag, "BAD Missing mechanism").await?;
                    continue;
                }
                let mechanism = args[0].to_ascii_uppercase();
                if mechanism != "PLAIN" {
                    write_tagged(&mut reader, tag, "NO Unsupported authentication mechanism")
                        .await?;
                    continue;
                }

                let mut response = args.get(1).cloned();
                if response.is_none() {
                    write_raw(&mut reader, b"+ \r\n").await?;
                    let mut resp_line = String::new();
                    let bytes = reader.read_line(&mut resp_line).await?;
                    if bytes == 0 {
                        break;
                    }
                    let resp_trim = resp_line.trim_end_matches(&['\r', '\n'][..]);
                    if resp_trim == "*" {
                        write_tagged(&mut reader, tag, "BAD AUTHENTICATE cancelled").await?;
                        continue;
                    }
                    response = Some(resp_trim.to_string());
                }

                let response = response.unwrap_or_default();
                if response.is_empty() || response == "=" {
                    write_tagged(&mut reader, tag, "NO Invalid credentials").await?;
                    continue;
                }
                let Ok(decoded) = general_purpose::STANDARD.decode(response.as_bytes()) else {
                    write_tagged(&mut reader, tag, "BAD Invalid authentication response").await?;
                    continue;
                };
                let decoded = String::from_utf8_lossy(&decoded);
                let mut parts = decoded.split('\0');
                let _authzid = parts.next().unwrap_or("");
                let authcid = parts.next().unwrap_or("");
                let passwd = parts.next().unwrap_or("");

                if auth.authenticate(authcid, passwd) {
                    current_user = authcid.to_string();
                    state = SessionState::Authenticated;
                    write_tagged(&mut reader, tag, "OK AUTHENTICATE completed").await?;
                } else {
                    write_tagged(&mut reader, tag, "NO Invalid credentials").await?;
                }
            }
            "LOGIN" => {
                if state != SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Already authenticated").await?;
                    continue;
                }
                let args = parse_imap_args(&trimmed);
                if args.len() < 2 {
                    write_tagged(&mut reader, tag, "BAD Missing credentials").await?;
                    continue;
                }
                let user = args[0].clone();
                let pass = args[1].clone();
                if auth.authenticate(&user, &pass) {
                    current_user = user;
                    state = SessionState::Authenticated;
                    write_tagged(&mut reader, tag, "OK LOGIN completed").await?;
                } else {
                    write_tagged(&mut reader, tag, "NO Invalid credentials").await?;
                }
            }
            "LIST" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let _reference = parts.next().unwrap_or("\"\"");
                let pattern = parts.next().unwrap_or("\"\"");
                let pattern = parse_mailbox_name(pattern);
                if pattern.is_empty() {
                    write_raw(&mut reader, b"* LIST (\\Noselect) \"/\" \"\"\r\n").await?;
                    write_tagged(&mut reader, tag, "OK LIST completed").await?;
                    continue;
                }
                let mailboxes = store
                    .lock()
                    .expect("store lock poisoned")
                    .list_mailboxes(&current_user);
                let subs = store
                    .lock()
                    .expect("store lock poisoned")
                    .list_subscriptions(&current_user);
                for name in mailboxes {
                    if !list_match(&name, &pattern) {
                        continue;
                    }
                    let subscribed = subs.iter().any(|s| s.eq_ignore_ascii_case(&name));
                    let attrs = if subscribed {
                        "(\\Subscribed \\HasNoChildren)"
                    } else {
                        "(\\HasNoChildren)"
                    };
                    write_raw(
                        &mut reader,
                        format!("* LIST {} \"/\" {}\r\n", attrs, name).as_bytes(),
                    )
                    .await?;
                }
                write_tagged(&mut reader, tag, "OK LIST completed").await?;
            }
            "LSUB" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let _reference = parts.next().unwrap_or("\"\"");
                let pattern = parts.next().unwrap_or("\"\"");
                let pattern = parse_mailbox_name(pattern);
                let mailboxes = store
                    .lock()
                    .expect("store lock poisoned")
                    .list_subscriptions(&current_user);
                for name in mailboxes {
                    if !list_match(&name, &pattern) {
                        continue;
                    }
                    write_raw(
                        &mut reader,
                        format!("* LSUB (\\HasNoChildren) \"/\" {}\r\n", name).as_bytes(),
                    )
                    .await?;
                }
                write_tagged(&mut reader, tag, "OK LSUB completed").await?;
            }
            "SELECT" | "EXAMINE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if let Some(name) = parts.next() {
                    current_mailbox = parse_mailbox_name(name);
                } else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                    continue;
                }
                let exists_mailbox = store
                    .lock()
                    .expect("store lock poisoned")
                    .list_mailboxes(&current_user)
                    .iter()
                    .any(|m| m.eq_ignore_ascii_case(&current_mailbox));
                if !exists_mailbox {
                    write_tagged(&mut reader, tag, "NO Mailbox does not exist").await?;
                    continue;
                }
                let messages = store
                    .lock()
                    .expect("store lock poisoned")
                    .list(&current_user, &current_mailbox);
                let exists = messages.len();
                let flags = b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n";
                write_raw(&mut reader, flags).await?;
                write_raw(&mut reader, format!("* {} EXISTS\r\n", exists).as_bytes()).await?;
                write_raw(
                    &mut reader,
                    b"* LIST (\\Subscribed \\HasNoChildren) \"/\" INBOX\r\n",
                )
                .await?;
                write_raw(
                    &mut reader,
                    b"* OK PERMANENTFLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n",
                )
                .await?;
                write_raw(
                    &mut reader,
                    format!("* OK UIDVALIDITY {}\r\n", UIDVALIDITY).as_bytes(),
                )
                .await?;
                write_raw(
                    &mut reader,
                    format!("* OK UIDNEXT {}\r\n", exists + 1).as_bytes(),
                )
                .await?;
                if upper_cmd == "EXAMINE" {
                    state = SessionState::SelectedReadOnly;
                    write_tagged(&mut reader, tag, "OK [READ-ONLY] EXAMINE completed").await?;
                } else {
                    state = SessionState::Selected;
                    write_tagged(&mut reader, tag, "OK [READ-WRITE] SELECT completed").await?;
                }
            }
            "UID" => {
                if let Some(subcmd) = parts.next() {
                    if subcmd.eq_ignore_ascii_case("FETCH") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected
                            && state != SessionState::SelectedReadOnly
                        {
                            write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                            continue;
                        }
                        let range = parts.next().unwrap_or("1:*");
                        let req = parse_fetch_request(&trimmed);
                        send_fetch_uid(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            range,
                            &req,
                        )
                        .await?;
                        write_tagged(&mut reader, tag, "OK UID FETCH completed").await?;
                    } else if subcmd.eq_ignore_ascii_case("SEARCH") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected
                            && state != SessionState::SelectedReadOnly
                        {
                            write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                            continue;
                        }
                        let args = parse_search_args(&trimmed);
                        if args.is_empty() {
                            write_tagged(&mut reader, tag, "BAD Missing search criteria").await?;
                            continue;
                        }
                        send_search(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            &args,
                            true,
                        )
                        .await?;
                        write_tagged(&mut reader, tag, "OK UID SEARCH completed").await?;
                    } else if subcmd.eq_ignore_ascii_case("STORE") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected {
                            write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                            continue;
                        }
                        let result = handle_store(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            &trimmed,
                            true,
                        )
                        .await?;
                        match result {
                            StoreResult::Ok => {
                                write_tagged(&mut reader, tag, "OK UID STORE completed").await?;
                            }
                            StoreResult::Bad => {
                                write_tagged(&mut reader, tag, "BAD Invalid STORE arguments")
                                    .await?;
                            }
                        }
                    } else if subcmd.eq_ignore_ascii_case("COPY") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected
                            && state != SessionState::SelectedReadOnly
                        {
                            write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                            continue;
                        }
                        let result = handle_copy(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            &trimmed,
                        );
                        match result {
                            CopyResult::Ok { copyuid } => {
                                if let Some(copyuid) = copyuid {
                                    let response = format!(
                                        "OK [COPYUID {} {} {}] UID COPY completed",
                                        UIDVALIDITY, copyuid.source, copyuid.dest
                                    );
                                    write_tagged(&mut reader, tag, &response).await?;
                                } else {
                                    write_tagged(&mut reader, tag, "OK UID COPY completed").await?;
                                }
                            }
                            CopyResult::No => {
                                write_tagged(&mut reader, tag, "NO UID COPY failed").await?;
                            }
                            CopyResult::Bad => {
                                write_tagged(&mut reader, tag, "BAD Invalid UID COPY").await?;
                            }
                        }
                    } else if subcmd.eq_ignore_ascii_case("MOVE") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected {
                            write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                            continue;
                        }
                        let result = handle_move(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            &trimmed,
                        )
                        .await?;
                        match result {
                            MoveResult::Ok => {
                                write_tagged(&mut reader, tag, "OK UID MOVE completed").await?;
                            }
                            MoveResult::No => {
                                write_tagged(&mut reader, tag, "NO UID MOVE failed").await?;
                            }
                            MoveResult::Bad => {
                                write_tagged(&mut reader, tag, "BAD Invalid UID MOVE").await?;
                            }
                        }
                    } else if subcmd.eq_ignore_ascii_case("EXPUNGE") {
                        if state == SessionState::NotAuthenticated {
                            write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                            continue;
                        }
                        if state != SessionState::Selected {
                            write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                            continue;
                        }
                        let args = parse_imap_args(&trimmed);
                        let Some(seqset) = args.get(1).map(String::as_str) else {
                            write_tagged(&mut reader, tag, "BAD Missing UID set").await?;
                            continue;
                        };
                        if seqset.is_empty() {
                            write_tagged(&mut reader, tag, "BAD Missing UID set").await?;
                            continue;
                        }
                        handle_uid_expunge(
                            &mut reader,
                            &store,
                            &current_user,
                            &current_mailbox,
                            seqset,
                        )
                        .await?;
                        write_tagged(&mut reader, tag, "OK UID EXPUNGE completed").await?;
                    } else {
                        write_tagged(&mut reader, tag, "BAD Unsupported UID command").await?;
                    }
                } else {
                    write_tagged(&mut reader, tag, "BAD Missing UID subcommand").await?;
                }
            }
            "FETCH" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected && state != SessionState::SelectedReadOnly {
                    write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                    continue;
                }
                let range = parts.next().unwrap_or("1:*");
                let req = parse_fetch_request(&trimmed);
                send_fetch_seq(
                    &mut reader,
                    &store,
                    &current_user,
                    &current_mailbox,
                    range,
                    &req,
                )
                .await?;
                write_tagged(&mut reader, tag, "OK FETCH completed").await?;
            }
            "SEARCH" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected && state != SessionState::SelectedReadOnly {
                    write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                    continue;
                }
                let args = parse_search_args(&trimmed);
                if args.is_empty() {
                    write_tagged(&mut reader, tag, "BAD Missing search criteria").await?;
                } else {
                    send_search(
                        &mut reader,
                        &store,
                        &current_user,
                        &current_mailbox,
                        &args,
                        false,
                    )
                    .await?;
                    write_tagged(&mut reader, tag, "OK SEARCH completed").await?;
                }
            }
            "IDLE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD IDLE takes no arguments").await?;
                    continue;
                }
                write_raw(&mut reader, b"+ idling\r\n").await?;

                // Subscribe to mailbox notifications
                let mut notifier_rx = mailbox_notifier.subscribe();
                let has_mailbox_selected =
                    state == SessionState::Selected || state == SessionState::SelectedReadOnly;

                loop {
                    line.clear();
                    tokio::select! {
                        // Check for client sending DONE
                        result = reader.read_line(&mut line) => {
                            let bytes = result?;
                            if bytes == 0 {
                                break;
                            }
                            let trimmed = line.trim_end_matches(&['\r', '\n'][..]);
                            if trimmed.eq_ignore_ascii_case("DONE") {
                                break;
                            }
                        }
                        // Check for mailbox notifications (only if a mailbox is selected)
                        result = notifier_rx.recv(), if has_mailbox_selected => {
                            if let Ok(event) = result {
                                // Only notify if this is the user's selected mailbox
                                if event.user.eq_ignore_ascii_case(&current_user)
                                    && event.mailbox.eq_ignore_ascii_case(&current_mailbox)
                                {
                                    write_raw(
                                        &mut reader,
                                        format!("* {} EXISTS\r\n", event.new_count).as_bytes(),
                                    ).await?;
                                }
                            }
                        }
                    }
                }
                write_tagged(&mut reader, tag, "OK IDLE completed").await?;
            }
            "STORE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected {
                    write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                    continue;
                }
                let result = handle_store(
                    &mut reader,
                    &store,
                    &current_user,
                    &current_mailbox,
                    &trimmed,
                    false,
                )
                .await?;
                match result {
                    StoreResult::Ok => {
                        write_tagged(&mut reader, tag, "OK STORE completed").await?;
                    }
                    StoreResult::Bad => {
                        write_tagged(&mut reader, tag, "BAD Invalid STORE arguments").await?;
                    }
                }
            }
            "COPY" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected && state != SessionState::SelectedReadOnly {
                    write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                    continue;
                }
                let result = handle_copy(
                    &mut reader,
                    &store,
                    &current_user,
                    &current_mailbox,
                    &trimmed,
                );
                match result {
                    CopyResult::Ok { copyuid } => {
                        if let Some(copyuid) = copyuid {
                            let response = format!(
                                "OK [COPYUID {} {} {}] COPY completed",
                                UIDVALIDITY, copyuid.source, copyuid.dest
                            );
                            write_tagged(&mut reader, tag, &response).await?;
                        } else {
                            write_tagged(&mut reader, tag, "OK COPY completed").await?;
                        }
                    }
                    CopyResult::No => {
                        write_tagged(&mut reader, tag, "NO COPY failed").await?;
                    }
                    CopyResult::Bad => {
                        write_tagged(&mut reader, tag, "BAD Invalid COPY arguments").await?;
                    }
                }
            }
            "EXPUNGE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected {
                    write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                    continue;
                }
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD EXPUNGE takes no arguments").await?;
                    continue;
                }
                handle_expunge(&mut reader, &store, &current_user, &current_mailbox).await?;
                write_tagged(&mut reader, tag, "OK EXPUNGE completed").await?;
            }
            "MOVE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected {
                    write_tagged(&mut reader, tag, "NO Mailbox is read-only").await?;
                    continue;
                }
                let result = handle_move(
                    &mut reader,
                    &store,
                    &current_user,
                    &current_mailbox,
                    &trimmed,
                )
                .await?;
                match result {
                    MoveResult::Ok => {
                        write_tagged(&mut reader, tag, "OK MOVE completed").await?;
                    }
                    MoveResult::No => {
                        write_tagged(&mut reader, tag, "NO MOVE failed").await?;
                    }
                    MoveResult::Bad => {
                        write_tagged(&mut reader, tag, "BAD Invalid MOVE arguments").await?;
                    }
                }
            }
            "GETQUOTAROOT" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let mailbox = parts.next().unwrap_or("INBOX");
                let mailbox = parse_mailbox_name(mailbox);
                write_raw(
                    &mut reader,
                    format!("* QUOTAROOT {} \"\"\r\n", mailbox).as_bytes(),
                )
                .await?;
                write_raw(&mut reader, b"* QUOTA \"\" (STORAGE 0 10240)\r\n").await?;
                write_tagged(&mut reader, tag, "OK GETQUOTAROOT completed").await?;
            }
            "GETQUOTA" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let root = parts.next().unwrap_or("\"\"");
                write_raw(
                    &mut reader,
                    format!("* QUOTA {} (STORAGE 0 10240)\r\n", root).as_bytes(),
                )
                .await?;
                write_tagged(&mut reader, tag, "OK GETQUOTA completed").await?;
            }
            "APPEND" => {
                let result = handle_append(
                    &mut reader,
                    &store,
                    &current_user,
                    &trimmed,
                    state != SessionState::NotAuthenticated,
                    &mailbox_notifier,
                )
                .await?;
                match result {
                    AppendResult::Ok { uid } => {
                        let response =
                            format!("OK [APPENDUID {} {}] APPEND completed", UIDVALIDITY, uid);
                        write_tagged(&mut reader, tag, &response).await?;
                    }
                    AppendResult::No => {
                        write_tagged(&mut reader, tag, "NO APPEND failed").await?;
                    }
                    AppendResult::Bad => {
                        write_tagged(&mut reader, tag, "BAD Invalid APPEND").await?;
                    }
                }
            }
            "CREATE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if let Some(name) = parts.next() {
                    let mailbox = parse_mailbox_name(name);
                    let created = store
                        .lock()
                        .expect("store lock poisoned")
                        .create_mailbox(&current_user, &mailbox);
                    if created {
                        write_tagged(&mut reader, tag, "OK CREATE completed").await?;
                    } else {
                        write_tagged(&mut reader, tag, "NO CREATE failed").await?;
                    }
                } else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                }
            }
            "SUBSCRIBE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let Some(name) = parts.next() else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                    continue;
                };
                let ok = store
                    .lock()
                    .expect("store lock poisoned")
                    .subscribe(&current_user, &parse_mailbox_name(name));
                if ok {
                    write_tagged(&mut reader, tag, "OK SUBSCRIBE completed").await?;
                } else {
                    write_tagged(&mut reader, tag, "NO SUBSCRIBE failed").await?;
                }
            }
            "UNSUBSCRIBE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let Some(name) = parts.next() else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                    continue;
                };
                let ok = store
                    .lock()
                    .expect("store lock poisoned")
                    .unsubscribe(&current_user, &parse_mailbox_name(name));
                if ok {
                    write_tagged(&mut reader, tag, "OK UNSUBSCRIBE completed").await?;
                } else {
                    write_tagged(&mut reader, tag, "NO UNSUBSCRIBE failed").await?;
                }
            }
            "NAMESPACE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD NAMESPACE takes no arguments").await?;
                    continue;
                }
                write_raw(&mut reader, b"* NAMESPACE ((\"\" \"/\")) NIL NIL\r\n").await?;
                write_tagged(&mut reader, tag, "OK NAMESPACE completed").await?;
            }
            "STATUS" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let Some(name) = parts.next() else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                    continue;
                };
                let items = parts.collect::<Vec<_>>().join(" ");
                if items.is_empty() {
                    write_tagged(&mut reader, tag, "BAD Missing status items").await?;
                    continue;
                }
                let mailbox = parse_mailbox_name(name);
                let (exists, messages) = {
                    let mut guard = store.lock().expect("store lock poisoned");
                    let exists = guard
                        .list_mailboxes(&current_user)
                        .iter()
                        .any(|m| m.eq_ignore_ascii_case(&mailbox));
                    let messages = guard.list(&current_user, &mailbox);
                    drop(guard);
                    (exists, messages)
                };
                if !exists {
                    write_tagged(&mut reader, tag, "NO Mailbox does not exist").await?;
                    continue;
                }
                let mut status_items = Vec::new();
                let items_upper = items.to_ascii_uppercase();
                if items_upper.contains("MESSAGES") {
                    status_items.push(format!("MESSAGES {}", messages.len()));
                }
                if items_upper.contains("UIDNEXT") {
                    status_items.push(format!("UIDNEXT {}", messages.len() + 1));
                }
                if items_upper.contains("UIDVALIDITY") {
                    status_items.push(format!("UIDVALIDITY {}", UIDVALIDITY));
                }
                if items_upper.contains("UNSEEN") {
                    let unseen = messages.iter().filter(|m| !m.seen).count();
                    status_items.push(format!("UNSEEN {}", unseen));
                }
                write_raw(
                    &mut reader,
                    format!("* STATUS {} ({})\r\n", mailbox, status_items.join(" ")).as_bytes(),
                )
                .await?;
                write_tagged(&mut reader, tag, "OK STATUS completed").await?;
            }
            "CLOSE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected {
                    write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                    continue;
                }
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD CLOSE takes no arguments").await?;
                    continue;
                }
                let _ = store
                    .lock()
                    .expect("store lock poisoned")
                    .expunge_deleted(&current_user, &current_mailbox);
                state = SessionState::Authenticated;
                write_tagged(&mut reader, tag, "OK CLOSE completed").await?;
            }
            "UNSELECT" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if state != SessionState::Selected && state != SessionState::SelectedReadOnly {
                    write_tagged(&mut reader, tag, "BAD No mailbox selected").await?;
                    continue;
                }
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD UNSELECT takes no arguments").await?;
                    continue;
                }
                state = SessionState::Authenticated;
                write_tagged(&mut reader, tag, "OK UNSELECT completed").await?;
            }
            "RENAME" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                let Some(old) = parts.next() else {
                    write_tagged(&mut reader, tag, "BAD Missing source").await?;
                    continue;
                };
                let Some(new) = parts.next() else {
                    write_tagged(&mut reader, tag, "BAD Missing destination").await?;
                    continue;
                };
                let ok = store.lock().expect("store lock poisoned").rename_mailbox(
                    &current_user,
                    &parse_mailbox_name(old),
                    &parse_mailbox_name(new),
                );
                if ok {
                    write_tagged(&mut reader, tag, "OK RENAME completed").await?;
                } else {
                    write_tagged(&mut reader, tag, "NO RENAME failed").await?;
                }
            }
            "DELETE" => {
                if state == SessionState::NotAuthenticated {
                    write_tagged(&mut reader, tag, "BAD Not authenticated").await?;
                    continue;
                }
                if let Some(name) = parts.next() {
                    let ok = store
                        .lock()
                        .expect("store lock poisoned")
                        .delete_mailbox(&current_user, &parse_mailbox_name(name));
                    if ok {
                        write_tagged(&mut reader, tag, "OK DELETE completed").await?;
                    } else {
                        write_tagged(&mut reader, tag, "NO DELETE failed").await?;
                    }
                } else {
                    write_tagged(&mut reader, tag, "BAD Missing mailbox").await?;
                }
            }
            "NOOP" => {
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD NOOP takes no arguments").await?;
                } else {
                    write_tagged(&mut reader, tag, "OK NOOP").await?;
                }
            }
            "LOGOUT" => {
                if parts.next().is_some() {
                    write_tagged(&mut reader, tag, "BAD LOGOUT takes no arguments").await?;
                } else {
                    write_raw(&mut reader, b"* BYE Logging out\r\n").await?;
                    write_tagged(&mut reader, tag, "OK LOGOUT completed").await?;
                    break;
                }
            }
            _ => {
                write_tagged(&mut reader, tag, "BAD Unsupported command").await?;
            }
        }
    }

    Ok(())
}

enum StoreResult {
    Ok,
    Bad,
}

async fn handle_store(
    writer: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    mailbox: &str,
    line: &str,
    is_uid: bool,
) -> io::Result<StoreResult> {
    let store_args = parse_store_args(line);
    if store_args.target.is_empty() || store_args.mode.is_none() {
        return Ok(StoreResult::Bad);
    }
    let silent = store_args.silent;
    let flagset = store_args.flags;

    let mut updated_seq = Vec::new();
    if is_uid {
        let max_uid = store
            .lock()
            .expect("store lock poisoned")
            .list(user, mailbox)
            .len();
        let max_uid = u32::try_from(max_uid).unwrap_or(u32::MAX);
        let targets = parse_sequence_set(&store_args.target, max_uid);
        for uid in targets {
            let mut guard = store.lock().expect("store lock poisoned");
            guard.apply_flags_by_uid(user, mailbox, uid, store_args.mode.unwrap(), &flagset);
            let seq = uid_to_seq(&guard, user, mailbox, uid);
            drop(guard);
            updated_seq.push(seq);
        }
    } else {
        let max_seq = store
            .lock()
            .expect("store lock poisoned")
            .list(user, mailbox)
            .len();
        let max_seq = u32::try_from(max_seq).unwrap_or(u32::MAX);
        let targets = parse_sequence_set(&store_args.target, max_seq);
        for seq in targets {
            let mut guard = store.lock().expect("store lock poisoned");
            guard.apply_flags_by_seq(user, mailbox, seq, store_args.mode.unwrap(), &flagset);
            drop(guard);
            updated_seq.push(Some(seq));
        }
    }

    if !silent {
        let messages = store
            .lock()
            .expect("store lock poisoned")
            .list(user, mailbox);
        for seq in updated_seq.into_iter().flatten() {
            if let Some(message) = messages.get(seq.saturating_sub(1) as usize) {
                let flags = format_flags(message);
                let req = FetchRequest {
                    include_flags: true,
                    ..Default::default()
                };
                write_fetch_response(writer, seq, message, &flags, &req).await?;
            }
        }
    }
    Ok(StoreResult::Ok)
}

enum AppendResult {
    Ok { uid: u32 },
    No,
    Bad,
}

async fn handle_append(
    reader: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    line: &str,
    authenticated: bool,
    mailbox_notifier: &MailboxNotifier,
) -> io::Result<AppendResult> {
    let flags = parse_flag_list_from_line(line);
    let mailbox = parse_append_mailbox(line);
    let internal_date = parse_internal_date_from_line(line).unwrap_or_else(current_internal_date);
    let mailbox_exists = {
        let mut guard = store.lock().expect("store lock poisoned");
        guard
            .list_mailboxes(user)
            .iter()
            .any(|m| m.eq_ignore_ascii_case(&mailbox))
    };
    let Some(literal) = parse_literal_size(line) else {
        return Ok(AppendResult::Bad);
    };

    if !literal.non_sync {
        write_raw(reader, b"+ Ready for literal data\r\n").await?;
    }

    let mut data = vec![0u8; literal.size];
    reader.read_exact(&mut data).await?;

    // Consume the trailing CRLF after literal
    let mut crlf = [0u8; 2];
    let _ = reader.read_exact(&mut crlf).await;

    if !authenticated {
        return Ok(AppendResult::Bad);
    }
    if !mailbox_exists {
        return Ok(AppendResult::No);
    }
    let (new_count, uid) = {
        let mut guard = store.lock().expect("store lock poisoned");
        let uid = if flags.seen || flags.flagged || flags.deleted || flags.answered || flags.draft {
            guard.append_with_flags(user, &mailbox, data, internal_date, &flags)
        } else {
            guard.append(user, &mailbox, data, internal_date)
        };
        let count = guard.list(user, &mailbox).len();
        drop(guard);
        (count, uid)
    };
    // Notify any IDLE connections about the new message
    let _ = mailbox_notifier.send(MailboxEvent {
        user: user.to_string(),
        mailbox: mailbox.clone(),
        new_count,
    });
    if uid == 0 {
        Ok(AppendResult::No)
    } else {
        Ok(AppendResult::Ok { uid })
    }
}

struct StoreArgs {
    target: String,
    mode: Option<FlagOp>,
    flags: FlagSet,
    silent: bool,
}

fn parse_store_args(line: &str) -> StoreArgs {
    let mut parts = line.split_whitespace();
    let _tag = parts.next();
    let cmd = parts.next().unwrap_or("");
    let mut target = String::new();
    let mut mode_token = "";

    if cmd.eq_ignore_ascii_case("UID") {
        let next = parts.next().unwrap_or("");
        if next.eq_ignore_ascii_case("STORE") {
            target = parts.next().unwrap_or("").to_string();
            mode_token = parts.next().unwrap_or("");
        }
    } else {
        target = parts.next().unwrap_or("").to_string();
        mode_token = parts.next().unwrap_or("");
    }

    let mut mode = None;
    let mut silent = false;
    let upper = mode_token.to_ascii_uppercase();
    if upper.starts_with("+FLAGS") {
        mode = Some(FlagOp::Add);
    } else if upper.starts_with("-FLAGS") {
        mode = Some(FlagOp::Remove);
    } else if upper.starts_with("FLAGS") {
        mode = Some(FlagOp::Replace);
    }
    if upper.contains(".SILENT") {
        silent = true;
    }

    let flags = parse_flag_list_from_line(line);
    StoreArgs {
        target,
        mode,
        flags,
        silent,
    }
}

enum CopyResult {
    Ok { copyuid: Option<CopyUid> },
    No,
    Bad,
}

struct CopyUid {
    source: String,
    dest: String,
}

enum MoveResult {
    Ok,
    No,
    Bad,
}

fn format_uid_set(uids: &[u32]) -> String {
    uids.iter()
        .map(|uid| uid.to_string())
        .collect::<Vec<_>>()
        .join(",")
}

async fn handle_expunge(
    writer: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    mailbox: &str,
) -> io::Result<()> {
    let expunged = {
        let mut guard = store.lock().expect("store lock poisoned");
        guard.expunge_deleted(user, mailbox)
    };
    for seq in expunged {
        write_raw(writer, format!("* {} EXPUNGE\r\n", seq).as_bytes()).await?;
    }
    Ok(())
}

async fn handle_uid_expunge(
    writer: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    mailbox: &str,
    seqset: &str,
) -> io::Result<()> {
    let expunged = {
        let mut guard = store.lock().expect("store lock poisoned");
        let messages = guard.list(user, mailbox);
        let max_uid = messages.iter().map(|m| m.uid).max().unwrap_or(0);
        let uids = parse_sequence_set(seqset, max_uid);
        guard.expunge_deleted_by_uid(user, mailbox, &uids)
    };
    for seq in expunged {
        write_raw(writer, format!("* {} EXPUNGE\r\n", seq).as_bytes()).await?;
    }
    Ok(())
}

fn handle_copy(
    _writer: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    mailbox: &str,
    line: &str,
) -> CopyResult {
    let Some((seqset, dest, is_uid)) = parse_copy_target(line) else {
        return CopyResult::Bad;
    };
    let mut guard = store.lock().expect("store lock poisoned");
    let exists = guard
        .list_mailboxes(user)
        .iter()
        .any(|m| m.eq_ignore_ascii_case(&dest));
    if !exists {
        return CopyResult::No;
    }
    let max = u32::try_from(guard.list(user, mailbox).len()).unwrap_or(u32::MAX);
    let seqs = parse_sequence_set(&seqset, max);
    let seqs = if is_uid {
        guard.seqs_from_uids(user, mailbox, &seqs)
    } else {
        seqs
    };
    if seqs.is_empty() {
        return CopyResult::Ok { copyuid: None };
    }

    let source_messages = guard.list(user, mailbox);
    let mut source_uids = Vec::new();
    for seq in &seqs {
        let index = seq.saturating_sub(1) as usize;
        if let Some(message) = source_messages.get(index) {
            source_uids.push(message.uid);
        }
    }

    let dest_before_len = guard.list(user, &dest).len();
    guard.copy_by_seq_set(user, mailbox, &seqs, &dest);
    let dest_after = guard.list(user, &dest);
    let dest_uids: Vec<u32> = dest_after
        .iter()
        .skip(dest_before_len)
        .map(|m| m.uid)
        .collect();

    let copyuid = if !source_uids.is_empty() && source_uids.len() == dest_uids.len() {
        Some(CopyUid {
            source: format_uid_set(&source_uids),
            dest: format_uid_set(&dest_uids),
        })
    } else {
        None
    };
    drop(guard);
    CopyResult::Ok { copyuid }
}

async fn handle_move(
    writer: &mut ImapReader,
    store: &Arc<Mutex<Store>>,
    user: &str,
    mailbox: &str,
    line: &str,
) -> io::Result<MoveResult> {
    let Some((seqset, dest, is_uid)) = parse_move_target(line) else {
        return Ok(MoveResult::Bad);
    };
    let expunged = {
        let mut guard = store.lock().expect("store lock poisoned");
        let exists = guard
            .list_mailboxes(user)
            .iter()
            .any(|m| m.eq_ignore_ascii_case(&dest));
        if !exists {
            return Ok(MoveResult::No);
        }
        let max = u32::try_from(guard.list(user, mailbox).len()).unwrap_or(u32::MAX);
        let seqs = parse_sequence_set(&seqset, max);
        let seqs = if is_uid {
            guard.seqs_from_uids(user, mailbox, &seqs)
        } else {
            seqs
        };
        guard.move_by_seq_set(user, mailbox, &seqs, &dest)
    };
    for seq in &expunged {
        write_raw(writer, format!("* {} EXPUNGE\r\n", seq).as_bytes()).await?;
    }
    Ok(MoveResult::Ok)
}