use aes::Aes128;
use base64::{engine::general_purpose::STANDARD, Engine};
use cbc::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use hmac::{Hmac, Mac};
use image::{png::PngEncoder, ColorType, ImageBuffer, Rgba};
use sha2::Sha256;

use crate::error::BrowserCoreError;

type Aes128CbcEnc = cbc::Encryptor<Aes128>;
type Aes128CbcDec = cbc::Decryptor<Aes128>;
type HmacSha256 = Hmac<Sha256>;

const WIDTH: u32 = 350;
const HEIGHT: u32 = 350;
const ANCHORS: [u32; 11] = [0, 35, 70, 105, 140, 175, 210, 245, 280, 315, 349];
const HEADER: usize = 4;
const CAPACITY: usize = ANCHORS.len() * WIDTH as usize - HEADER;

const COLORS: &[(char, (u8, u8, u8))] = &[
    ('a', (204, 180, 194)),
    ('A', (255, 255, 255)),
    ('b', (197, 186, 201)),
    ('B', (221, 206, 212)),
    ('c', (181, 185, 193)),
    ('C', (184, 201, 223)),
    ('d', (224, 218, 192)),
    ('D', (185, 191, 195)),
    ('e', (181, 197, 198)),
    ('E', (193, 206, 255)),
    ('f', (252, 193, 211)),
    ('F', (183, 192, 229)),
    ('g', (180, 191, 192)),
    ('G', (187, 219, 189)),
    ('h', (195, 187, 234)),
    ('H', (182, 216, 189)),
    ('i', (197, 183, 248)),
    ('I', (200, 182, 204)),
    ('j', (255, 235, 196)),
    ('J', (194, 186, 228)),
    ('k', (199, 238, 239)),
    ('K', (208, 247, 234)),
    ('l', (244, 214, 189)),
    ('L', (187, 243, 239)),
    ('m', (188, 231, 238)),
    ('M', (187, 197, 227)),
    ('n', (186, 240, 191)),
    ('N', (187, 198, 206)),
    ('o', (205, 193, 184)),
    ('O', (191, 187, 197)),
    ('p', (194, 200, 206)),
    ('P', (195, 183, 229)),
    ('q', (182, 219, 196)),
    ('Q', (238, 216, 184)),
    ('r', (199, 181, 208)),
    ('R', (239, 231, 198)),
    ('s', (189, 188, 230)),
    ('S', (242, 192, 230)),
    ('t', (199, 199, 199)),
    ('T', (188, 190, 230)),
    ('u', (230, 180, 253)),
    ('U', (241, 247, 247)),
    ('v', (242, 190, 199)),
    ('V', (230, 247, 234)),
    ('w', (197, 186, 249)),
    ('W', (194, 247, 249)),
    ('x', (242, 182, 246)),
    ('X', (188, 222, 193)),
    ('y', (188, 194, 183)),
    ('Y', (197, 195, 197)),
    ('z', (187, 249, 240)),
    ('Z', (233, 231, 242)),
    ('0', (195, 184, 218)),
    ('1', (232, 180, 196)),
    ('2', (191, 193, 196)),
    ('3', (185, 186, 186)),
    ('4', (191, 247, 180)),
    ('5', (187, 199, 248)),
    ('6', (248, 198, 184)),
    ('7', (243, 195, 184)),
    ('8', (232, 192, 208)),
    ('9', (239, 197, 183)),
    ('/', (199, 187, 241)),
    ('+', (195, 216, 223)),
    ('=', (193, 211, 184)),
];

pub fn decrypt_private_key(image_base64: &str, password: &str) -> Result<String, BrowserCoreError> {
    let encrypted = decode_image(image_base64)?;
    let packed = STANDARD.decode(encrypted).map_err(|_| {
        BrowserCoreError::InvalidRequest("legacy wallet encryption is invalid".into())
    })?;
    if packed.len() < 49 {
        return Err(BrowserCoreError::InvalidRequest(
            "legacy wallet payload is too short".into(),
        ));
    }
    let key = normalized_key(password);
    let (iv, remainder) = packed.split_at(16);
    let (expected_hmac, ciphertext) = remainder.split_at(32);
    let mut mac = HmacSha256::new_from_slice(&key).unwrap();
    mac.update(ciphertext);
    mac.verify_slice(expected_hmac).map_err(|_| {
        BrowserCoreError::InvalidRequest("wallet password or integrity check failed".into())
    })?;
    let plaintext = Aes128CbcDec::new_from_slices(&key, iv)
        .unwrap()
        .decrypt_padded_vec_mut::<Pkcs7>(ciphertext)
        .map_err(|_| BrowserCoreError::InvalidRequest("wallet decryption failed".into()))?;
    String::from_utf8(plaintext)
        .map_err(|_| BrowserCoreError::InvalidRequest("wallet private key is not UTF-8".into()))
}

pub fn encrypt_private_key(private_key: &str, password: &str) -> Result<String, BrowserCoreError> {
    if private_key.len() < 10 {
        return Err(BrowserCoreError::InvalidPrivateKey);
    }
    let key = normalized_key(password);
    let iv_text = STANDARD.encode(&private_key.as_bytes()[..10]);
    let iv = iv_text.as_bytes();
    let ciphertext = Aes128CbcEnc::new_from_slices(&key, iv)
        .unwrap()
        .encrypt_padded_vec_mut::<Pkcs7>(private_key.as_bytes());
    let mut mac = HmacSha256::new_from_slice(&key).unwrap();
    mac.update(&ciphertext);
    let mut packed = Vec::with_capacity(48 + ciphertext.len());
    packed.extend_from_slice(iv);
    packed.extend_from_slice(&mac.finalize().into_bytes());
    packed.extend_from_slice(&ciphertext);
    encode_image(&STANDARD.encode(packed))
}

fn normalized_key(password: &str) -> [u8; 16] {
    let mut key = [0u8; 16];
    let bytes = password.as_bytes();
    key[..bytes.len().min(16)].copy_from_slice(&bytes[..bytes.len().min(16)]);
    key
}

fn vivid((r, g, b): (u8, u8, u8)) -> (u8, u8, u8) {
    let average = (r as f32 + g as f32 + b as f32) / 3.0;
    let apply = |value: u8| {
        let saturated = average + ((value as f32 - average) * 2.35);
        (((saturated - 128.0) * 1.12) + 128.0 - 18.0)
            .clamp(0.0, 255.0)
            .round() as u8
    };
    (apply(r), apply(g), apply(b))
}

fn color(character: char) -> Option<(u8, u8, u8)> {
    COLORS
        .iter()
        .find(|(candidate, _)| *candidate == character)
        .map(|(_, value)| vivid(*value))
}

fn character(rgb: (u8, u8, u8)) -> Option<char> {
    COLORS
        .iter()
        .find(|(_, value)| vivid(*value) == rgb)
        .map(|(character, _)| *character)
}

fn decode_image(image_base64: &str) -> Result<String, BrowserCoreError> {
    let bytes = STANDARD
        .decode(image_base64)
        .map_err(|_| BrowserCoreError::InvalidRequest("wallet image is not valid base64".into()))?;
    let image = image::load_from_memory_with_format(&bytes, image::ImageFormat::Png)
        .map_err(|_| BrowserCoreError::InvalidRequest("wallet image is not a PNG".into()))?
        .to_rgba8();
    if image.width() != WIDTH || image.height() != HEIGHT {
        return Err(BrowserCoreError::InvalidRequest(
            "wallet image dimensions are invalid".into(),
        ));
    }
    let mut text = String::with_capacity(ANCHORS.len() * WIDTH as usize);
    for row in ANCHORS {
        for x in 0..WIDTH {
            let pixel = image.get_pixel(x, row).0;
            text.push(character((pixel[0], pixel[1], pixel[2])).ok_or_else(|| {
                BrowserCoreError::InvalidRequest(
                    "wallet image contains an unknown anchor color".into(),
                )
            })?);
        }
    }
    let length = text
        .get(..HEADER)
        .and_then(|value| value.parse::<usize>().ok())
        .ok_or_else(|| {
            BrowserCoreError::InvalidRequest("wallet image length header is invalid".into())
        })?;
    text.get(HEADER..HEADER + length)
        .map(str::to_owned)
        .ok_or_else(|| {
            BrowserCoreError::InvalidRequest("wallet image payload is incomplete".into())
        })
}

fn encode_image(ciphertext: &str) -> Result<String, BrowserCoreError> {
    if ciphertext.len() > CAPACITY {
        return Err(BrowserCoreError::InvalidRequest(
            "encrypted wallet exceeds image capacity".into(),
        ));
    }
    let payload = format!("{:04}{}", ciphertext.len(), ciphertext);
    let chars: Vec<char> = payload.chars().collect();
    let slot_count = ANCHORS.len() * WIDTH as usize;
    let slots: Vec<char> = (0..slot_count)
        .map(|index| chars[index % chars.len()])
        .collect();
    let mut image: ImageBuffer<Rgba<u8>, Vec<u8>> = ImageBuffer::new(WIDTH, HEIGHT);
    for x in 0..WIDTH as usize {
        for anchor in 0..ANCHORS.len() - 1 {
            let upper_y = ANCHORS[anchor];
            let lower_y = ANCHORS[anchor + 1];
            let upper = color(slots[anchor * WIDTH as usize + x]).unwrap();
            let lower = color(slots[(anchor + 1) * WIDTH as usize + x]).unwrap();
            let distance = lower_y - upper_y;
            for y in upper_y..=lower_y {
                let ratio = (y - upper_y) as f32 / distance as f32;
                let blend = |start: u8, end: u8| {
                    (start as f32 + ((end as f32 - start as f32) * ratio)).round() as u8
                };
                image.put_pixel(
                    x as u32,
                    y,
                    Rgba([
                        blend(upper.0, lower.0),
                        blend(upper.1, lower.1),
                        blend(upper.2, lower.2),
                        255,
                    ]),
                );
            }
        }
    }
    let mut png = Vec::new();
    PngEncoder::new(&mut png)
        .encode(image.as_raw(), WIDTH, HEIGHT, ColorType::Rgba8)
        .map_err(|_| BrowserCoreError::InvalidRequest("failed to encode wallet image".into()))?;
    Ok(STANDARD.encode(png))
}

#[cfg(test)]
mod tests {
    use super::{decrypt_private_key, encrypt_private_key};

    #[test]
    fn legacy_private_key_image_round_trips() {
        let private_key = "ab".repeat(1281);
        let encrypted = encrypt_private_key(&private_key, "correct horse battery staple").unwrap();
        let decrypted = decrypt_private_key(&encrypted, "correct horse battery staple").unwrap();
        assert_eq!(decrypted, private_key);
    }

    #[test]
    fn legacy_private_key_image_rejects_wrong_password() {
        let private_key = "cd".repeat(1281);
        let encrypted = encrypt_private_key(&private_key, "correct horse battery staple").unwrap();
        assert!(decrypt_private_key(&encrypted, "this password is incorrect").is_err());
    }
}
