1use crate::fits::Fits;
2use crate::hdu::{ExtensionHDU, HDU};
3use crate::header::header::BLOCK_NUM_BYTES;
4use crate::header::{ExtensionType, Header};
5use crate::slice_ascii_table_hdu::SliceAsciiTableHDU;
6use crate::slice_bin_table_hdu::SliceBinTableHDU;
7use crate::slice_image_hdu::SliceImageHDU;
8use log::debug;
9use std::error::Error;
10use std::io::{Cursor, Seek, SeekFrom};
11use std::sync::Arc;
12
13#[derive(Debug, Clone)]
19pub struct FitsSlice {
20 primary_hdu: SliceImageHDU,
21 extension_hdus: Vec<ExtensionHDU<Self>>,
22}
23
24impl FitsSlice {
25 pub fn from_slice(data: &[u8]) -> Result<Self, Box<dyn Error + Send + Sync>> {
31 Self::from_vec(data.to_vec())
32 }
33
34 pub fn from_vec(data: Vec<u8>) -> Result<Self, Box<dyn Error + Send + Sync>> {
40 #[cfg(feature = "gzip")]
41 let data = if is_gzipped(&data) {
42 decompress(&data)?
43 } else {
44 data
45 };
46
47 let data = Arc::new(data);
48 let mut reader: Box<dyn crate::util::ReadSeek> =
49 Box::new(Cursor::new(crate::util::SharedBytes(Arc::clone(&data))));
50
51 let header =
52 Header::from_reader(&mut reader)?.ok_or("Could not read primary FITS header")?;
53 header.validate_primary()?;
54 debug!("Read primary header: {:?}", header);
55
56 let data_offset = header.bytes_len();
57 let primary_hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
58
59 let mut extension_hdus = vec![];
60 let mut offset = primary_hdu.byte_size();
61
62 loop {
63 if offset as usize >= data.len() {
64 break;
65 }
66
67 reader.seek(SeekFrom::Start(offset))?;
68
69 let Some(header) = Header::from_reader(&mut reader)? else {
70 break;
71 };
72 header.validate_extension()?;
73
74 let extension_type = header
75 .extension()
76 .ok_or("This is not a valid fits extension. Card XTENSION is missing or invalid")?;
77
78 let data_offset = offset as usize + header.bytes_len();
79
80 match extension_type {
81 ExtensionType::Image => {
82 let hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
83 offset += hdu.byte_size();
84 extension_hdus.push(ExtensionHDU::Image(hdu));
85 }
86 ExtensionType::BinTable if header.is_compressed_image() => {
90 let hdu = SliceImageHDU::new(header, Arc::clone(&data), data_offset);
91 offset += hdu.byte_size();
92 extension_hdus.push(ExtensionHDU::Image(hdu));
93 }
94 ExtensionType::BinTable => {
95 let hdu = SliceBinTableHDU::new(header, Arc::clone(&data), data_offset);
96 offset += hdu.byte_size();
97 extension_hdus.push(ExtensionHDU::BinTable(hdu));
98 }
99 ExtensionType::AsciiTable => {
100 let hdu = SliceAsciiTableHDU::new(header, Arc::clone(&data), data_offset);
101 offset += hdu.byte_size();
102 extension_hdus.push(ExtensionHDU::AsciiTable(hdu));
103 }
104 }
105 }
106
107 Ok(Self {
108 primary_hdu,
109 extension_hdus,
110 })
111 }
112
113 #[cfg(feature = "tokio")]
119 pub async fn from_vec_async(data: Vec<u8>) -> Result<Self, Box<dyn Error + Send + Sync>> {
120 tokio::task::spawn_blocking(move || Self::from_vec(data)).await?
121 }
122
123 pub fn new() -> Self {
125 Self {
126 primary_hdu: SliceImageHDU::empty(),
127 extension_hdus: vec![],
128 }
129 }
130}
131
132impl Default for FitsSlice {
133 fn default() -> Self {
134 Self::new()
135 }
136}
137
138impl Fits for FitsSlice {
139 type ImageHDU = SliceImageHDU;
140 type BinTableHDU = SliceBinTableHDU;
141 type AsciiTableHDU = SliceAsciiTableHDU;
142
143 fn primary_hdu(&self) -> &Self::ImageHDU {
144 &self.primary_hdu
145 }
146
147 fn primary_hdu_mut(&mut self) -> &mut Self::ImageHDU {
148 &mut self.primary_hdu
149 }
150
151 fn extension_count(&self) -> usize {
152 self.extension_hdus.len()
153 }
154
155 fn extension_hdu(&self, index: usize) -> Option<&ExtensionHDU<Self>> {
156 self.extension_hdus.get(index)
157 }
158
159 fn extension_hdu_mut(&mut self, index: usize) -> Option<&mut ExtensionHDU<Self>> {
160 self.extension_hdus.get_mut(index)
161 }
162
163 fn extension_hdus(&self) -> impl Iterator<Item = &ExtensionHDU<Self>> {
164 self.extension_hdus.iter()
165 }
166
167 fn extension_hdus_mut(&mut self) -> impl Iterator<Item = &mut ExtensionHDU<Self>> {
168 self.extension_hdus.iter_mut()
169 }
170
171 fn push_extension(&mut self, extension: ExtensionHDU<Self>) {
172 self.extension_hdus.push(extension);
173 }
174
175 fn remove_extension(&mut self, index: usize) -> Option<ExtensionHDU<Self>> {
176 (index < self.extension_hdus.len()).then(|| self.extension_hdus.remove(index))
177 }
178
179 fn to_vec(&self) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
180 let mut bytes = Vec::new();
181
182 if self.primary_hdu.header().is_compressed_image() {
188 append_hdu(&mut bytes, &empty_primary_header(), &[], 0)?;
189 append_hdu(
190 &mut bytes,
191 &self
192 .primary_hdu
193 .header()
194 .conformed(Some(ExtensionType::BinTable)),
195 self.primary_hdu.data_bytes(),
196 0,
197 )?;
198 } else {
199 append_hdu(
200 &mut bytes,
201 &self.primary_hdu.header().conformed(None),
202 self.primary_hdu.data_bytes(),
203 0,
204 )?;
205 }
206
207 for extension in &self.extension_hdus {
208 match extension {
209 ExtensionHDU::Image(hdu) => append_hdu(
212 &mut bytes,
213 &hdu.header()
214 .conformed(Some(extension_type_of(hdu.header()))),
215 hdu.data_bytes(),
216 0,
217 )?,
218 ExtensionHDU::BinTable(hdu) => append_hdu(
219 &mut bytes,
220 &hdu.header().conformed(Some(ExtensionType::BinTable)),
221 hdu.data_bytes(),
222 0,
223 )?,
224 ExtensionHDU::AsciiTable(hdu) => append_hdu(
227 &mut bytes,
228 &hdu.header().conformed(Some(ExtensionType::AsciiTable)),
229 hdu.data_bytes(),
230 b' ',
231 )?,
232 }
233 }
234
235 Ok(bytes)
236 }
237}
238
239#[cfg(feature = "gzip")]
241fn is_gzipped(data: &[u8]) -> bool {
242 data.starts_with(&[0x1f, 0x8b])
243}
244
245#[cfg(feature = "gzip")]
246fn decompress(data: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
247 use std::io::Read;
248
249 let mut decoder = flate2::read::GzDecoder::new(data);
250 let mut decompressed = Vec::new();
251 decoder.read_to_end(&mut decompressed)?;
252
253 Ok(decompressed)
254}
255
256fn append_hdu(
262 bytes: &mut Vec<u8>,
263 header: &Header,
264 data: &[u8],
265 padding: u8,
266) -> Result<(), Box<dyn Error + Send + Sync>> {
267 header.validate_against_data(data.len())?;
268
269 let mut data = data.to_vec();
270 let overhang = data.len() % BLOCK_NUM_BYTES;
271 if overhang != 0 {
272 data.resize(data.len() + BLOCK_NUM_BYTES - overhang, padding);
273 }
274
275 bytes.extend_from_slice(&header.checksummed_bytes(&data));
276 bytes.extend_from_slice(&data);
277
278 Ok(())
279}
280
281fn extension_type_of(header: &Header) -> ExtensionType {
286 if header.is_compressed_image() {
287 ExtensionType::BinTable
288 } else {
289 ExtensionType::Image
290 }
291}
292
293fn empty_primary_header() -> Header {
296 let mut header = Header::default();
297
298 let _ = header.set_card(crate::header::card_keys::EXTEND, true);
301
302 header.conformed(None)
303}