use rusqlite::Connection;
use aes_gcm::{Aes256Gcm, Key, Nonce};
use aes_gcm::aead::{Aead, KeyInit};
use pbkdf2::pbkdf2;
use hmac::Hmac;
use sha2::Sha256;
use rand::Rng;

pub fn init_db() -> Connection {
    let conn = Connection::open("client.db").unwrap();
    conn.execute(
        "CREATE TABLE IF NOT EXISTS hosts (
            host TEXT PRIMARY KEY,
            last_used DATETIME DEFAULT CURRENT_TIMESTAMP
        )",
        [],
    ).unwrap();
    conn.execute(
        "CREATE TABLE IF NOT EXISTS credentials (
            host TEXT PRIMARY KEY,
            username TEXT NOT NULL,
            encrypted_password BLOB NOT NULL
        )",
        [],
    ).unwrap();
    conn
}

pub fn get_last_host(conn: &Connection) -> Option<String> {
    let mut stmt = conn.prepare("SELECT host FROM hosts ORDER BY last_used DESC LIMIT 1").unwrap();
    let mut rows = stmt.query([]).unwrap();
    if let Ok(Some(row)) = rows.next() {
        row.get(0).ok()
    } else {
        None
    }
}

pub fn save_host(conn: &Connection, host: &str) {
    conn.execute(
        "INSERT OR REPLACE INTO hosts (host, last_used) VALUES (?1, CURRENT_TIMESTAMP)",
        [host],
    ).unwrap();
}

pub fn save_credentials(conn: &Connection, host: &str, username: &str, password: &str) -> Result<(), aes_gcm::Error> {
    let encrypted = encrypt_password(password)?;
    conn.execute(
        "INSERT OR REPLACE INTO credentials (host, username, encrypted_password) VALUES (?1, ?2, ?3)",
        (host, username, &encrypted),
    ).unwrap();
    Ok(())
}

pub fn get_credentials(conn: &Connection, host: &str) -> Option<(String, String)> {
    let mut stmt = conn.prepare("SELECT username, encrypted_password FROM credentials WHERE host = ?1").unwrap();
    let mut rows = stmt.query([host]).unwrap();
    if let Ok(Some(row)) = rows.next() {
        let username: String = row.get(0).unwrap();
        let encrypted: Vec<u8> = row.get(1).unwrap();
        decrypt_password(&encrypted).map(|password| (username, password))
    } else {
        None
    }
}

fn derive_key() -> Key<Aes256Gcm> {
    let mut key = [0u8; 32];
    let salt = b"client_salt"; // Fixed salt for simplicity
    pbkdf2::<Hmac<Sha256>>(b"client_key", salt, 10000, &mut key).expect("error deriving key");
    key.into()
}

fn encrypt_password(password: &str) -> Result<Vec<u8>, aes_gcm::Error> {
    let key = derive_key();
    let cipher = Aes256Gcm::new(&key);
    let mut nonce_bytes = [0u8; 12];
    rand::thread_rng().fill(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);
    let ciphertext = cipher.encrypt(nonce, password.as_bytes())?;
    let mut result = nonce_bytes.to_vec();
    result.extend(ciphertext);
    Ok(result)
}

fn decrypt_password(encrypted: &[u8]) -> Option<String> {
    if encrypted.len() < 12 {
        return None;
    }
    let key = derive_key();
    let cipher = Aes256Gcm::new(&key);
    let nonce = Nonce::from_slice(&encrypted[..12]);
    cipher.decrypt(nonce, &encrypted[12..]).ok().and_then(|plaintext| String::from_utf8(plaintext).ok())
}