pub mod config;
mod api;
mod auth;
mod crypto;
mod imap_server;
mod smtp_server;
mod srp;
mod store;

use std::collections::HashMap;
use std::sync::Arc;

use api::{ApiClient, LABEL_INBOX};
use store::MessageStore;
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tracing::{error, info};
// ── Shared bridge state ───────────────────────────────────────────────────────

/// State shared between the IMAP and SMTP server tasks.
pub struct BridgeState {
    pub store: MessageStore,
    pub key_pool: Vec<crypto::PrivateKey>,
    pub http_client: reqwest::Client,
    pub session: auth::Session,
    pub config: config::Config,
    pub sender_addresses: Vec<(String, String)>,
    pub own_public_keys: HashMap<String, String>,
    pub sender_key_indices: HashMap<String, usize>,
}

pub type SharedState = Arc<Mutex<BridgeState>>;

// ── Public entry point ────────────────────────────────────────────────────────

/// Authenticate with ProtonMail, bind IMAP and SMTP ports, then return `Ok(())`.
/// The servers keep running in background Tokio tasks on a dedicated OS thread.
/// Returns `Err` if authentication or port binding fails.
///
/// This is the integration entry point for tuimail. It is synchronous so the
/// caller (tuimail's main thread) needs no Tokio runtime of its own.
pub fn start(config: config::Config) -> Result<(), String> {
    let (ready_tx, ready_rx) = std::sync::mpsc::channel::<Result<(), String>>();
    std::thread::spawn(move || {
        let rt = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .build()
            .expect("tokio runtime");
        rt.block_on(async move {
            match run(config).await {
                Ok(()) => {
                    let _ = ready_tx.send(Ok(()));
                    // Keep the runtime alive so spawned server tasks keep running.
                    std::future::pending::<()>().await;
                }
                Err(e) => {
                    let _ = ready_tx.send(Err(e));
                }
            }
        });
    });
    ready_rx
        .recv()
        .map_err(|_| "bridge thread exited unexpectedly".to_string())?
}

// ── Async startup ─────────────────────────────────────────────────────────────

async fn run(config: config::Config) -> Result<(), String> {
    let client = auth::build_client().map_err(|e| format!("HTTP client: {e}"))?;
    let session = auth::authenticate(&client, &config.proton)
        .await
        .map_err(|e| format!("Authentication failed: {e}"))?;
    info!("Authenticated as {}", config.proton.username);

    let api = ApiClient::new(&client, &session);

    let (key_pool, sender_addresses, own_public_keys, sender_key_indices) =
        unlock_key_pool(&api, &config).await?;
    info!("{} key(s) in pool", key_pool.len());
    crypto::dump_key_pool(&key_pool);

    let (messages, total) = api
        .list_messages(LABEL_INBOX, 0, 50)
        .await
        .map_err(|e| format!("list_messages: {e}"))?;
    let mut store = MessageStore::new();
    store.load_all(messages);
    info!("Inbox: {} messages ({total} total)", store.count());

    // Bind ports before spawning tasks — ports are occupied when start() returns.
    let imap_listener = TcpListener::bind(("127.0.0.1", config.bridge.imap_port))
        .await
        .map_err(|e| format!("bind IMAP :{}: {e}", config.bridge.imap_port))?;
    let smtp_listener = TcpListener::bind(("127.0.0.1", config.bridge.smtp_port))
        .await
        .map_err(|e| format!("bind SMTP :{}: {e}", config.bridge.smtp_port))?;

    let state: SharedState = Arc::new(Mutex::new(BridgeState {
        store,
        key_pool,
        http_client: client,
        session,
        config,
        sender_addresses,
        own_public_keys,
        sender_key_indices,
    }));

    let imap_state = Arc::clone(&state);
    tokio::spawn(async move {
        if let Err(e) = imap_server::run_with_listener(imap_state, imap_listener).await {
            eprintln!("IMAP server error: {e}");
        }
    });
    tokio::spawn(async move {
        if let Err(e) = smtp_server::run_with_listener(state, smtp_listener).await {
            error!("SMTP server error: {e}");
        }
    });

    Ok(()) // ← ports are bound; start() unblocks here
}

// ── Key pool setup ────────────────────────────────────────────────────────────

async fn unlock_key_pool(
    api: &ApiClient<'_>,
    config: &config::Config,
) -> Result<
    (
        Vec<crypto::PrivateKey>,
        Vec<(String, String)>,
        HashMap<String, String>,
        HashMap<String, usize>,
    ),
    String,
> {
    let (user_keys_res, addresses_res, key_salts_res) = tokio::join!(
        api.get_user_keys(),
        api.get_addresses(),
        api.get_key_salts(),
    );

    let user_keys = user_keys_res.map_err(|e| format!("get_user_keys: {e}"))?;
    let addresses = addresses_res.map_err(|e| format!("get_addresses: {e}"))?;

    let key_salts = match key_salts_res {
        Ok(salts) => salts,
        Err(e) => {
            if e.contains("403") {
                auth::Session::delete();
                return Err(format!(
                    "Session 'locked' scope expired — session.json deleted. \
                     Re-run to enter TOTP and regain 'locked' scope. ({e})"
                ));
            }
            eprintln!("get_key_salts: {e} — continuing without salt");
            vec![]
        }
    };

    let password = config
        .proton
        .mailbox_password
        .as_deref()
        .unwrap_or(&config.proton.password);

    let passphrase_for = |key_id: &str| -> String {
        let salt = key_salts
            .iter()
            .find(|s| s.id == key_id)
            .and_then(|s| s.key_salt.as_deref())
            .unwrap_or("");
        crypto::derive_key_passphrase(password, salt).unwrap_or_else(|_| password.to_string())
    };

    let mut user_private_keys: Vec<crypto::PrivateKey> = Vec::new();
    for uk in &user_keys {
        if uk.active != 1 {
            continue;
        }
        let salt = key_salts
            .iter()
            .find(|s| s.id == uk.id)
            .and_then(|s| s.key_salt.as_deref())
            .unwrap_or("");
        let pp = crypto::derive_key_passphrase(password, salt)
            .unwrap_or_else(|_| password.to_string());
        match crypto::PrivateKey::unlock(&uk.private_key, &pp) {
            Ok(k) => {
                info!("user key {} unlocked", &uk.id[..8.min(uk.id.len())]);
                user_private_keys.push(k);
            }
            Err(e) => error!("user key {} FAILED: {e}", &uk.id[..8.min(uk.id.len())]),
        }
    }

    let mut key_pool: Vec<crypto::PrivateKey> = Vec::new();
    let mut own_public_keys: HashMap<String, String> = HashMap::new();
    let mut sender_key_indices: HashMap<String, usize> = HashMap::new();
    for addr in &addresses {
        let mut first_active_key_done = false;
        for ak in &addr.keys {
            if ak.active != 1 {
                continue;
            }
            let pp = if let Some(token) = &ak.token {
                let mut token_pp = None;
                for uk in &user_private_keys {
                    if let Ok(p) = crypto::decrypt_token(token, uk) {
                        token_pp = Some(p);
                        break;
                    }
                }
                token_pp.unwrap_or_else(|| passphrase_for(&ak.id))
            } else {
                passphrase_for(&ak.id)
            };
            match crypto::PrivateKey::unlock(&ak.private_key, &pp) {
                Ok(k) => {
                    info!(
                        "address key {} ({}) unlocked primary={}",
                        &ak.id[..8.min(ak.id.len())],
                        addr.email,
                        ak.primary
                    );
                    if !first_active_key_done {
                        match k.public_key_armored() {
                            Ok(pk) => {
                                info!(
                                    "  [own-key] {} → first-active key primary={}",
                                    addr.email, ak.primary
                                );
                                own_public_keys.insert(addr.email.clone(), pk);
                                sender_key_indices.insert(addr.email.clone(), key_pool.len());
                            }
                            Err(e) => info!("extract pub key for {}: {e}", addr.email),
                        }
                        first_active_key_done = true;
                    }
                    key_pool.push(k);
                }
                Err(e) => error!(
                    "address key {} ({}): {e}",
                    &ak.id[..8.min(ak.id.len())],
                    addr.email
                ),
            }
        }
    }

    key_pool.extend(user_private_keys);

    let sender_addresses: Vec<(String, String)> = addresses
        .iter()
        .map(|a| (a.id.clone(), a.email.clone()))
        .collect();

    info!("{} own public key(s) from address keys", own_public_keys.len());
    Ok((key_pool, sender_addresses, own_public_keys, sender_key_indices))
}