use aes_gcm::Aes256Gcm;
use aes_gcm::aead::{Aead, KeyInit};

use base64::Engine as _;
use rustdns::{Class, Message, Type};
use sha2::{Digest, Sha256};
use std::error::Error;
use std::fmt::Display;
use std::net::UdpSocket;
use std::sync::{OnceLock, RwLock};
use std::time::Duration;

static DEFAULT_MASTER_KEY: OnceLock<[u8; 32]> = OnceLock::new();
static ANONYMOUS_TENANT_SPACE: OnceLock<String> = OnceLock::new();
static ACTIVE_DNS_SERVER: OnceLock<&'static str> = OnceLock::new();
static TEMPORARY_SALT: RwLock<Option<String>> = RwLock::new(None);
static IS_SCOPE_MODE: RwLock<bool> = RwLock::new(false);

use aliyun_rs::Client;
use std::sync::LazyLock;
static ALIYUN_GLOBAL_CLIENT: LazyLock<Client> = LazyLock::new(|| {
    let id = std::env::var("ALIYUN_ACCESS_KEY_ID").expect("ALIYUN_ACCESS_KEY_ID");
    let secret = std::env::var("ALIYUN_ACCESS_KEY_SECRET").expect("ALIYUN_ACCESS_KEY_SECRET");

    Client::new(&id, &secret)
});

#[derive(Debug, Clone)]
pub struct DnsRecordPayload {
    pub record_type: &'static str,
    pub record: String,
    pub domain_name: &'static str,
    pub value: String,
}

pub trait DnsWriteProvider {
    fn add_record(&self, record: &DnsRecordPayload) -> Result<(), Box<dyn Error>>;
}

pub trait WriteRecord {
    async fn write(self) -> Result<(), Box<dyn Error>>;

    fn write_as<P: DnsWriteProvider>(self, provider: &P) -> Result<(), Box<dyn Error>>;
}

impl WriteRecord for Vec<DnsRecordPayload> {
    /// asynchronously writes the DNS records using the default provider (alidns).
    async fn write(self) -> Result<(), Box<dyn Error>> {
        let dns = ALIYUN_GLOBAL_CLIENT.dns();

        for r in self {
            let resp = dns
                .add_record(r.record_type, &r.record, r.domain_name, &r.value)
                .await?;
            println!("Record added: {:?}", resp);
        }
        Ok(())
    }

    fn write_as<P: DnsWriteProvider>(self, provider: &P) -> Result<(), Box<dyn Error>> {
        let mut errors = false;
        for r in self {
            match provider.add_record(&r) {
                Ok(_) => println!("Record added: {:?}", r),
                Err(e) => {
                    eprintln!("Error adding record: {:?}", e);
                    errors = true;
                    continue;
                }
            }
        }

        if errors {
            return Err("One or more records failed to be added.".into());
        }
        Ok(())
    }
}

fn digest_master_key(raw_key: &str) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(raw_key.as_bytes());
    let result = hasher.finalize();
    let mut bytes = [0u8; 32];
    bytes.copy_from_slice(&result[..32]);
    bytes
}

fn derive_tenant_fingerprint(key_bytes: &[u8; 32]) -> String {
    let mut space_hasher = Sha256::new();
    space_hasher.update(key_bytes);
    let space_result = space_hasher.finalize();
    let hex_str: String = space_result.iter().map(|b| format!("{:02x}", b)).collect();
    hex_str[..12].to_string()
}

fn mask_key_domain(raw_key: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(raw_key.as_bytes());
    let result = hasher.finalize();
    let hex_str: String = result.iter().map(|b| format!("{:02x}", b)).collect();
    hex_str[..12].to_string()
}

pub struct KeyGuard;

impl Drop for KeyGuard {
    fn drop(&mut self) {
        if let Ok(mut salt) = TEMPORARY_SALT.write() {
            *salt = None;
        }
        if let Ok(mut scope_mode) = IS_SCOPE_MODE.write() {
            *scope_mode = false;
        }
    }
}

fn derive_twin_key(
    purpose: &str,
    _provided_salt: Option<&str>,
) -> Result<[u8; 32], Box<dyn Error>> {
    let base_master_key = DEFAULT_MASTER_KEY
        .get()
        .ok_or("conf.run error:  root key not initialized.")?;

    let active_salt = {
        if *IS_SCOPE_MODE.read().unwrap() {
            TEMPORARY_SALT.read().unwrap().clone()
        } else {
            TEMPORARY_SALT.write().unwrap().take()
        }
    };

    let mut hasher = Sha256::new();
    hasher.update(base_master_key);
    hasher.update(purpose.as_bytes());
    if let Some(code) = active_salt {
        hasher.update(code.as_bytes());
    }

    let result = hasher.finalize();
    let mut derived_key = [0u8; 32];
    derived_key.copy_from_slice(&result[..32]);
    Ok(derived_key)
}
pub fn use_googledns() {
    let _ = ACTIVE_DNS_SERVER.set("8.8.8.8:53");
}

pub fn use_opendns() {
    let _ = ACTIVE_DNS_SERVER.set("208.67.222.222:53");
}
fn fetch_dns_txt(target_domain: &str) -> Result<Vec<(usize, String)>, Box<dyn Error>> {
    let mut m = Message::default();
    m.add_question(target_domain, Type::TXT, Class::Internet);
    let question_bytes = m.to_vec()?;

    let socket = UdpSocket::bind("0.0.0.0:0")?;
    socket.set_read_timeout(Some(Duration::from_secs_f64(2.0)))?;
    let target_dns = ACTIVE_DNS_SERVER.get().unwrap_or(&"223.5.5.5:53");
    socket.connect(target_dns)?;
    socket.send(&question_bytes)?;

    let mut resp_buf = [0u8; 4096];
    let len = socket.recv(&mut resp_buf)?;
    let answer_message = Message::from_slice(&resp_buf[0..len])?;

    let mut raw_chunks = Vec::new();
    for record in answer_message.answers {
        if record.r#type() == Type::TXT {
            let txt_content = record.resource.to_string();
            let clean_txt = txt_content.replace('"', "");

            let parts: Vec<&str> = clean_txt.splitn(2, ':').collect();
            if parts.len() == 2 {
                if let Ok(index) = parts[0].trim().parse::<usize>() {
                    raw_chunks.push((index, parts[1].replace(',', "").replace(' ', "")));
                    continue;
                }
            }
            raw_chunks.push((0, clean_txt.replace(',', "").replace(' ', "")));
        }
    }
    Ok(raw_chunks)
}

pub fn register(
    master_key_str: &str,
    raw_username: &str,
) -> Result<Vec<DnsRecordPayload>, Box<dyn Error>> {
    let key_bytes = digest_master_key(master_key_str);
    let tenant_fingerprint = derive_tenant_fingerprint(&key_bytes);

    let mut crypto_hasher = Sha256::new();
    crypto_hasher.update(&key_bytes);
    crypto_hasher.update(b"CONF_RUN_CRYPTO_AXIS");
    let crypto_result = crypto_hasher.finalize();
    let mut crypto_key = [0u8; 32];
    crypto_key.copy_from_slice(&crypto_result[..32]);

    let encrypted_base64 = encrypt(raw_username.trim().to_lowercase().as_bytes(), &crypto_key)?;
    let mut dispatch_list = Vec::new();
    dispatch_list.push(DnsRecordPayload {
        record_type: "TXT",
        record: tenant_fingerprint.clone(),
        domain_name: "conf.run",
        value: format!("0:{}", encrypted_base64),
    });

    Ok(dispatch_list)
}

pub fn init_masterkey(master_key_str: &str) {
    let key_bytes = digest_master_key(master_key_str);
    let tenant_fingerprint = derive_tenant_fingerprint(&key_bytes);
    let _ = DEFAULT_MASTER_KEY.set(key_bytes);
    let _ = ANONYMOUS_TENANT_SPACE.set(format!("{}", tenant_fingerprint));
}

pub fn key<T: Display>(short_code: T) {
    if let Ok(mut scope_mode) = IS_SCOPE_MODE.write() {
        *scope_mode = false;
    }
    if let Ok(mut salt) = TEMPORARY_SALT.write() {
        *salt = Some(short_code.to_string());
    }
}

#[must_use = "use let _guard = conf_run::scope_key(...)"]
pub fn scope_key<T: Display>(short_code: T) -> KeyGuard {
    if let Ok(mut scope_mode) = IS_SCOPE_MODE.write() {
        *scope_mode = true;
    }
    if let Ok(mut salt) = TEMPORARY_SALT.write() {
        *salt = Some(short_code.to_string());
    }
    KeyGuard
}

pub fn var(key_domain: &str) -> Result<String, Box<dyn Error>> {
    let crypto_key = derive_twin_key("CONF_RUN_CRYPTO_AXIS", None)?;
    let tenant_space = ANONYMOUS_TENANT_SPACE
        .get()
        .ok_or("conf.run error: Tenant space not initialized.")?;

    let current_domain = format!("{}.{}.conf.run", mask_key_domain(key_domain), tenant_space);
    let mut raw_chunks = fetch_dns_txt(&current_domain)?;
    if raw_chunks.is_empty() {
        return Err(format!("Key [{}] not found", key_domain).into());
    }

    raw_chunks.sort_by_key(|chunk| chunk.0);
    let mut full_payload = String::new();
    for chunk in raw_chunks {
        full_payload.push_str(&chunk.1);
    }

    let decrypted_bytes = decrypt(&full_payload, &crypto_key)?;
    Ok(String::from_utf8(decrypted_bytes)?)
}

pub fn set(key_domain: &str, value: &str) -> Result<Vec<DnsRecordPayload>, Box<dyn Error>> {
    let crypto_key = derive_twin_key("CONF_RUN_CRYPTO_AXIS", None)?;
    let tenant_space = ANONYMOUS_TENANT_SPACE
        .get()
        .ok_or("conf.run error: Tenant space not initialized.")?;

    let masked_key = mask_key_domain(key_domain);
    let encrypted_base64 = encrypt(value.as_bytes(), &crypto_key)?;
    let chunks: Vec<&[u8]> = encrypted_base64.as_bytes().chunks(180).collect();

    let mut dispatch_list = Vec::new();
    for (index, chunk) in chunks.iter().enumerate() {
        let chunk_str = std::str::from_utf8(chunk)?;

        dispatch_list.push(DnsRecordPayload {
            record_type: "TXT",
            record: format!("{}.{}", masked_key, tenant_space),
            domain_name: "conf.run",
            value: format!("{}:{}", index, chunk_str),
        });

        println!(
            "[conf.run] record -> {}.{} | value -> {}:{}",
            masked_key, tenant_space, index, chunk_str
        );
    }
    Ok(dispatch_list)
}

pub fn list() -> Result<String, Box<dyn Error>> {
    let crypto_key = derive_twin_key("CONF_RUN_CRYPTO_AXIS", None)?;
    let tenant_space = ANONYMOUS_TENANT_SPACE
        .get()
        .ok_or("conf.run error: Tenant space not initialized.")?;

    let target_root_domain = format!("{}.conf.run", tenant_space);
    let mut raw_chunks = fetch_dns_txt(&target_root_domain)?;
    println!(
        "[conf.run] list -> {} | total chunks: {}",
        target_root_domain,
        raw_chunks.len()
    );
    if raw_chunks.is_empty() {
        return Err("Tenant root metadata record is empty.".into());
    }

    raw_chunks.sort_by_key(|chunk| chunk.0);
    let mut full_payload = String::new();
    for chunk in raw_chunks {
        full_payload.push_str(&chunk.1);
    }
    let decrypted_bytes = decrypt(&full_payload, &crypto_key)?;

    Ok(String::from_utf8(decrypted_bytes)?)
}

pub fn verify_login(master_key_str: &str) -> Option<String> {
    let key_bytes = digest_master_key(master_key_str);
    let tenant_fingerprint = derive_tenant_fingerprint(&key_bytes);
    let target_dns_path = format!("{}.conf.run", tenant_fingerprint);
    let _ = DEFAULT_MASTER_KEY.set(key_bytes);
    let _ = ANONYMOUS_TENANT_SPACE.set(tenant_fingerprint);

    let crypto_key = match derive_twin_key("CONF_RUN_CRYPTO_AXIS", None) {
        Ok(k) => k,
        Err(_) => return None,
    };

    if let Ok(raw_chunks) = fetch_dns_txt(&target_dns_path) {
        if raw_chunks.is_empty() {
            return None;
        }
        let mut sorted_chunks = raw_chunks;
        sorted_chunks.sort_by_key(|chunk| chunk.0);
        let mut full_payload = String::new();
        for chunk in sorted_chunks {
            full_payload.push_str(&chunk.1);
        }

        if let Ok(decrypted_bytes) = decrypt(&full_payload, &crypto_key) {
            if let Ok(decrypted_name) = String::from_utf8(decrypted_bytes) {
                return Some(decrypted_name.trim().to_lowercase());
            }
        }
    }
    None
}

fn encrypt(data: &[u8], key: &[u8; 32]) -> Result<String, Box<dyn Error>> {
    let cipher = Aes256Gcm::new(key.into());
    let mut nonce_bytes = [0u8; 12];
    rand::fill(&mut nonce_bytes);
    let ciphertext = cipher
        .encrypt(&nonce_bytes.into(), data)
        .map_err(|e| format!("{}", e))?;
    let mut final_payload = nonce_bytes.to_vec();
    final_payload.extend_from_slice(&ciphertext);
    Ok(base64::engine::general_purpose::STANDARD_NO_PAD.encode(&final_payload))
}

fn decrypt(base64_str: &str, key: &[u8; 32]) -> Result<Vec<u8>, Box<dyn Error>> {
    let cipher = Aes256Gcm::new(key.into());
    let mut clean_str = base64_str
        .replace('"', "")
        .replace(' ', "")
        .replace(',', "");
    while clean_str.len() % 4 != 0 {
        clean_str.push('=');
    }

    let encrypted_bytes = match base64::engine::general_purpose::STANDARD.decode(&clean_str) {
        Ok(b) => b,
        Err(_) => base64::engine::general_purpose::STANDARD_NO_PAD.decode(&clean_str)?,
    };
    if encrypted_bytes.len() < 12 {
        return Err("Payload too short".into());
    }
    let (nonce_bytes, ciphertext) = encrypted_bytes.split_at(12);
    let nonce_array: &[u8; 12] = nonce_bytes.try_into().map_err(|_| "Invalid nonce")?;
    Ok(cipher
        .decrypt(nonce_array.into(), ciphertext)
        .map_err(|_| "Decryption failed")?)
}
