use darling::ToTokens;
use proc_macro2::TokenStream;
use quote::{TokenStreamExt, quote};

use super::options::*;

pub struct PopulateNested<'a> {
    column: &'a Column,
    ident: &'a syn::Ident,
    generics: &'a syn::Generics,
    error: &'a syn::Type,
    id: &'a syn::Ident,
    table_name: &'a str,
    events_table_name: &'a str,
    repo_types_mod: syn::Ident,
}

impl<'a> PopulateNested<'a> {
    pub fn new(column: &'a Column, opts: &'a RepositoryOptions) -> Self {
        Self {
            column,
            ident: &opts.ident,
            generics: &opts.generics,
            error: opts.err(),
            id: opts.id(),
            table_name: opts.table_name(),
            events_table_name: opts.events_table_name(),
            repo_types_mod: opts.repo_types_mod(),
        }
    }
}

impl ToTokens for PopulateNested<'_> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let ty = self.column.ty();
        let ident = self.ident;
        let error = self.error;
        let repo_types_mod = &self.repo_types_mod;
        let accessor = self.column.parent_accessor();

        let query = format!(
            "WITH entities AS (SELECT * FROM {} WHERE ({} = ANY($1))) SELECT i.id AS \"entity_id: {}\", e.sequence, e.event, CASE WHEN $2 THEN e.context ELSE NULL::jsonb END as \"context: es_entity::ContextData\", e.recorded_at FROM entities i JOIN {} e ON i.id = e.id ORDER BY e.id, e.sequence",
            self.table_name,
            self.column.name(),
            self.id,
            self.events_table_name,
        );

        let (impl_generics, ty_generics, where_clause) = self.generics.split_for_impl();

        tokens.append_all(quote! {
            impl #impl_generics es_entity::PopulateNested<#ty> for #ident #ty_generics #where_clause {
                async fn populate_in_op<OP, P>(
                    op: &mut OP,
                    mut lookup: std::collections::HashMap<#ty, &mut P>,
                ) -> Result<(), #error>
                where
                    OP: es_entity::AtomicOperation,
                    P: Parent<<Self as EsRepo>::Entity>
                {
                    let parent_ids: Vec<_> = lookup.keys().collect();
                    let rows = {
                        sqlx::query_as!(
                            #repo_types_mod::Repo__DbEvent,
                            #query,
                            parent_ids.as_slice() as &[&#ty],
                            <#repo_types_mod::Repo__Event as EsEvent>::event_context(),
                        ).fetch_all(op.as_executor()).await?
                    };
                    let n = rows.len();
                    let (mut res, _) = es_entity::EntityEvents::load_n::<<Self as EsRepo>::Entity>(rows.into_iter(), n)?;
                    Self::load_all_nested_in_op(op, &mut res).await?;
                    for entity in res.into_iter() {
                        let parent = lookup.get_mut(&entity.#accessor).expect("parent not present");
                        parent.inject_children(std::iter::once(entity));
                    }
                    Ok(())
                }
            }
        });
    }
}
