use crate::{error::{PdfMakeError, Result}, model::{Alignment, ImageCover, ImageSource, VerticalAlignment}};
use base64::Engine;
use serde::{Deserialize, Serialize};
use std::sync::Arc;

#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum ImageFormat {
    Png,
    Jpeg,
}

#[derive(Debug, Clone, PartialEq)]
pub struct LoadedImage {
    pub data: Arc<Vec<u8>>,
    pub format: ImageFormat,
    pub width: u32,
    pub height: u32,
}

pub fn load_image(source: &ImageSource) -> Result<LoadedImage> {
    let data: Arc<Vec<u8>> = match source {
        ImageSource::Bytes(bytes) => Arc::new(bytes.clone()),
        ImageSource::Path(path) => Arc::new(std::fs::read(path).map_err(|source| PdfMakeError::ImageRead { path: path.clone(), source })?),
        ImageSource::Base64(encoded) => {
            let decoded = decode_base64_image(encoded)?;
            Arc::new(decoded)
        }
    };
    parse_image_data(data)
}

fn decode_base64_image(encoded: &str) -> Result<Vec<u8>> {
    let b64_data = if let Some(comma_pos) = encoded.find(',') {
        let prefix = &encoded[..comma_pos];
        if !prefix.contains("base64") {
            return Err(PdfMakeError::InvalidBase64("missing base64 marker in data URI".into()));
        }
        &encoded[comma_pos + 1..]
    } else {
        encoded
    };

    base64::engine::general_purpose::STANDARD
        .decode(b64_data.trim())
        .map_err(|e| PdfMakeError::InvalidBase64(e.to_string()))
}

pub fn parse_image_data(data: Arc<Vec<u8>>) -> Result<LoadedImage> {
    if let Some((width, height)) = parse_png_dimensions(&data) {
        return Ok(LoadedImage { data, format: ImageFormat::Png, width, height });
    }
    if let Some((width, height)) = parse_jpeg_dimensions(&data) {
        return Ok(LoadedImage { data, format: ImageFormat::Jpeg, width, height });
    }
    Err(PdfMakeError::UnsupportedImageFormat)
}

fn parse_png_dimensions(data: &[u8]) -> Option<(u32, u32)> {
    if data.len() < 24 || &data[0..8] != b"\x89PNG\r\n\x1a\n" || &data[12..16] != b"IHDR" {
        return None;
    }
    let width = u32::from_be_bytes([data[16], data[17], data[18], data[19]]);
    let height = u32::from_be_bytes([data[20], data[21], data[22], data[23]]);
    if width == 0 || height == 0 { None } else { Some((width, height)) }
}

fn parse_jpeg_dimensions(data: &[u8]) -> Option<(u32, u32)> {
    if data.len() < 4 || data[0] != 0xFF || data[1] != 0xD8 {
        return None;
    }
    let mut i = 2;
    while i + 9 < data.len() {
        if data[i] != 0xFF {
            i += 1;
            continue;
        }
        while i < data.len() && data[i] == 0xFF {
            i += 1;
        }
        if i >= data.len() { return None; }
        let marker = data[i];
        i += 1;
        if marker == 0xD9 || marker == 0xDA { return None; }
        if i + 2 > data.len() { return None; }
        let segment_len = u16::from_be_bytes([data[i], data[i + 1]]) as usize;
        if segment_len < 2 || i + segment_len > data.len() { return None; }
        if matches!(marker, 0xC0 | 0xC1 | 0xC2 | 0xC3 | 0xC5 | 0xC6 | 0xC7 | 0xC9 | 0xCA | 0xCB | 0xCD | 0xCE | 0xCF) {
            if segment_len < 7 { return None; }
            let height = u16::from_be_bytes([data[i + 3], data[i + 4]]) as u32;
            let width = u16::from_be_bytes([data[i + 5], data[i + 6]]) as u32;
            if width == 0 || height == 0 { return None; }
            let orientation = get_exif_orientation(data);
            if matches!(orientation, Some(5..=8)) {
                return Some((height, width));
            }
            return Some((width, height));
        }
        i += segment_len;
    }
    None
}

pub fn get_exif_orientation(data: &[u8]) -> Option<u16> {
    if data.len() < 4 || data[0] != 0xFF || data[1] != 0xD8 {
        return None;
    }
    let mut i = 2;
    while i + 4 < data.len() {
        if data[i] != 0xFF {
            i += 1;
            continue;
        }
        while i < data.len() && data[i] == 0xFF {
            i += 1;
        }
        if i >= data.len() { return None; }
        let marker = data[i];
        i += 1;
        if marker == 0xD9 || marker == 0xDA { return None; }
        if i + 2 > data.len() { return None; }
        let segment_len = u16::from_be_bytes([data[i], data[i + 1]]) as usize;
        if segment_len < 2 || i + segment_len > data.len() { return None; }
        if marker == 0xE1 {
            let app1_start = i + 2;
            let app1_end = i + segment_len;
            if app1_end - app1_start < 6 { return None; }
            if &data[app1_start..app1_start + 6] != b"Exif\0\0" { return None; }
            let tiff_start = app1_start + 6;
            if app1_end - tiff_start < 8 { return None; }
            let byte_order = u16::from_be_bytes([data[tiff_start], data[tiff_start + 1]]);
            let (read_u16, read_u32): (fn(&[u8], usize) -> u16, fn(&[u8], usize) -> u32) = match byte_order {
                0x4D4D => (|d, o| u16::from_be_bytes([d[o], d[o + 1]]), |d, o| u32::from_be_bytes([d[o], d[o + 1], d[o + 2], d[o + 3]])),
                0x4949 => (|d, o| u16::from_le_bytes([d[o], d[o + 1]]), |d, o| u32::from_le_bytes([d[o], d[o + 1], d[o + 2], d[o + 3]])),
                _ => return None,
            };
            let ifd0_offset = read_u32(data, tiff_start + 4) as usize;
            let ifd0_pos = tiff_start + ifd0_offset;
            if ifd0_pos + 2 > app1_end { return None; }
            let num_entries = read_u16(data, ifd0_pos) as usize;
            let entries_start = ifd0_pos + 2;
            if entries_start + num_entries * 12 > app1_end { return None; }
            for entry_idx in 0..num_entries {
                let entry_pos = entries_start + entry_idx * 12;
                let tag = read_u16(data, entry_pos);
                let type_val = read_u16(data, entry_pos + 2);
                if tag == 274 && type_val == 3 {
                    return Some(read_u16(data, entry_pos + 8));
                }
            }
            return None;
        }
        i += segment_len;
    }
    None
}

pub fn resolve_display_size(intrinsic_width: u32, intrinsic_height: u32, width: Option<f32>, height: Option<f32>, fit: Option<crate::model::Fit>) -> Result<(f32, f32)> {
    if intrinsic_width == 0 || intrinsic_height == 0 {
        return Err(PdfMakeError::InvalidImageDimensions);
    }
    let iw = intrinsic_width as f32;
    let ih = intrinsic_height as f32;
    if let Some(fit) = fit {
        if fit.width <= 0.0 || fit.height <= 0.0 { return Err(PdfMakeError::InvalidImageDimensions); }
        let scale = (fit.width / iw).min(fit.height / ih);
        return Ok((iw * scale, ih * scale));
    }
    match (width, height) {
        (Some(w), Some(h)) if w > 0.0 && h > 0.0 => Ok((w, h)),
        (Some(w), None) if w > 0.0 => Ok((w, w * ih / iw)),
        (None, Some(h)) if h > 0.0 => Ok((h * iw / ih, h)),
        (None, None) => Ok((iw, ih)),
        _ => Err(PdfMakeError::InvalidImageDimensions),
    }
}

pub fn calculate_cover(
    img_width: u32, img_height: u32,
    cover: &ImageCover,
) -> (f32, f32, f32, f32) {
    let img_ratio = img_width as f32 / img_height as f32;
    let cover_ratio = cover.width / cover.height;

    let (scaled_w, scaled_h) = if img_ratio > cover_ratio {
        (cover.height * img_ratio, cover.height)
    } else {
        (cover.width, cover.width / img_ratio)
    };

    let x_offset = match cover.align {
        Alignment::Left => 0.0,
        Alignment::Center => (cover.width - scaled_w) / 2.0,
        Alignment::Right => cover.width - scaled_w,
        Alignment::Justify => 0.0,
    };
    let y_offset = match cover.valign {
        VerticalAlignment::Top => 0.0,
        VerticalAlignment::Middle => (cover.height - scaled_h) / 2.0,
        VerticalAlignment::Bottom => cover.height - scaled_h,
    };

    (x_offset, y_offset, scaled_w, scaled_h)
}