// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::cell::Cell;
use core::ops::{Index, IndexMut};
use kernel::common::cells::{OptionalCell, TakeCell};
use kernel::hil;

pub static mut TXBUFFER: [u8; PAGE_SIZE + 5] = [0; PAGE_SIZE + 5];
pub static mut RXBUFFER: [u8; PAGE_SIZE + 5] = [0; PAGE_SIZE + 5];
const PAGE_SIZE: usize = 512;

pub struct NexusSpiflashPage(pub [u8; PAGE_SIZE]);

impl Default for NexusSpiflashPage {
    fn default() -> Self {
        Self { 0: [0; PAGE_SIZE] }
    }
}

impl Index<usize> for NexusSpiflashPage {
    type Output = u8;
    fn index(&self, idx: usize) -> &u8 {
        &self.0[idx]
    }
}

impl IndexMut<usize> for NexusSpiflashPage {
    fn index_mut(&mut self, idx: usize) -> &mut u8 {
        &mut self.0[idx]
    }
}

impl AsMut<[u8]> for NexusSpiflashPage {
    fn as_mut(&mut self) -> &mut [u8] {
        &mut self.0
    }
}

enum Opcodes {
    Read = 0x03,
}

#[derive(Clone, Copy, PartialEq)]
enum State {
    Idle,
    WriteOpcode,
    ReadPage { offset: usize },
}

pub struct NexusSpiflash<'a, S: hil::spi::SpiMasterDevice + 'a> {
    spi: &'a S,
    state: Cell<State>,
    txbuffer: TakeCell<'static, [u8]>,
    rxbuffer: TakeCell<'static, [u8]>,
    client: OptionalCell<&'a dyn hil::flash::Client<NexusSpiflash<'a, S>>>,
    client_page: TakeCell<'static, NexusSpiflashPage>,
}

impl<'a, S: hil::spi::SpiMasterDevice + 'a> NexusSpiflash<'a, S> {
    pub fn new(
        spi: &'a S,
        txbuffer: &'static mut [u8],
        rxbuffer: &'static mut [u8],
    ) -> NexusSpiflash<'a, S> {
        NexusSpiflash {
            spi: spi,
            state: Cell::new(State::Idle),
            txbuffer: TakeCell::new(txbuffer),
            rxbuffer: TakeCell::new(rxbuffer),
            client: OptionalCell::empty(),
            client_page: TakeCell::empty(),
        }
    }
}

impl<'a, S: hil::spi::SpiMasterDevice + 'a, C: hil::flash::Client<Self>>
    hil::flash::HasClient<'a, C> for NexusSpiflash<'a, S>
{
    fn set_client(&'a self, client: &'a C) {
        self.client.set(client);
    }
}

impl<'a, S: hil::spi::SpiMasterDevice + 'a> hil::flash::Flash for NexusSpiflash<'a, S> {
    type Page = NexusSpiflashPage;
    fn read_page(
        &self,
        page: usize,
        read_buffer: &'static mut <Self as kernel::hil::flash::Flash>::Page,
    ) -> core::result::Result<
        (),
        (
            kernel::ReturnCode,
            &'static mut <Self as kernel::hil::flash::Flash>::Page,
        ),
    > {
        self.txbuffer.take().map(|txbuffer| {
            self.rxbuffer.take().map(move |rxbuffer| {
                txbuffer[0] = Opcodes::Read as u8;
                // TODO(atv): See if we can make 4-byte addressing work. Otherwise limited to 2**24 bytes.
                let byte = page * PAGE_SIZE;
                txbuffer[1] = ((byte >> 16) & 0xFF) as u8;
                txbuffer[2] = ((byte >> 08) & 0xFF) as u8;
                txbuffer[3] = ((byte >> 00) & 0xFF) as u8;
                self.state.set(State::WriteOpcode);
                self.spi.hold_low();
                self.spi
                    .read_write_bytes(txbuffer, Some(rxbuffer), 4 as usize);
            });
        });
        self.client_page.replace(read_buffer);
        Ok(())
    }
    fn write_page(
        &self,
        _: usize,
        _: &'static mut <Self as kernel::hil::flash::Flash>::Page,
    ) -> core::result::Result<
        (),
        (
            kernel::ReturnCode,
            &'static mut <Self as kernel::hil::flash::Flash>::Page,
        ),
    > {
        todo!()
    }
    fn erase_page(&self, _: usize) -> kernel::ReturnCode {
        todo!()
    }
}

impl<'a, S: hil::spi::SpiMasterDevice + 'a> hil::time::AlarmClient for NexusSpiflash<'a, S> {
    fn alarm(&self) {
        todo!()
    }
}

impl<'a, S: hil::spi::SpiMasterDevice + 'a> hil::spi::SpiMasterClient for NexusSpiflash<'a, S> {
    fn read_write_done(
        &self,
        write_buffer: &'static mut [u8],
        read_buffer: core::option::Option<&'static mut [u8]>,
        _len: usize,
    ) {
        match self.state.get() {
            State::WriteOpcode => {
                read_buffer.map(move |read_buffer| {
                    self.state.set(State::ReadPage { offset: 0 });
                    self.spi
                        .read_write_bytes(write_buffer, Some(read_buffer), PAGE_SIZE / 2);
                });
            }
            State::ReadPage { offset } => {
                read_buffer.map(move |read_buffer| {
                    self.client_page.map(|page| {
                        for i in 0..(PAGE_SIZE / 2) {
                            page[i + offset] = read_buffer[i];
                        }
                    });
                    let next_offset = offset + (PAGE_SIZE / 2);
                    if next_offset == PAGE_SIZE {
                        self.state.set(State::Idle);
                        self.client_page.take().map(|page| {
                            self.client.map(move |client| {
                                self.txbuffer.replace(read_buffer);
                                self.rxbuffer.replace(write_buffer);
                                client.read_complete(page, hil::flash::Error::CommandComplete);
                            });
                        });
                    } else {
                        self.state.set(State::ReadPage {
                            offset: next_offset,
                        });
                        // release_low will cause our SPI host to release CS after the read/write.
                        if next_offset + (PAGE_SIZE / 2) == PAGE_SIZE {
                            self.spi.release_low();
                        }
                        self.spi
                            .read_write_bytes(write_buffer, Some(read_buffer), PAGE_SIZE / 2);
                    }
                });
            }
            State::Idle => {
                panic!();
            }
        }
    }
}