use swc_common::{Spanned, DUMMY_SP};
use swc_ecma_ast::{
    ArrowExpr, BindingIdent, BlockStmtOrExpr, Class, Expr, Function, Ident, Lit, ReturnStmt, Stmt,
    TsKeywordTypeKind, TsParenthesizedType, TsType, TsTypeAliasDecl, TsTypeAnn,
    TsUnionOrIntersectionType, TsUnionType, UnaryExpr, UnaryOp,
};
use swc_ecma_visit::{Visit, VisitWith};

use super::{
    util::types::{ts_keyword_type, type_ann},
    FastDts,
};

impl FastDts {
    pub(crate) fn infer_type_from_expr(&mut self, e: &Expr) -> Option<Box<TsType>> {
        match e {
            Expr::Ident(ident) if ident.sym.as_str() == "undefined" => {
                Some(ts_keyword_type(TsKeywordTypeKind::TsUndefinedKeyword))
            }
            Expr::Lit(lit) => match lit {
                Lit::Str(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsStringKeyword)),
                Lit::Bool(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsBooleanKeyword)),
                Lit::Num(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsNumberKeyword)),
                Lit::BigInt(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsBigIntKeyword)),
                Lit::Null(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsNullKeyword)),
                Lit::Regex(_) | Lit::JSXText(_) => None,
            },
            Expr::Tpl(_) => Some(ts_keyword_type(TsKeywordTypeKind::TsStringKeyword)),
            Expr::Fn(fn_expr) => self.transform_fn_to_ts_type(
                &fn_expr.function,
                fn_expr.ident.as_ref().map(|ident| ident.span),
            ),
            Expr::Arrow(arrow_expr) => self.transform_arrow_expr_to_ts_type(arrow_expr),
            Expr::Array(arr) => {
                self.array_inferred(arr.span);
                Some(ts_keyword_type(TsKeywordTypeKind::TsUnknownKeyword))
            }
            Expr::Object(obj) => self.transform_object_to_ts_type(obj, false),
            Expr::Class(class) => {
                self.inferred_type_of_class_expression(class.span());
                Some(ts_keyword_type(TsKeywordTypeKind::TsUnknownKeyword))
            }
            Expr::Paren(expr) => self.infer_type_from_expr(&expr.expr),
            Expr::TsNonNull(non_null) => self.infer_type_from_expr(&non_null.expr),
            Expr::TsSatisfies(satisifies) => self.infer_type_from_expr(&satisifies.expr),
            Expr::TsConstAssertion(assertion) => self.transform_expr_to_ts_type(&assertion.expr),
            Expr::TsAs(ts_as) => Some(ts_as.type_ann.clone()),
            Expr::TsTypeAssertion(type_assertion) => Some(type_assertion.type_ann.clone()),
            Expr::Unary(unary) if Self::can_infer_unary_expr(unary) => {
                self.infer_type_from_expr(&unary.arg)
            }
            _ => None,
        }
    }

    pub(crate) fn infer_function_return_type(
        &mut self,
        function: &Function,
    ) -> Option<Box<TsTypeAnn>> {
        if function.return_type.is_some() {
            return function.return_type.clone();
        }

        if function.is_async || function.is_generator {
            return None;
        }

        function
            .body
            .as_ref()
            .and_then(|body| ReturnTypeInferrer::infer(self, &body.stmts))
            .map(type_ann)
    }

    pub(crate) fn infer_arrow_return_type(&mut self, arrow: &ArrowExpr) -> Option<Box<TsTypeAnn>> {
        if arrow.return_type.is_some() {
            return arrow.return_type.clone();
        }

        if arrow.is_async || arrow.is_generator {
            return None;
        }

        match arrow.body.as_ref() {
            BlockStmtOrExpr::BlockStmt(block_stmt) => {
                ReturnTypeInferrer::infer(self, &block_stmt.stmts)
            }
            BlockStmtOrExpr::Expr(expr) => self.infer_type_from_expr(expr),
        }
        .map(type_ann)
    }

    pub(crate) fn need_to_infer_type_from_expression(expr: &Expr) -> bool {
        match expr {
            Expr::Lit(lit) => !(lit.is_str() || lit.is_num() || lit.is_big_int() || lit.is_bool()),
            Expr::Tpl(tpl) => !tpl.exprs.is_empty(),
            Expr::Unary(unary) => !Self::can_infer_unary_expr(unary),
            _ => true,
        }
    }

    pub(crate) fn can_infer_unary_expr(unary: &UnaryExpr) -> bool {
        let is_arithmetic = matches!(unary.op, UnaryOp::Plus | UnaryOp::Minus);
        let is_number_lit = match unary.arg.as_ref() {
            Expr::Lit(lit) => lit.is_num() || lit.is_big_int(),
            _ => false,
        };
        is_arithmetic && is_number_lit
    }
}

#[derive(Default)]
pub struct ReturnTypeInferrer {
    value_bindings: Vec<Ident>,
    type_bindings: Vec<Ident>,
    return_expr: Option<Option<Box<Expr>>>,
    return_expr_count: u8,
}

impl ReturnTypeInferrer {
    pub fn infer(transformer: &mut FastDts, stmts: &[Stmt]) -> Option<Box<TsType>> {
        let mut visitor = ReturnTypeInferrer::default();
        stmts.visit_children_with(&mut visitor);

        let expr = visitor.return_expr??;
        let Some(mut expr_type) = transformer.infer_type_from_expr(&expr) else {
            return if expr.is_fn_expr() || expr.is_arrow() {
                Some(ts_keyword_type(TsKeywordTypeKind::TsUnknownKeyword))
            } else {
                None
            };
        };

        if let Some((ref_name, is_value)) = match expr_type.as_ref() {
            TsType::TsTypeRef(type_ref) => Some((type_ref.type_name.as_ident()?, false)),
            TsType::TsTypeQuery(type_query) => {
                Some((type_query.expr_name.as_ts_entity_name()?.as_ident()?, true))
            }
            _ => None,
        } {
            let is_defined = if is_value {
                visitor.value_bindings.contains(ref_name)
            } else {
                visitor.type_bindings.contains(ref_name)
            };

            if is_defined {
                transformer.type_containing_private_name(ref_name.sym.as_str(), ref_name.span);
            }
        }

        if visitor.return_expr_count > 1 {
            if expr_type.is_ts_fn_or_constructor_type() {
                expr_type = Box::new(TsType::TsParenthesizedType(TsParenthesizedType {
                    span: DUMMY_SP,
                    type_ann: expr_type,
                }));
            }

            expr_type = Box::new(TsType::TsUnionOrIntersectionType(
                TsUnionOrIntersectionType::TsUnionType(TsUnionType {
                    span: DUMMY_SP,
                    types: vec![
                        expr_type,
                        ts_keyword_type(TsKeywordTypeKind::TsUndefinedKeyword),
                    ],
                }),
            ))
        }

        Some(expr_type)
    }
}

impl Visit for ReturnTypeInferrer {
    fn visit_arrow_expr(&mut self, _node: &ArrowExpr) {}

    fn visit_function(&mut self, _node: &Function) {}

    fn visit_class(&mut self, _node: &Class) {}

    fn visit_binding_ident(&mut self, node: &BindingIdent) {
        self.value_bindings.push(node.id.clone());
    }

    fn visit_ts_type_alias_decl(&mut self, node: &TsTypeAliasDecl) {
        self.type_bindings.push(node.id.clone());
    }

    fn visit_return_stmt(&mut self, node: &ReturnStmt) {
        self.return_expr_count += 1;
        if self.return_expr_count > 1 {
            if let Some(return_expr) = &self.return_expr {
                if return_expr.is_some() {
                    self.return_expr = None;
                    return;
                }
            } else {
                return;
            }
        }
        self.return_expr = Some(node.arg.clone());
    }
}
