use cubecl_core as cubecl;
use cubecl_core::prelude::*;

use crate::components::AttentionPrecision;
use crate::components::tile::TileAttention;

#[derive(CubeType)]
/// Key and Value inputs to the Tile Attention
///
/// Key and Value share the same trait because they may
/// be the same reused underlying fragment
pub enum KeyValueTile<AP: AttentionPrecision, TA: TileAttention<AP>> {
    Reuse(ReuseKV<AP, TA>),
    Key(Key<AP, TA>),
    Value(Value<AP, TA>),
}

#[cube]
impl<AP: AttentionPrecision, TA: TileAttention<AP>> KeyValueTile<AP, TA> {
    pub fn new_key_value(#[comptime] config: TA::Config) -> Self {
        Self::new_Reuse(ReuseKV::new(config))
    }

    pub fn new_key(#[comptime] config: TA::Config) -> Self {
        Self::new_Key(Key::new(config))
    }

    pub fn new_value(#[comptime] config: TA::Config) -> Self {
        Self::new_Value(Value::new(config))
    }

    /// Get the underlying key as readable
    pub fn key(&self) -> &TA::KeyValue {
        match self {
            KeyValueTile::Reuse(reuse_kv) => &reuse_kv.fragment,
            KeyValueTile::Key(key) => &key.fragment,
            KeyValueTile::Value(_) => panic!("Tried to access key on value-only fragment"),
        }
    }

    /// Get the underlying key as writable
    pub fn key_mut(&mut self) -> &mut TA::KeyValue {
        match self {
            KeyValueTile::Reuse(reuse_kv) => &mut reuse_kv.fragment,
            KeyValueTile::Key(key) => &mut key.fragment,
            KeyValueTile::Value(_) => panic!("Tried to access key on value-only fragment"),
        }
    }

    /// Get the underlying value as readable
    pub fn value(&self) -> &TA::KeyValue {
        match self {
            KeyValueTile::Reuse(reuse_kv) => &reuse_kv.fragment,
            KeyValueTile::Key(_) => panic!("Tried to access value on key-only fragment"),
            KeyValueTile::Value(value) => &value.fragment,
        }
    }

    /// Get the underlying value as writable
    pub fn value_mut(&mut self) -> &mut TA::KeyValue {
        match self {
            KeyValueTile::Reuse(reuse_kv) => &mut reuse_kv.fragment,
            KeyValueTile::Key(_) => panic!("Tried to access value on key-only fragment"),
            KeyValueTile::Value(value) => &mut value.fragment,
        }
    }
}

#[derive(CubeType)]
pub struct ReuseKV<AP: AttentionPrecision, TA: TileAttention<AP>> {
    pub fragment: TA::KeyValue,
}

#[cube]
impl<AP: AttentionPrecision, TA: TileAttention<AP>> ReuseKV<AP, TA> {
    pub fn new(#[comptime] config: TA::Config) -> Self {
        let fragment = TA::allocate_key_value(config);
        ReuseKV::<AP, TA> { fragment }
    }
}

#[derive(CubeType)]
pub struct Key<AP: AttentionPrecision, TA: TileAttention<AP>> {
    pub fragment: TA::KeyValue,
}

#[cube]
impl<AP: AttentionPrecision, TA: TileAttention<AP>> Key<AP, TA> {
    pub fn new(#[comptime] config: TA::Config) -> Self {
        Key::<AP, TA> {
            fragment: TA::allocate_key(config),
        }
    }
}

#[derive(CubeType)]
pub struct Value<AP: AttentionPrecision, TA: TileAttention<AP>> {
    pub fragment: TA::KeyValue,
}

#[cube]
impl<AP: AttentionPrecision, TA: TileAttention<AP>> Value<AP, TA> {
    pub fn new(#[comptime] config: TA::Config) -> Self {
        Value::<AP, TA> {
            fragment: TA::allocate_value(config),
        }
    }
}
