use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};

use serde::{ser, Serialize};
use md5::{Md5, Digest};
use aes::cipher::{KeyIvInit, BlockDecryptMut, block_padding::NoPadding, StreamCipher};
use sha2::Sha256;
use hkdf::Hkdf;

const APP_SECRET: &str = "979549437fcc4a3faad4867b5cd25dcb";

type Aes128CbcDec = cbc::Decryptor<aes::Aes128>;
type Aes128Ctr64BE = ctr::Ctr64BE<aes::Aes128>;

#[derive(Serialize)]
pub struct SignedReq {
    pub request_ts: u64,
    pub request_sig: String,
}

pub fn gen_signed_req<T>(object: &str, method: &str, params: &T) -> SignedReq
where T: ser::Serialize + ?Sized {
    let timestamp = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();

    let mut signature = format!("{object}{method}");

    let params = serde_urlencoded::to_string(params).unwrap();
    let params: BTreeMap<String, String> = serde_urlencoded::from_str(&params).unwrap();
    for (key, val) in params {
        signature.push_str(&key);
        signature.push_str(&val);
    };

    signature.push_str(&timestamp.to_string());
    signature.push_str(APP_SECRET);

    let mut hasher = Md5::new();
    hasher.update(signature);
    let signature = hex::encode(hasher.finalize());

    SignedReq {request_ts: timestamp, request_sig: signature}
}

pub fn get_session_key(infos: &str, out: &mut [u8; 16]) {
    let mut infos = infos.split('.');
    let salt = infos.next().unwrap();
    let info = infos.next().unwrap();

    let salt = base64::decode_config(salt, base64::URL_SAFE_NO_PAD).unwrap();
    let info = base64::decode_config(info, base64::URL_SAFE_NO_PAD).unwrap();

    let hk = Hkdf::<Sha256>::new(Some(&salt), &hex::decode(APP_SECRET).unwrap());
    hk.expand(&info, out).unwrap();
}

pub fn get_track_key(key_str: &str, session_key: &[u8; 16], out: &mut [u8; 16]) {
    let mut key = key_str.split('.');
    key.next();

    let mut trk_key = [0u8; 32];
    let mut iv = [0u8; 16];

    base64::decode_config_slice(key.next().unwrap(), base64::URL_SAFE_NO_PAD, &mut trk_key).unwrap();
    base64::decode_config_slice(key.next().unwrap(), base64::URL_SAFE_NO_PAD, &mut iv).unwrap();

    Aes128CbcDec::new(session_key.into(), &iv.into())
        .decrypt_padded_b2b_mut::<NoPadding>(&trk_key[..16], out)
        .unwrap();
}

pub fn decrypt_frame(frame: &[u8], key: &[u8; 16], iv_raw: &[u8]) -> Vec<u8> {
    let mut out = vec![0u8; frame.len()];
    
    let mut iv = [0u8; 16];
    iv[..iv_raw.len()].copy_from_slice(iv_raw);

    Aes128Ctr64BE::new(key.into(), &iv.into())
        .apply_keystream_b2b(frame, &mut out)
        .unwrap();
    out
}