use std::{collections::HashMap, fmt::Debug};
use anyhow::{Context, Result};

use crate::{
    cpu::crossbar::Crossbar,
    instruction_set::Instructions,
    memory_manager::{CoreMemory, MemoryStorable, type_traits::TryToUsize},
};

pub mod crossbar;

#[derive(Debug, Clone)]
pub struct CPU {
    cores: Box<[Core]>,
}

impl CPU {
    pub fn new(num_cores: impl TryToUsize) -> Self {
        let num_cores = num_cores.try_into().expect("num_cores can not be negative");
        let mut cores: Vec<Core> = std::iter::repeat_with(Core::new)
            .take(num_cores + 1)
            .collect();
        Self {
            cores: cores.into(),
        }
    }

    pub fn reserve_crossbar(
        &mut self,
        num_crossbar: impl TryToUsize,
        byte_width: impl TryToUsize,
        height: impl TryToUsize,
    ) {
        let num_crossbar = num_crossbar
            .try_into()
            .expect("num_crossbar can not be negative");
        let byte_width = byte_width
            .try_into()
            .expect("byte_width can not be negative");
        let height = height.try_into().expect("height can not be negative");
        for core in &mut self.cores {
            core.reserve_crossbar(num_crossbar, byte_width, height);
        }
    }

    pub fn host(&mut self) -> &mut Core {
        &mut self.cores[0]
    }

    pub fn core(&mut self, index: impl TryToUsize) -> &mut Core {
        let index = index.try_into().expect("can not be negative");
        &mut self.cores[index]
    }

    pub fn num_core(&self) -> usize {
        self.cores.len()
    }

    pub(crate) fn host_and_cores(&mut self, core: impl TryToUsize) -> (&mut Core, &mut Core) {
        let core = core.try_into().expect("core can not be negative");
        assert_ne!(
            core, 0,
            "Retriving a core with the host that will always be at position 0"
        );
        let (host, cores) = self.cores.split_at_mut(1);
        let host = host.get_mut(0).expect("Missing the host??");
        let core = cores
            .get_mut(core - 1)
            .expect("Requested a core not present ");
        (host, core)
    }

    pub fn get_multiple_cores<const N: usize>(&mut self, indices: [usize; N]) -> [&mut Core; N] {
        self.cores.get_disjoint_mut(indices).unwrap()
    }
}

#[derive(Debug, Clone)]
pub struct Core {
    crossbars: Vec<Crossbar>,
    memory: CoreMemory,
    registers: [i32; 32],
}

impl Core {
    fn new() -> Self {
        Self {
            crossbars: Vec::new(),
            memory: CoreMemory::new(),
            registers: [0; 32],
        }
    }

    pub fn reserve_crossbar(
        &mut self,
        num_crossbar: impl TryToUsize,
        width: impl TryToUsize,
        height: impl TryToUsize,
    ) {
        let num_crossbar = num_crossbar
            .try_into()
            .expect("num_crossbar can not be negative");
        let width = width.try_into().expect("width can not be negative");
        let height = height.try_into().expect("height can not be negative");
        for _ in 0..num_crossbar {
            let mut crossbar = CoreMemory::new();
            crossbar.set_capacity(width * height);
            self.crossbars.push(Crossbar::new(width, height, crossbar));
        }
    }

    pub fn execute_load<T>(&mut self) -> Result<Vec<&[T]>>
    where
        T: MemoryStorable,
    {
        self.memory.execute_load()
    }

    pub fn execute_store<T>(&mut self, address: impl TryToUsize, element: &[T]) -> Result<()>
    where
        T: MemoryStorable,
    {
        let address = address.try_into().context("address can not be negative")?;
        self.memory.execute_store(address, element)
    }

    pub fn reserve_load(
        &mut self,
        address: impl TryToUsize,
        size: impl TryToUsize,
    ) -> Result<&mut CoreMemory> {
        let address = address.try_into().context("address can not be negative")?;
        let size = size.try_into().context("size can not be negative")?;
        self.memory.reserve_load(address, size)
    }

    pub fn set_register(&mut self, index: impl TryToUsize, value: i32) {
        let index = index.try_into().expect("index can not be negative");
        assert!(
            value >= 0,
            "Register cannot be negative if happens remove this and go check where it's used as usize"
        );
        self.registers[index] = value;
    }

    pub fn register(&mut self, index: impl TryToUsize) -> i32 {
        let index = index.try_into().expect("index can not be negative");
        self.registers[index]
    }

    pub fn load<T>(&mut self, address: impl TryToUsize, size: impl TryToUsize) -> Result<Vec<&[T]>>
    where
        T: MemoryStorable,
    {
        let address = address.try_into().context("address can not be negative")?;
        let size = size.try_into().context("size can not be negative")?;
        self.memory.load(address, size)
    }

    pub fn get_memory_crossbar(&mut self) -> (&mut CoreMemory, &mut Vec<Crossbar>) {
        let Self {
            crossbars,
            memory,
            registers,
        } = self;
        (memory, crossbars)
    }

    pub fn memset(&mut self, address: impl TryToUsize, size: impl TryToUsize, val: u8) -> Result<()> {
        let address = address.try_into().context("address can not be negative")?;
        let size = size.try_into().context("size can not be negative")?;
        self.memory.memset(address, size, val)
    }
}
