use super::{next_node_functional, ParseNode, ParseNodeFunctional};
use crate::ast_node::{ASTContext, Node, TokenSlice};
use crate::syntax_error::SyntaxError;
use crate::tokens::Token;
use crate::types::IJType;
use anyhow::Result;
use std::fmt::Debug;
use std::rc::Rc;

#[derive(Debug, PartialEq, Clone)]
pub struct AsFunction;

impl AsFunction {
    fn next_node_functional_part(
        slice: TokenSlice,
        context: &mut ASTContext,
        needed_outputs: Option<&[IJType]>,
    ) -> Result<(Vec<Rc<Node>>, TokenSlice)> {
        let (nodes, rest) = next_node_functional(slice, context, needed_outputs)?;

        if !rest.is_empty() && context.get_token_at_index(rest.start)? == &Token::TypeDeclaration {
            let rest = rest.move_start(1)?;
            let (fn_type, type_end) = IJType::parse_tokens(&context.get_tokens_from_slice(rest))?;
            let rest = rest.move_start(type_end)?;
            if let IJType::Function(_) = fn_type {
            } else {
                return Err(context.add_context_to_syntax_error(
                    SyntaxError::ExpectedFunction(fn_type.to_string()).into(),
                    slice,
                ));
            }

            let node = nodes
                .iter()
                .find(|n| n.output_type.type_match_function(&fn_type).unwrap_or(false))
                .ok_or_else(|| {
                    SyntaxError::TypeError(
                        fn_type.to_string(),
                        nodes
                            .iter()
                            .map(|n| n.output_type.to_string())
                            .collect::<Vec<_>>()
                            .join(", "),
                    )
                })?;

            return Ok((
                vec![Rc::new(Node::new_skip_number_inference(
                    node.op.clone(),
                    node.input_types.clone(),
                    fn_type,
                    node.operands.clone(),
                    context,
                )?)],
                rest,
            ));
        }

        Ok((nodes, rest))
    }
}

impl ParseNode for AsFunction {
    fn next_node(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
    ) -> Result<(Rc<Node>, TokenSlice)> {
        let (nodes, rest) = AsFunction::next_node_functional_part(slice, context, None)?;
        if nodes.len() != 1 {
            Err(context.add_context_to_syntax_error(
                SyntaxError::FunctionTypeAmbiguous(
                    context.token_slice_to_string(slice),
                    nodes
                        .iter()
                        .map(|n| n.output_type.to_string())
                        .collect::<Vec<_>>()
                        .join(", "),
                )
                .into(),
                slice,
            ))
        } else {
            Ok((nodes[0].clone(), rest))
        }
    }
}

impl ParseNodeFunctional for AsFunction {
    fn next_node_functional_impl(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
        needed_outputs: Option<&[IJType]>,
    ) -> Result<(Vec<Rc<Node>>, TokenSlice)> {
        let slice = slice.move_start(1)?;
        AsFunction::next_node_functional_part(slice, context, needed_outputs)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::function_enums::BinaryMathFunctionEnum;
    use crate::operations::Operation;
    use crate::parser::{parse_str, parse_str_no_context};
    use crate::types::FunctionSignature;

    #[test]
    fn test_as_function_simple() -> Result<()> {
        let mut context = ASTContext::new();
        parse_str("var f: Fn(N,N->N)", &mut context)?;

        let node = parse_str("~f", &mut context)?;
        assert_eq!(node.op, Operation::Function("f".to_string()));
        assert_eq!(node.output_type, IJType::number_function(2));
        Ok(())
    }

    #[test]
    fn test_as_function_with_type() -> Result<()> {
        let mut context = ASTContext::new();
        parse_str("var f: Fn(N<a>,N<b>->N<c>)", &mut context)?;

        let node = parse_str("~f", &mut context)?;
        assert_eq!(node.op, Operation::Function("f".to_string()));
        assert_eq!(
            node.output_type,
            IJType::Function(FunctionSignature::new(
                vec![
                    IJType::Number(Some("a".to_string())),
                    IJType::Number(Some("b".to_string()))
                ],
                IJType::Number(Some("c".to_string()))
            ))
        );
        Ok(())
    }

    #[test]
    fn test_as_function_declaration() -> Result<()> {
        let (node, _) = parse_str_no_context("~+:Fn(N,N->N)")?;

        assert_eq!(
            node.op,
            Operation::BinaryFunction(BinaryMathFunctionEnum::Add)
        );
        assert_eq!(node.output_type, IJType::number_function(2));
        Ok(())
    }

    #[test]
    fn test_as_function_ambiguous() -> Result<()> {
        assert!(parse_str_no_context("~+").is_err());

        Ok(())
    }

    #[test]
    fn test_as_function_function() -> Result<()> {
        let (node, _) = parse_str_no_context("~(/ ~+)")?;
        assert_eq!(node.op, Operation::Reduce);

        assert_eq!(
            node.output_type,
            IJType::Function(FunctionSignature::new(
                vec![IJType::Tensor(None)],
                IJType::Number(None)
            ))
        );
        Ok(())
    }

    #[test]
    fn test_as_function_double() -> Result<()> {
        let (node, _) = parse_str_no_context("~~+:Fn(N,N->N)")?;

        assert_eq!(
            node.op,
            Operation::BinaryFunction(BinaryMathFunctionEnum::Add)
        );
        assert_eq!(node.output_type, IJType::number_function(2));
        Ok(())
    }
}
