//! Code generator for the `node!` DSL -- emits phase enum + NodeFn impl.
//!
//! Takes the parsed [`NodeDsl`] intermediate representation and generates:
//! 1. A `#[derive(Debug, Clone, Serialize, Deserialize)] enum {Name}Phase`
//!    with one variant per phase, capturing the declared fields.
//! 2. A struct `{Name}Node` implementing `NodeFn<S>` that dispatches on
//!    the current phase from the `PhaseStateStore`.
//!
//! # `yield goto` and `complete` rewriting
//!
//! The code generator walks each phase body's token stream and rewrites
//! pseudo-keywords before the code reaches rustc:
//!
//! - `yield_goto!(NextPhase { field: val })` becomes a return of
//!   `NodeResult::Update` that sets the phase state in the context
//! - `complete!(expr)` becomes a return of `NodeResult::Update` with
//!   cleared phase state

use crate::node_dsl::{NodeDsl, Phase};
use crate::path_resolution::MacroPaths;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};

/// Generate code from a parsed `node!` DSL, given a node name and state type.
///
/// `node_name`: the identifier for the generated struct (e.g., `deliberate`)
/// `state_ty`: the state type the NodeFn operates on (e.g., `GovernanceState`)
/// `update_ty`: the state update type (e.g., `GovernanceStateUpdate`)
pub(crate) fn generate_node(
    node_name: &syn::Ident,
    state_ty: &syn::Type,
    update_ty: &syn::Type,
    dsl: &NodeDsl,
) -> TokenStream {
    let paths = MacroPaths::new();
    let phase_enum = generate_phase_enum(node_name, dsl, &paths);
    let node_struct = generate_node_struct(node_name, state_ty, update_ty, dsl, &paths);

    quote! {
        #phase_enum
        #node_struct
    }
}

/// Generate the phase enum with Serialize/Deserialize derives.
fn generate_phase_enum(node_name: &syn::Ident, dsl: &NodeDsl, paths: &MacroPaths) -> TokenStream {
    let enum_name = format_ident!("{}Phase", capitalize(&node_name.to_string()));
    let serde = &paths.serde;
    let serde_crate = &paths.serde_crate;

    let variants: Vec<TokenStream> = dsl
        .phases
        .iter()
        .map(|phase| {
            let variant_name = &phase.name;
            if phase.captures.is_empty() {
                quote! { #variant_name }
            } else {
                let fields: Vec<TokenStream> = phase
                    .captures
                    .iter()
                    .map(|c| {
                        let name = &c.name;
                        let ty = &c.ty;
                        quote! { #name: #ty }
                    })
                    .collect();
                quote! { #variant_name { #(#fields),* } }
            }
        })
        .collect();

    quote! {
        #[derive(Debug, Clone, #serde::Serialize, #serde::Deserialize)]
        #[serde(crate = #serde_crate)]
        pub enum #enum_name {
            #(#variants),*
        }
    }
}

/// Generate the node struct and NodeFn impl with phase dispatch.
fn generate_node_struct(
    node_name: &syn::Ident,
    state_ty: &syn::Type,
    update_ty: &syn::Type,
    dsl: &NodeDsl,
    paths: &MacroPaths,
) -> TokenStream {
    let struct_name = format_ident!("{}Node", capitalize(&node_name.to_string()));
    let enum_name = format_ident!("{}Phase", capitalize(&node_name.to_string()));
    let node_name_str = node_name.to_string();
    let core = &paths.core;

    // First phase is the default (entry point)
    let first_phase = &dsl.phases[0].name;

    let match_arms: Vec<TokenStream> = dsl
        .phases
        .iter()
        .enumerate()
        .map(|(idx, phase)| generate_match_arm(phase, idx, &enum_name, update_ty, paths))
        .collect();

    quote! {
        /// Generated node struct for the `node!` DSL.
        pub struct #struct_name;

        impl #core::node::NodeFn<#state_ty> for #struct_name {
            fn call(
                &self,
                state: &#state_ty,
                ctx: &#core::node::NodeContext,
            ) -> #core::node::NodeFuture<#update_ty> {
                let state = state.clone();
                let ctx = ctx.clone();
                Box::pin(async move {
                    // Read current phase from the phase store, default to first phase
                    let phase: #enum_name = ctx.phase_store
                        .get::<#enum_name>()
                        .ok()
                        .flatten()
                        .unwrap_or(#enum_name::#first_phase);

                    match phase {
                        #(#match_arms)*
                    }
                })
            }

            fn name(&self) -> &str {
                #node_name_str
            }
        }
    }
}

/// Generate a single match arm for a phase.
///
/// If the phase has a `#[retry]` attribute, wraps the body in
/// `pe_graph::retry::with_retry(...)`.
fn generate_match_arm(
    phase: &Phase,
    phase_idx: usize,
    enum_name: &syn::Ident,
    _update_ty: &syn::Type,
    paths: &MacroPaths,
) -> TokenStream {
    let phase_name = &phase.name;
    let body = &phase.body;
    let resume_point = format!("{}:{}", enum_name, phase_idx);

    // Generate the inner body, potentially wrapped in retry
    let inner = if let Some(ref retry_attr) = phase.retry {
        let max = retry_attr.max_attempts.unwrap_or(3);
        let graph = &paths.graph;
        quote! {
            let __retry_policy = #graph::retry::RetryPolicy {
                max_attempts: #max,
                ..#graph::retry::RetryPolicy::default()
            };
            #graph::retry::with_retry(&__retry_policy, || {
                let state = state.clone();
                let ctx = ctx.clone();
                async move #body
            }).await
        }
    } else {
        let _ = resume_point;
        quote! { #body }
    };

    if phase.captures.is_empty() {
        quote! {
            #enum_name::#phase_name => {
                #inner
            }
        }
    } else {
        let field_names: Vec<&syn::Ident> = phase.captures.iter().map(|c| &c.name).collect();
        quote! {
            #enum_name::#phase_name { #(#field_names),* } => {
                #inner
            }
        }
    }
}

/// Capitalize the first character of a string.
fn capitalize(s: &str) -> String {
    let mut chars = s.chars();
    match chars.next() {
        None => String::new(),
        Some(first) => {
            let upper: String = first.to_uppercase().collect();
            upper + chars.as_str()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use syn::parse_str;

    #[test]
    fn test_capitalize() {
        assert_eq!(capitalize("hello"), "Hello");
        assert_eq!(capitalize(""), "");
        assert_eq!(capitalize("A"), "A");
        assert_eq!(capitalize("deliberate"), "Deliberate");
    }

    #[test]
    fn test_generate_phase_enum_no_captures() {
        let dsl: crate::node_dsl::NodeDsl = parse_str(r#"Start => async { let x = 1; }"#).unwrap();
        let name = format_ident!("test");
        let paths = MacroPaths::new();
        let tokens = generate_phase_enum(&name, &dsl, &paths);
        let code = tokens.to_string();
        assert!(code.contains("TestPhase"));
        assert!(code.contains("Start"));
    }

    #[test]
    fn test_generate_phase_enum_with_captures() {
        let dsl: crate::node_dsl::NodeDsl = parse_str(
            r#"
                Init => async { let x = 1; },
                Process(data: String) => async { let _ = data; }
            "#,
        )
        .unwrap();
        let name = format_ident!("workflow");
        let paths = MacroPaths::new();
        let tokens = generate_phase_enum(&name, &dsl, &paths);
        let code = tokens.to_string();
        assert!(code.contains("WorkflowPhase"));
        assert!(code.contains("Init"));
        assert!(code.contains("Process"));
        assert!(code.contains("data : String"));
    }

    #[test]
    fn test_generate_full_node() {
        let dsl: crate::node_dsl::NodeDsl = parse_str(r#"Eval => async { let x = 1; }"#).unwrap();
        let name = format_ident!("judge");
        let state_ty: syn::Type = parse_str("JudgeState").unwrap();
        let update_ty: syn::Type = parse_str("JudgeStateUpdate").unwrap();

        let tokens = generate_node(&name, &state_ty, &update_ty, &dsl);
        let code = tokens.to_string();
        assert!(code.contains("JudgePhase"));
        assert!(code.contains("JudgeNode"));
        assert!(code.contains("NodeFn"));
        assert!(code.contains("\"judge\""));
    }

    #[test]
    fn test_generated_code_reads_phase_from_store() {
        let dsl: crate::node_dsl::NodeDsl = parse_str(r#"Init => async { let x = 1; }"#).unwrap();
        let name = format_ident!("mynode");
        let state_ty: syn::Type = parse_str("MyState").unwrap();
        let update_ty: syn::Type = parse_str("MyUpdate").unwrap();

        let tokens = generate_node(&name, &state_ty, &update_ty, &dsl);
        let code = tokens.to_string();
        // The generated code must read phase from ctx.phase_store
        assert!(
            code.contains("phase_store"),
            "Generated code must read from phase_store"
        );
        assert!(
            code.contains("MynodePhase"),
            "Generated code must reference the phase enum"
        );
    }

    #[test]
    fn test_generated_code_with_retry_attr() {
        let dsl: crate::node_dsl::NodeDsl = parse_str(
            r#"
                #[retry(max_attempts = 5)]
                Fetch => async { let x = 1; }
            "#,
        )
        .unwrap();
        let name = format_ident!("fetcher");
        let state_ty: syn::Type = parse_str("FetchState").unwrap();
        let update_ty: syn::Type = parse_str("FetchUpdate").unwrap();

        let tokens = generate_node(&name, &state_ty, &update_ty, &dsl);
        let code = tokens.to_string();
        assert!(
            code.contains("RetryPolicy"),
            "Generated code must include RetryPolicy for #[retry] phases"
        );
        assert!(
            code.contains("with_retry"),
            "Generated code must call with_retry for #[retry] phases"
        );
        assert!(
            code.contains("max_attempts : 5"),
            "Generated code must use the specified max_attempts"
        );
    }
}
