reefer 0.3.0

use crate::FrameType;
use crate::basis_grammar::Alias;
use crate::clifford::SynBasis;
use crate::optimizer::Optimizer;
use crate::spec::AlgebraSpec;
use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use quote::quote;
use std::collections::BTreeMap;
use std::mem::take;
use syn::parse::{Parse, ParseStream};

/// Expand the `algebra` attribute macro into a fully generated module definition.
///
/// The function parses the attribute payload, constructs an internal algebra specification,
/// and then rewrites the provided module by expanding `basis!` and `shape!` helper macros.
pub fn do_algebra_macro(
    attr_ts: TokenStream,
    mod_ts: TokenStream,
) -> Result<TokenStream, syn::Error> {
    let spec: AlgebraSpec = syn::parse2(quote!(#attr_ts #mod_ts))?;

    let mod_ = spec.module;
    // println!("{}", prettyplease::unparse(&syn::parse_quote!(#mod_)));
    Ok(quote!(#mod_))
}

/// Entry point for the `build_expr` procedural macro.
///
/// The function parses an [`AlgebraSpec`] alongside an expression closure, performs light-weight
/// optimisations, and returns the resulting token stream.
pub fn do_build_expr_macro(ts: TokenStream) -> Result<TokenStream, syn::Error> {
    let build_expr: BuildExpr = syn::parse2(ts)?;
    let optimized_expr = build_expr.optimize()?;
    let fun = optimized_expr.fun;

    // set span to call_site for now, not 100% on why atm, probably messed up span hygiene somewhere
    let call_site = Span::call_site();
    // Inject operator trait imports into the closure body so generated postfix
    // method calls (e.g. `a.add(b)`) compile even when consumers don't
    // explicitly `use std::ops` themselves. We do this by replacing the
    // closure body with a small block that contains the `use` items followed
    // by the original expression.
    let mut fun = fun.clone();
    let original_body = *fun.body.clone();
    fun.body = Box::new(syn::parse_quote!({
        use core::ops::{Add, Mul, Div, Sub, Neg};
        #original_body
    }));
    let mut ts = fun.into_token_stream();
    ts = ts
        .into_iter()
        .map(|mut token| {
            token.set_span(token.span().resolved_at(call_site));
            token
        })
        .collect();
    Ok(ts)
}

pub type BasisElem = SynBasis<FrameType>;

/// Parsed representation of the `build_expr!` macro input.
pub struct BuildExpr {
    /// Algebra specification parsed from the macro preamble.
    pub spec: AlgebraSpec,
    /// Closure that will be optimised into a struct literal or expression tree.
    pub fun: syn::ExprClosure,
    /// Mapping from alias identifiers to their multivector representations.
    pub alias_mapping: BTreeMap<Alias, BasisElem>,
}
impl Parse for BuildExpr {
    /// Parse the macro input into an [`AlgebraSpec`] followed by a closure expression.
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mut build = BuildExpr {
            spec: input.parse()?,
            fun: input.parse()?,
            alias_mapping: BTreeMap::new(),
        };
        build.alias_mapping = build.spec.build_alias_mapping()?;
        Ok(build)
    }
}

impl BuildExpr {
    /// Apply compile-time optimisations to the captured closure body.
    pub fn optimize(mut self) -> syn::Result<Self> {
        // Split optimisation into two passes to prepare for future cross-term
        // rewrites (CSE etc.). Pass1 mirrors the existing lightweight
        // optimisation (evaluate closure, drop symbolic zeros, build shape
        // literal). Pass2 is currently a no-op stub and will be implemented
        // later.
        let alias_map = take(&mut self.alias_mapping);
        let mut build = Optimizer::new(self.spec, alias_map)?;
        // Pass 1: current behaviour (moved to optimizer_passes::run_pass1)
        build = build.run_pass1(&mut self.fun)?;
        // Pass 2: stub (future CSE / egg-driven cross-term rewrites)
        build = build.run_pass2(&mut self.fun)?;
        self.spec = build.spec;
        Ok(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use proc_macro2::TokenStream;
    use quote::{ToTokens, quote};
    use syn::{Expr, ExprClosure};

    /// Valid attribute and module inputs should expand without returning an error.
    #[test]
    fn test_do_algebra_macro_valid() {
        let attr_ts: TokenStream = quote! { f32, 2, 1 };
        let mod_ts: TokenStream = quote! {
            mod testmod {
                basis!(e1 = P0 + N0);
                basis!(e2 = P1);
                shape!(Bar { e1, e2 });
            }
        };
        let result = do_algebra_macro(attr_ts, mod_ts);
        result.unwrap();
    }

    /// Invalid attribute inputs should surface parsing errors.
    #[test]
    fn test_do_algebra_macro_invalid_attr() {
        let attr_ts: TokenStream = quote! { invalid, attr };
        let mod_ts: TokenStream = quote! {
            mod testmod {
                basis! { e1 = P0 + P1 }
            }
        };
        let result = do_algebra_macro(attr_ts, mod_ts);
        assert!(result.is_err());
    }

    /// Invalid module syntax must result in an error.
    #[test]
    fn test_do_algebra_macro_invalid_mod() {
        let attr_ts: TokenStream = quote! { f32, 2, 1 };
        let mod_ts: TokenStream = quote! { invalid mod syntax };
        let result = do_algebra_macro(attr_ts, mod_ts);
        assert!(result.is_err());
    }
    fn build_fixture(spec_tokens: TokenStream, closure_src: &str) -> BuildExpr {
        let spec: AlgebraSpec = syn::parse2(spec_tokens).expect("spec");
        let fun: ExprClosure = syn::parse_str(closure_src).expect("closure");
        let alias_mapping = spec.build_alias_mapping().expect("alias map");
        BuildExpr {
            spec,
            fun,
            alias_mapping,
        }
    }

    #[test]
    fn constant_closure_rewrites_to_shape_literal() {
        let spec_tokens = quote! {
            f32, 3, 0
            mod fixture {
                basis!(e1 = P0);
                basis!(e2 = P1);
                basis!(e3 = P2);

                shape!(Vector { e1, e2, e3 });
            }
        };

        let build_expr = build_fixture(spec_tokens, "|| e1 + 2.0 * e2 - 3.0 * e3");
        let optimized = build_expr.optimize().expect("optimize");

        assert!(optimized.fun.inputs.is_empty(), "closure inputs mutated");

        match optimized.fun.body.as_ref() {
            Expr::Struct(struct_expr) => {
                let path = struct_expr.path.to_token_stream().to_string();
                assert!(
                    path.ends_with("fixture :: Vector"),
                    "unexpected path: {path}"
                );
                assert_eq!(struct_expr.fields.len(), 3, "expected three vector fields");
            }
            other => panic!("expected struct literal, got {:?}", other),
        }
    }

    #[test]
    fn unknown_alias_passes_through_as_scalar() {
        let spec_tokens = quote! {
            f32, 3, 0
            mod fixture {
                basis!(e1 = P0);
                basis!(e2 = P1);
                basis!(e3 = P2);

                shape!(Vector { e1, e2, e3 });
                shape!(Scalar { 1 });
            }
        };

        let build_expr = build_fixture(spec_tokens, "|| unknown_alias");
        let optimized = build_expr.optimize().expect("optimize");
        assert!(optimized.fun.inputs.is_empty(), "closure inputs mutated");
        match optimized.fun.body.as_ref() {
            Expr::Struct(struct_expr) => {
                let path = struct_expr.path.to_token_stream().to_string();
                assert!(
                    path.ends_with("fixture :: Scalar"),
                    "unexpected path: {path}"
                );
                assert_eq!(struct_expr.fields.len(), 1, "expected one scalar field");
            }
            other => panic!("expected struct literal, got {:?}", other),
        }
    }

    #[test]
    fn closure_args_properly_bound() {
        let spec_tokens = quote! {
            f32, 1, 1
            mod fixture {
                basis!(t = P0);
                basis!(e1 = N0);

                shape!(Event { t, e1 });
            }
        };

        let build_expr = build_fixture(spec_tokens, "|ev: fixture::Event| ev");
        let optimized = build_expr.optimize().expect("optimize");
        assert_eq!(optimized.fun.inputs.len(), 1, "closure inputs mutated");
        match optimized.fun.body.as_ref() {
            Expr::Struct(struct_expr) => {
                let path = struct_expr.path.to_token_stream().to_string();
                assert!(
                    path.ends_with("fixture :: Event"),
                    "unexpected path: {path}"
                );
                assert_eq!(struct_expr.fields.len(), 2, "expected two event fields");
                let field_names: Vec<_> = struct_expr
                    .fields
                    .iter()
                    .map(|f| match &f.member {
                        syn::Member::Named(ident) => ident.to_string(),
                        syn::Member::Unnamed(_) => panic!("expected named field"),
                    })
                    .collect();
                assert!(field_names.contains(&"t".to_string()), "missing field: t");
                assert!(field_names.contains(&"e1".to_string()), "missing field: e1");
            }
            other => panic!("expected struct literal, got {:?}", other),
        }
    }

    #[test]
    fn closure_with_unknown_alias_and_arg() {
        let spec_tokens = quote! {
            f32, 1, 1
            mod fixture {
                basis!(t = P0);
                basis!(e1 = N0);

                shape!(Event { t, e1 });
                shape!(Scalar { 1 });
            }
        };

        let build_expr = build_fixture(spec_tokens, "|ev: fixture::Event| unknown_alias * t + ev");
        let optimized = build_expr.optimize().expect("optimize");
        assert_eq!(optimized.fun.inputs.len(), 1, "closure inputs mutated");

        let expected = "fixture :: Event { t : (unknown_alias) . add (ev . t) , e1 : ev . e1 }";
        let actual = optimized.fun.body.to_token_stream().to_string();
        assert_eq!(actual, expected, "optimized closure body mismatch");
    }
}