1use image::{DynamicImage, ImageBuffer, Luma};
4use std::fmt;
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub enum ShareType {
9    Binary,
11    Grayscale,
13    Color,
15}
16
17#[derive(Debug, Clone)]
19pub struct Share {
20    pub image: DynamicImage,
22    pub share_type: ShareType,
24    pub index: usize,
26    pub total_shares: usize,
28    pub original_width: u32,
30    pub original_height: u32,
31    pub block_size: usize,
33    pub is_meaningful: bool,
35}
36
37impl Share {
38    pub fn new(
40        image: DynamicImage,
41        index: usize,
42        total_shares: usize,
43        original_width: u32,
44        original_height: u32,
45        block_size: usize,
46        is_meaningful: bool,
47    ) -> Self {
48        let share_type = match &image {
49            DynamicImage::ImageLuma8(_) => ShareType::Binary,
50            DynamicImage::ImageLuma16(_) => ShareType::Grayscale,
51            DynamicImage::ImageRgb8(_) | DynamicImage::ImageRgba8(_) => ShareType::Color,
52            _ => ShareType::Grayscale,
53        };
54
55        Self {
56            image,
57            share_type,
58            index,
59            total_shares,
60            original_width,
61            original_height,
62            block_size,
63            is_meaningful,
64        }
65    }
66
67    pub fn dimensions(&self) -> (u32, u32) {
69        (self.image.width(), self.image.height())
70    }
71
72    pub fn to_binary(&self) -> ImageBuffer<Luma<u8>, Vec<u8>> {
74        match &self.image {
75            DynamicImage::ImageLuma8(img) => {
76                let mut binary = ImageBuffer::new(img.width(), img.height());
78                for (x, y, pixel) in img.enumerate_pixels() {
79                    let value = if pixel[0] > 127 { 255 } else { 0 };
80                    binary.put_pixel(x, y, Luma([value]));
81                }
82                binary
83            }
84            _ => {
85                let gray = self.image.to_luma8();
87                let mut binary = ImageBuffer::new(gray.width(), gray.height());
88                for (x, y, pixel) in gray.enumerate_pixels() {
89                    let value = if pixel[0] > 127 { 255 } else { 0 };
90                    binary.put_pixel(x, y, Luma([value]));
91                }
92                binary
93            }
94        }
95    }
96
97    pub fn is_compatible(&self, other: &Share) -> bool {
99        self.dimensions() == other.dimensions()
100            && self.total_shares == other.total_shares
101            && self.original_width == other.original_width
102            && self.original_height == other.original_height
103            && self.block_size == other.block_size
104    }
105
106    pub fn stack_binary(&self, other: &Share) -> ImageBuffer<Luma<u8>, Vec<u8>> {
108        let binary1 = self.to_binary();
109        let binary2 = other.to_binary();
110
111        let (width, height) = self.dimensions();
112        let mut result = ImageBuffer::new(width, height);
113
114        for y in 0..height {
115            for x in 0..width {
116                let p1 = binary1.get_pixel(x, y)[0];
117                let p2 = binary2.get_pixel(x, y)[0];
118                let value = if p1 == 0 && p2 == 0 { 0 } else { 255 };
121                result.put_pixel(x, y, Luma([value]));
122            }
123        }
124
125        result
126    }
127
128    pub fn stack_and(&self, other: &Share) -> ImageBuffer<Luma<u8>, Vec<u8>> {
130        let binary1 = self.to_binary();
131        let binary2 = other.to_binary();
132
133        let (width, height) = self.dimensions();
134        let mut result = ImageBuffer::new(width, height);
135
136        for y in 0..height {
137            for x in 0..width {
138                let p1 = binary1.get_pixel(x, y)[0];
139                let p2 = binary2.get_pixel(x, y)[0];
140                let value = if p1 == 0 || p2 == 0 { 0 } else { 255 };
142                result.put_pixel(x, y, Luma([value]));
143            }
144        }
145
146        result
147    }
148
149    pub fn save(&self, path: &str) -> Result<(), image::ImageError> {
151        self.image.save(path)
152    }
153}
154
155impl fmt::Display for Share {
156    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
157        write!(
158            f,
159            "Share {}/{}: {:?} {}x{} (original: {}x{}, block_size: {})",
160            self.index,
161            self.total_shares,
162            self.share_type,
163            self.image.width(),
164            self.image.height(),
165            self.original_width,
166            self.original_height,
167            self.block_size
168        )
169    }
170}
171
172pub fn stack_shares(shares: &[Share]) -> Option<ImageBuffer<Luma<u8>, Vec<u8>>> {
174    if shares.is_empty() {
175        return None;
176    }
177
178    let first = &shares[0];
180    for share in shares.iter().skip(1) {
181        if !first.is_compatible(share) {
182            return None;
183        }
184    }
185
186    let mut result = shares[0].to_binary();
188
189    for share in shares.iter().skip(1) {
190        let binary = share.to_binary();
191        for y in 0..result.height() {
192            for x in 0..result.width() {
193                let current = result.get_pixel(x, y)[0];
194                let new = binary.get_pixel(x, y)[0];
195                let value = if current == 0 && new == 0 { 0 } else { 255 };
196                result.put_pixel(x, y, Luma([value]));
197            }
198        }
199    }
200
201    Some(result)
202}
203
204pub fn progressive_stack(shares: &[Share]) -> Vec<ImageBuffer<Luma<u8>, Vec<u8>>> {
206    let mut results = Vec::new();
207
208    for i in 2..=shares.len() {
209        if let Some(stacked) = stack_shares(&shares[0..i]) {
210            results.push(stacked);
211        }
212    }
213
214    results
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn test_share_creation() {
223        let img = DynamicImage::new_luma8(100, 100);
224        let share = Share::new(img, 1, 3, 50, 50, 2, false);
225
226        assert_eq!(share.index, 1);
227        assert_eq!(share.total_shares, 3);
228        assert_eq!(share.share_type, ShareType::Binary);
229        assert_eq!(share.dimensions(), (100, 100));
230    }
231
232    #[test]
233    fn test_share_compatibility() {
234        let img1 = DynamicImage::new_luma8(100, 100);
235        let img2 = DynamicImage::new_luma8(100, 100);
236        let img3 = DynamicImage::new_luma8(200, 200);
237
238        let share1 = Share::new(img1, 1, 3, 50, 50, 2, false);
239        let share2 = Share::new(img2, 2, 3, 50, 50, 2, false);
240        let share3 = Share::new(img3, 3, 3, 50, 50, 2, false);
241
242        assert!(share1.is_compatible(&share2));
243        assert!(!share1.is_compatible(&share3));
244    }
245}