use core::{cell::RefCell, fmt::Formatter};

use alloc::rc::Rc;
use embedded_storage::{ReadStorage, Storage};
use esp_bootloader_esp_idf::partitions::PartitionTable;
use esp_hal::time::Instant;
use log::*;
use rmp::{decode::{RmpRead, RmpReadErr}, encode::{RmpWrite, RmpWriteErr, ValueWriteError}};

use crate::simdata::{BundleEventHeader, EventRecord, RmpData, SimDataError, StreamEvent, StreamHeader, StreamIndex, StreamType};

#[derive(Debug)]
pub struct SharedFlash<S> {
    storage: Rc<RefCell<S>>
}

impl<S> Clone for SharedFlash<S> {
    fn clone(&self) -> Self {
        Self {
            storage: Rc::clone(&self.storage)
        }
    }
}

impl<S> SharedFlash<S> {
    pub fn new(storage: S) -> Self {
        Self {
            storage: Rc::new(RefCell::new(storage))
        }
    }
}

impl<S: Storage> Storage for SharedFlash<S> {
    fn write(&mut self, offset: u32, bytes: &[u8]) -> Result<(), Self::Error> {
        self.storage.borrow_mut().write(offset, bytes)
    }
}

impl<S: ReadStorage> ReadStorage for SharedFlash<S> {
    type Error = S::Error;

    fn read(&mut self, offset: u32, bytes: &mut [u8]) -> Result<(), Self::Error> {
        self.storage.borrow_mut().read(offset, bytes)
    }

    fn capacity(&self) -> usize {
        self.storage.borrow().capacity()
    }
}

#[derive(Debug)]
pub enum StorageRangeError<E> {
    OutOfData,
    Storage(E)
}
impl<E: core::fmt::Debug + 'static> RmpReadErr for StorageRangeError<E> {}
impl<E: core::fmt::Debug + 'static> RmpWriteErr for StorageRangeError<E> {}

impl<E> core::fmt::Display for StorageRangeError<E> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.write_str("RmpErr")
    }
}

#[derive(Debug)]
pub struct StorageRange<S> {
    storage: S,
    start: usize,
    end: usize,
    offset: usize
}

impl<S: ReadStorage> StorageRange<S> {
    pub const fn new(storage: S, start: usize, end: usize) -> Self {
        assert!(start <= end);
        Self {
            storage,
            start,
            end,
            offset: 0
        }
    }

    pub const fn capacity(&self) -> usize {
        self.end - self.start
    }

    pub const fn pos(&self) -> usize {
        self.offset
    }

    pub fn seek(&mut self, offset: usize) -> Result<(), StorageRangeError<S::Error>> {
        self.offset += offset;
        if self.offset > self.end {
            Err(StorageRangeError::OutOfData)
        } else {
            Ok(())
        }
    }

    pub fn subset(&self, size: usize) -> Result<Self, StorageRangeError<S::Error>> where S: Clone + core::fmt::Debug {
        trace!("subset {:#02x}:{:#02x} -> {:#02x}:{:#02x}", self.start, self.end, self.start + self.offset, self.start + self.offset + size);
        if self.start + self.offset + size > self.end {
            Err(StorageRangeError::OutOfData)
        } else {
            Ok(Self {
                storage: self.storage.clone(),
                start: self.offset + self.start,
                end: self.start + self.offset + size,
                offset: 0
            })
        }
    }
}

impl<S: ReadStorage> RmpRead for StorageRange<S> where S::Error: core::fmt::Debug + 'static {
    type Error = StorageRangeError<S::Error>;

    fn read_exact_buf(&mut self, buf: &mut [u8]) -> Result<(), StorageRangeError<S::Error>> {
        let pos = self.start + self.offset;
        if pos > self.end {
            Err(StorageRangeError::OutOfData)
        } else {
            assert!(pos + buf.len() <= self.end);
            match self.storage.read(pos as u32, buf) {
                Ok(_) => {
                    self.offset += buf.len();
                    Ok(())
                },
                Err(err) => Err(StorageRangeError::Storage(err))
            }
        }
    }
}


impl<S: Storage> RmpWrite for StorageRange<S> where S::Error: core::fmt::Debug + 'static {
    type Error = StorageRangeError<S::Error>;

    fn write_bytes(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
        let pos = self.start + self.offset;
        if pos > self.end {
            Err(StorageRangeError::OutOfData)
        } else {
            assert!(pos + buf.len() <= self.end);
            match self.storage.write(pos as u32, buf) {
                Ok(_) => {
                    self.offset += buf.len();
                    Ok(())
                },
                Err(err) => Err(StorageRangeError::Storage(err))
            }
        }
    }
}

pub struct SimDataRecorder<S> {
    storage: StorageRange<S>,
    last_stamp: Instant
}

impl<S: Storage + Clone> SimDataRecorder<S> where S::Error: core::fmt::Debug + 'static {
    pub fn open(storage: S, partitions: PartitionTable<'_>) -> Result<Self, SimDataError<ValueWriteError<StorageRangeError<S::Error>>>> {
        let partition_type = esp_bootloader_esp_idf::partitions::PartitionType::Data(
                esp_bootloader_esp_idf::partitions::DataPartitionSubType::Undefined,
        );
        info!("Searching for sim data partition");
        let data_partition = partitions.iter().find(|partition| {
            partition.partition_type() == partition_type && partition.label_as_str() == "sim"
        }).ok_or(SimDataError::PartitionNotFound)?;

        let start = data_partition.offset() as usize;
        let end = data_partition.len() as usize + start;
        let mut writer = StorageRange::new(storage.clone(), start, end);
        warn!("Writing new simulation data at {start:#02x}:{end:#02x}");
        StreamIndex { count: 1 }.write_rmp(&mut writer)?;
        StreamHeader { id: StreamType::Bundle, size: 0 }.write_rmp(&mut writer)?;
        Ok(Self {
            storage: writer,
            last_stamp: Instant::now()
        })
    }

    pub fn write_next<T: EventRecord>(&mut self, event: T) -> Result<(), SimDataError<ValueWriteError<StorageRangeError<S::Error>>>> {
        BundleEventHeader { id: T::stream_id() }.write_rmp(&mut self.storage)?;
        let now = Instant::now();
        let event = StreamEvent {
            data: event,
            timecode: (now - self.last_stamp).as_millis() as f64 / 1000.0
        };
        event.write_rmp(&mut self.storage)?;
        self.last_stamp = now;
        Ok(())
    }
}