Unholy DI things

use std::sync::Arc;

use dptree::{case, deps, di::*, from_fn_with_description, HandlerDescription};
use teloxide::{
    dispatching::{dialogue::InMemStorage, HandlerExt, UpdateFilterExt, UpdateHandler},
    macros::BotCommands,
    prelude::*,
};

pub trait InsertContainer {
    fn insert_container(&mut self, cont: DependencyMap);
}

impl InsertContainer for DependencyMap {
    fn insert_container(&mut self, cont: DependencyMap) {
        DependencyMap::insert_container(self, cont);
    }
}

trait InjectContainer<'a, Input, Output, Descr> {
    fn inject_cont<Proj, Args>(self, proj: Proj) -> Handler<'a, Input, Output, Descr>
    where
        Input: Clone + InsertContainer,
        Asyncify<Proj>: Injectable<Input, DependencyMap, Args> + Send + Sync + 'a;
}

impl<'a, Input, Output, Descr> InjectContainer<'a, Input, Output, Descr>
    for Handler<'a, Input, Output, Descr>
where
    Input: Send + 'a,
    Output: 'a,
    Descr: HandlerDescription,
{
    #[must_use]
    #[track_caller]
    fn inject_cont<Proj, Args>(self, proj: Proj) -> Handler<'a, Input, Output, Descr>
    where
        Input: Clone + InsertContainer,
        Asyncify<Proj>: Injectable<Input, DependencyMap, Args> + Send + Sync + 'a,
    {
        self.chain(inject_cont_async_with_description(
            Descr::map(),
            Asyncify(proj),
        ))
    }
}

#[must_use]
pub fn inject_cont_async_with_description<'a, Projection, Input, Output, Args, Descr>(
    description: Descr,
    proj: Projection,
) -> Handler<'a, Input, Output, Descr>
where
    Input: Clone + InsertContainer,
    Projection: Injectable<Input, DependencyMap, Args> + Send + Sync + 'a,
    Input: Send + 'a,
    Output: 'a,
    Descr: HandlerDescription,
{
    let proj = Arc::new(proj);

    from_fn_with_description(description, move |container: Input, cont| {
        let proj = Arc::clone(&proj);

        async move {
            let proj = proj.inject(&container);
            let res = proj().await;
            std::mem::drop(proj);

            let mut intermediate = container.clone();
            intermediate.insert_container(res);
            match cont(intermediate).await {
                ControlFlow::Continue(_) => ControlFlow::Continue(container),
                ControlFlow::Break(result) => ControlFlow::Break(result),
            }
        }
    })
}

type HandlerResult = Result<(), Box<dyn std::error::Error + Send + Sync>>;
type DiDialogue = Dialogue<DiState, InMemStorage<DiState>>;

#[derive(Clone)]
struct DiState {
    di: DependencyMap,
}

impl Default for DiState {
    fn default() -> Self {
        DiState {
            di: DependencyMap::new(),
        }
    }
}

// Just for QoL, not necessary
impl From<DependencyMap> for DiState {
    fn from(value: DependencyMap) -> Self {
        DiState { di: value }
    }
}

#[derive(BotCommands, Default, Clone, Debug)]
#[command(rename_rule = "snake_case")]
enum Commands {
    #[default]
    IAmSure,
}

async fn handler(
    bot: Bot,
    message: Message,
    DiState { mut di }: DiState,
    didialogue: DiDialogue,
) -> HandlerResult {
    let val_before = di.insert(3_i32);
    let val: Arc<i32> = di.get();

    bot.send_message(
        message.chat.id,
        format!("Before: {:?}, Now: {:?}", val_before, val),
    )
    .await?;

    didialogue.update(di).await?;
    Ok(())
}

async fn i_am_sure(bot: Bot, message: Message, val: i32) -> HandlerResult {
    bot.send_message(message.chat.id, format!("{val}")).await?;
    Ok(())
}

fn handler_tree() -> UpdateHandler<Box<dyn std::error::Error + Send + Sync + 'static>> {
    dptree::entry()
        .enter_dialogue::<Update, InMemStorage<DiState>, DiState>()
        .inject_cont(|DiState { di }: DiState| di)
        .branch(
            Update::filter_message()
                .filter_command::<Commands>()
                .branch(case![Commands::IAmSure].endpoint(i_am_sure)),
        )
        .branch(Update::filter_message().endpoint(handler))
}

#[tokio::main]
async fn main() {
    dotenv::dotenv().ok();

    let bot = Bot::from_env();

    Dispatcher::builder(bot, handler_tree())
        .dependencies(deps![InMemStorage::<DiState>::new()])
        .enable_ctrlc_handler()
        .build()
        .dispatch()
        .await;
}