use oxc_ast::{
    AstKind,
    ast::{Argument, Expression},
};
use oxc_diagnostics::OxcDiagnostic;
use oxc_macros::declare_oxc_lint;
use oxc_span::Span;

use crate::{context::LintContext, rule::Rule, utils::PossibleJestNode};

fn add_type_parameter_to_module_mock_diagnostic(module_name: &str, span: Span) -> OxcDiagnostic {
    OxcDiagnostic::warn(
        "`jest.mock()` factories should not be used without an explicit type parameter.",
    )
    .with_help(format!(
        "Add a type parameter to the mock factory such as `typeof import({module_name:?})`"
    ))
    .with_label(span)
}

#[derive(Debug, Default, Clone)]
pub struct NoUntypedMockFactory;

declare_oxc_lint!(
    /// ### What it does
    ///
    /// This rule triggers a warning if `mock()` or `doMock()` is used without a generic
    /// type parameter or return type.
    ///
    /// ### Why is this bad?
    ///
    /// By default, `jest.mock` and `jest.doMock` allow any type to be returned by a
    /// mock factory. A generic type parameter can be used to enforce that the factory
    /// returns an object with the same shape as the original module, or some other
    /// strict type. Requiring a type makes it easier to use TypeScript to catch changes
    /// needed in test mocks when the source module changes.
    ///
    /// ### Examples
    ///
    /// Examples of **incorrect** code for this rule:
    /// ```typescript
    /// jest.mock('../moduleName', () => {
    ///     return jest.fn(() => 42);
    /// });
    ///
    /// jest.mock('./module', () => ({
    ///     ...jest.requireActual('./module'),
    ///     foo: jest.fn(),
    /// }));
    ///
    /// jest.mock('random-num', () => {
    ///     return jest.fn(() => 42);
    /// });
    /// ```
    ///
    /// Examples of **correct** code for this rule:
    /// ```typescript
    /// // Uses typeof import()
    /// jest.mock<typeof import('../moduleName')>('../moduleName', () => {
    ///     return jest.fn(() => 42);
    /// });
    ///
    /// jest.mock<typeof import('./module')>('./module', () => ({
    ///     ...jest.requireActual('./module'),
    ///     foo: jest.fn(),
    /// }));
    ///
    /// // Uses custom type
    /// jest.mock<() => number>('random-num', () => {
    ///     return jest.fn(() => 42);
    /// });
    ///
    /// // No factory
    /// jest.mock('random-num');
    ///
    /// // Virtual mock
    /// jest.mock(
    ///     '../moduleName',
    ///     () => {
    ///         return jest.fn(() => 42);
    ///     },
    ///     { virtual: true },
    /// );
    /// ```
    NoUntypedMockFactory,
    jest,
    style,
    conditional_fix,
    version = "0.2.15",
    short_description = "This rule triggers a warning if `mock()` or `doMock()` is used without a generic type parameter or return type.",
);

impl Rule for NoUntypedMockFactory {
    fn run_on_jest_node<'a, 'c>(
        &self,
        jest_node: &PossibleJestNode<'a, 'c>,
        ctx: &'c LintContext<'a>,
    ) {
        Self::run(jest_node, ctx);
    }

    fn should_run(&self, ctx: &crate::context::ContextHost) -> bool {
        ctx.source_type().is_typescript()
    }
}

impl NoUntypedMockFactory {
    fn run<'a>(possible_jest_node: &PossibleJestNode<'a, '_>, ctx: &LintContext<'a>) {
        let node = possible_jest_node.node;
        let AstKind::CallExpression(call_expr) = node.kind() else {
            return;
        };
        let Some(mem_expr) = call_expr.callee.as_member_expression() else {
            return;
        };
        let Some((property_span, property_name)) = mem_expr.static_property_info() else {
            return;
        };

        if call_expr.arguments.len() != 2 || (property_name != "mock" && property_name != "doMock")
        {
            return;
        }

        let Some(factory_node) = call_expr.arguments.get(1) else {
            return;
        };

        if call_expr.type_arguments.is_some() || Self::has_return_type(factory_node) {
            return;
        }

        let Some(name_node) = call_expr.arguments.first() else {
            return;
        };
        let Some(expr) = name_node.as_expression() else {
            return;
        };

        if let Expression::StringLiteral(string_literal) = expr {
            ctx.diagnostic_with_fix(
                add_type_parameter_to_module_mock_diagnostic(
                    string_literal.value.as_str(),
                    property_span,
                ),
                |fixer| {
                    let mut code = String::with_capacity(string_literal.value.len() + 20);
                    code.push_str("<typeof import('");
                    code.push_str(string_literal.value.as_str());
                    code.push_str("')>");

                    fixer.insert_text_after(&call_expr.callee, code)
                },
            );
        } else if let Expression::Identifier(ident) = expr {
            ctx.diagnostic(add_type_parameter_to_module_mock_diagnostic(
                ident.name.as_str(),
                property_span,
            ));
        }
    }

    fn has_return_type(argument: &Argument) -> bool {
        let Some(expr) = argument.as_expression() else {
            return false;
        };

        match expr {
            Expression::FunctionExpression(func_expr) => func_expr.return_type.is_some(),
            Expression::ArrowFunctionExpression(arrow_func_expr) => {
                arrow_func_expr.return_type.is_some()
            }
            _ => false,
        }
    }
}

#[test]
fn test() {
    use crate::tester::Tester;

    let pass = vec![
        ("jest.mock('random-number');", None),
        (
            "
                jest.mock<typeof import('../moduleName')>('../moduleName', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock<typeof import('./module')>('./module', () => ({
                    ...jest.requireActual('./module'),
                    foo: jest.fn()
                }));
            ",
            None,
        ),
        (
            "
                jest.mock<typeof import('foo')>('bar', () => ({
                    ...jest.requireActual('bar'),
                    foo: jest.fn()
                }));
            ",
            None,
        ),
        (
            "
                jest.doMock('./module', (): typeof import('./module') => ({
                    ...jest.requireActual('./module'),
                    foo: jest.fn()
                }));
            ",
            None,
        ),
        (
            "
                jest.mock('../moduleName', function (): typeof import('../moduleName') {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock('../moduleName', function (): (() => number) {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock<() => number>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest['doMock']<() => number>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock<any>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock(
                    '../moduleName',
                    () => {
                        return jest.fn(() => 42)
                    },
                    {virtual: true},
                );
            ",
            None,
        ),
        // Should not match
        (
            "
                mockito<() => number>('foo', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        ("test.skip('basic', async () => {});", None),
    ];

    let fail = vec![
        (
            "
                jest.mock('../moduleName', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock(\"./module\", () => ({
                    ...jest.requireActual('./module'),
                    foo: jest.fn()
                }));
            ",
            None,
        ),
        (
            "
                jest.mock('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.doMock('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest['mock']('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                const moduleToMock = 'random-num';
                jest.mock(moduleToMock, () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
    ];

    let fix = vec![
        (
            "
                jest.mock('../moduleName', () => {
                    return jest.fn(() => 42);
                });
            ",
            "
                jest.mock<typeof import('../moduleName')>('../moduleName', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.mock('./module', () => ({
                    ...jest.requireActual('./module'),
                    foo: jest.fn()
                }));
            ",
            "
                jest.mock<typeof import('./module')>('./module', () => ({
                    ...jest.requireActual('./module'),
                    foo: jest.fn()
                }));
            ",
            None,
        ),
        (
            "
                jest.mock('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            "
                jest.mock<typeof import('random-num')>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest.doMock('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            "
                jest.doMock<typeof import('random-num')>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        (
            "
                jest['mock']('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            "
                jest['mock']<typeof import('random-num')>('random-num', () => {
                    return jest.fn(() => 42);
                });
            ",
            None,
        ),
        // Test case for newline between jest.mock( and module name
        (
            "jest.mock(
  'foo',
  () => {}
);",
            "jest.mock<typeof import('foo')>(
  'foo',
  () => {}
);",
            None,
        ),
    ];

    Tester::new(NoUntypedMockFactory::NAME, NoUntypedMockFactory::PLUGIN, pass, fail)
        .with_jest_plugin(true)
        .expect_fix(fix)
        .test_and_snapshot();
}
