import { NodePath } from 'babel-traverse';
import {
    AssignmentExpression,
    CallExpression,
    callExpression,
    Identifier,
    ImportDeclaration,
    isArrowFunctionExpression,
    isIdentifier,
    isImportDefaultSpecifier,
    isImportSpecifier,
    isMemberExpression,
    stringLiteral,
    ExportDefaultDeclaration,
    isReturnStatement,
    Statement,
    VariableDeclaration,
    isJSX,
    Program,
    arrowFunctionExpression,
    identifier,
    blockStatement,
    variableDeclarator,
    variableDeclaration,
    memberExpression,
    expressionStatement,
    assignmentExpression,
    exportDefaultDeclaration,
    isProgram,
    returnStatement,
    ImportSpecifier,
    ifStatement,
    binaryExpression,
    nullLiteral, importDeclaration, importDefaultSpecifier
} from 'babel-types';
import nodePath from 'path';

function visitFile() {
    const cssModules: string[] = [];
    const cssFiles: string[] = [];
    const rootLevelReactFunctions: string[] = [];

    return {
        cssModules: cssModules,
        cssFiles: cssFiles,
        visitor: {
            AssignmentExpression(path: NodePath<AssignmentExpression>, state: any) {
                if (state.opts.side === 'client') {
                    if (isMemberExpression(path.node.left)) {
                        if (isIdentifier(path.node.left.property)) {
                            const ident = path.node.left.property as Identifier;
                            if (ident.name === 'getServerSideProps') {
                                path.remove();
                            }
                        }
                    }
                }
            },
            VariableDeclaration(path: NodePath<VariableDeclaration>, state: any) {
                if (state.opts.side === 'server') {
                    if (isProgram(path.parent) && path.node.declarations.length === 1) {
                        const decl = path.node.declarations[0];
                        if (!isIdentifier(decl.id) || !isArrowFunctionExpression(decl.init)) {
                            return;
                        }
                        const variableName = decl.id.name;
                        path.traverse({
                            Statement(path: NodePath<Statement>) {
                                if (isReturnStatement(path.node)) {
                                    // JSXFragment not yet added to this version of babel-types
                                    if (path.node.argument !== null && (isJSX(path.node.argument) || path.node.argument.type as any === 'JSXFragment')) {
                                        rootLevelReactFunctions.push(variableName);
                                    }
                                }
                            }
                        });
                    }
                }
            },
            CallExpression(path: NodePath<CallExpression>, state: any) {
                if (isMemberExpression(path.node.callee)) {
                    if (isIdentifier(path.node.callee.object) && isIdentifier(path.node.callee.property)) {
                        if (path.node.callee.object.name === 'OrisonUtil' && path.node.callee.property.name === 'invokeSided') {
                            const toUse = state.opts.side === 'server' ? path.node.arguments[0] : path.node.arguments[1];
                            if (!isArrowFunctionExpression(toUse)) {
                                throw new TypeError('Parameters to OrisonUtil.invokeSided must be arrow function expressions');
                            }
                            path.replaceWith(callExpression(toUse, []) as any);
                        }
                    }
                }
            },
            ImportDeclaration(path: NodePath<ImportDeclaration>, state: any) {
                if (state.opts.side === 'server') {
                    if (path.node.source.value.endsWith('.module.css')) {
                        if (isImportDefaultSpecifier(path.node.specifiers[0])) {
                            cssModules.push(path.node.specifiers[0].local.name);
                        } else {
                            throw new TypeError('Default import is required for CSS modules');
                        }
                    } else if (path.node.source.value.endsWith('.css')) {
                        if (path.node.specifiers.length === 0) {
                            let name = '__orison_css_global_' + cssFiles.length;
                            path.replaceWith(
                                importDeclaration(
                                    [
                                        importDefaultSpecifier(
                                            identifier(name)
                                        )
                                    ],
                                    path.node.source
                                )
                            );
                            cssFiles.push(name);
                        } else {
                            if (isImportDefaultSpecifier(path.node.specifiers[0])) {
                                if (path.node.specifiers[0].local.name.startsWith('__orison_css_global_')) {
                                    return;
                                }
                            }
                            throw new Error('You cannot import from a non-module CSS file into a variable; try "import \'' + path.node.source.value + '\';" instead');
                        }
                    }
                }
                if (path.node.source.value === 'orisonjs') {
                    for (let i = 0; i < path.node.specifiers.length; i++) {
                        const specifier = path.node.specifiers[i];
                        if (isImportSpecifier(specifier)) {
                            const removable = ['OrisonPage', 'OrisonUtil'];
                            if (removable.includes(specifier.imported.name)) {
                                if (path.node.specifiers.length === 1) {
                                    path.remove();
                                    break;
                                } else {
                                    path.node.specifiers.splice(i, 1);
                                    i = -1;
                                }
                            }
                        }
                    }
                    if (state.opts.side === 'client' && !path.removed) {
                        throw new TypeError(
                            'You cannot import anything from the "orisonjs" package on the client-side ' +
                            'that is not either a type definition or the OrisonUtil object, given: ' +
                            path.node.specifiers.filter(specifier => isImportSpecifier(specifier)).map(specifier => (specifier as ImportSpecifier).imported.name).join(', ')
                        );
                    }
                } else if (state.opts.nonRequire) {
                    if (path.node.source.value.startsWith('../') || path.node.source.value.startsWith('./')) {
                        const ext = nodePath.extname(path.node.source.value);
                        if (ext === '' || ext === 'js' || ext === 'ts' || ext === 'jsx' || ext === 'tsx') {
                            path.node.source = stringLiteral('NON_REQUIRE_' + path.node.source.value);
                        }
                    }
                }
            },
            ExportDefaultDeclaration(path: NodePath<ExportDefaultDeclaration>, state: any) {
                if (state.opts.side === 'server') {
                    if ((cssModules.length > 0 || cssFiles.length > 0) && rootLevelReactFunctions.length > 0 && isIdentifier(path.node.declaration)) {
                        for (const func of rootLevelReactFunctions) {
                            if (func === path.node.declaration.name) {
                                const blocks: Statement[] = [
                                    variableDeclaration('var', [
                                        variableDeclarator(
                                            identifier('__orison_style_context_import'),
                                            memberExpression(
                                                callExpression(identifier('require'), [stringLiteral('orisonjs/lib/styles')]),
                                                identifier('default')
                                            )
                                        )
                                    ]),
                                    variableDeclaration('var', [
                                        variableDeclarator(
                                            identifier('__orison_style_context'),
                                            callExpression(
                                                memberExpression(
                                                    identifier('React'),
                                                    identifier('useContext')
                                                ),
                                                [identifier('__orison_style_context_import')]
                                            )
                                        )
                                    ]),
                                ];
                                blocks.push(
                                    ifStatement(
                                        binaryExpression(
                                            '===',
                                            memberExpression(
                                                identifier('__orison_style_context'),
                                                identifier('cssSources')
                                            ),
                                            nullLiteral()
                                        ),
                                        blockStatement([
                                            returnStatement(
                                                callExpression(
                                                    path.node.declaration,
                                                    [identifier('__props')]
                                                )
                                            )
                                        ])
                                    )
                                );
                                const allCss = cssModules.concat(cssFiles);
                                for (const css of allCss) {
                                    blocks.push(
                                        expressionStatement(
                                            callExpression(
                                                memberExpression(
                                                    memberExpression(
                                                        identifier('__orison_style_context'),
                                                        identifier('cssSources')
                                                    ),
                                                    identifier('add')
                                                ),
                                                [
                                                    memberExpression(
                                                        identifier(css),
                                                        identifier('__orison_css_sources')
                                                    )
                                                ]
                                            )
                                        )
                                    );
                                }
                                blocks.push(
                                    returnStatement(
                                        callExpression(
                                            path.node.declaration,
                                            [identifier('__props')]
                                        )
                                    )
                                );

                                const wrapperFunction = blockStatement(blocks);
                                const wrapperFunctionHeader = arrowFunctionExpression([identifier('__props')], wrapperFunction);

                                if (!isProgram(path.parent)) {
                                    throw new Error('Expecting export default declaration in program root');
                                }

                                path.parent.body.push(
                                    variableDeclaration('var', [
                                        variableDeclarator(
                                            identifier('__wrapped_component'),
                                            wrapperFunctionHeader
                                        )
                                    ]),
                                    expressionStatement(
                                        assignmentExpression(
                                            '=',
                                            memberExpression(
                                                identifier('__wrapped_component'),
                                                identifier('getServerSideProps')
                                            ),
                                            memberExpression(
                                                path.node.declaration,
                                                identifier('getServerSideProps')
                                            )
                                        )
                                    ),
                                    exportDefaultDeclaration(
                                        identifier('__wrapped_component')
                                    )
                                );

                                path.parentPath.traverse({
                                    VariableDeclaration(varPath: NodePath<VariableDeclaration>) {
                                        if (varPath.node.declarations.length === 1 && isIdentifier(varPath.node.declarations[0].id)) {
                                            if (varPath.node.declarations[0].id.name === '__wrapped_component') {
                                                path.parentPath.scope.registerDeclaration(varPath);
                                            }
                                        }
                                    }
                                });
                                path.remove();

                                {
                                    let len = cssModules.length;
                                    for (let i = 0; i < len; i++) {
                                        cssModules.pop();
                                    }
                                }
                                {
                                    let len = cssFiles.length;
                                    for (let i = 0; i < len; i++) {
                                        cssFiles.pop();
                                    }
                                }
                                return;
                            }
                        }
                    }
                }
            }
        }
    };
}

export default function (api: any) {
    api.assertVersion(7);

    return {
        visitor: {
            Program(path: NodePath<Program>, state: any) {
                const { cssModules, cssFiles, visitor } = visitFile();
                path.traverse(visitor, state);

                if (cssModules.length > 0 || cssFiles.length > 0) {
                    throw new Error('Could not inject server-side CSS bindings. Make sure your default export is a React functional component.');
                }
            },
        }
    }
}