import {
    Plugin,
    Shape,
    type Source,
    ts,
    IdentifierShape,
    type StringShape,
    ObjectShape,
} from '@servicenow/sdk-build-core'

export class ArrowFunctionShape extends Shape {
    private readonly parameters: IdentifierShape[]
    private readonly statements: Shape[]
    private readonly returnValue: Shape | undefined
    private readonly implicitReturn: boolean

    constructor({
        source,
        parameters,
        statements,
        returnValue,
        implicitReturn,
    }: {
        source: Source
        parameters?: (IdentifierShape | StringShape | string)[]
        statements?: Shape[]
        returnValue?: Shape
        implicitReturn?: boolean
    }) {
        super({ source })
        this.parameters = (parameters ?? []).map((s) =>
            s instanceof IdentifierShape
                ? s
                : new IdentifierShape({
                      source,
                      name: typeof s === 'string' ? s : s.getValue(),
                  })
        )
        this.statements = statements ?? []
        this.returnValue = returnValue
        this.implicitReturn = implicitReturn ?? false
    }

    getParameters(): IdentifierShape[] {
        return this.parameters
    }

    getParameter(index: number): IdentifierShape {
        const parameter = this.parameters[index]
        if (!parameter) {
            throw new Error(`Arrow function does not have a parameter at index: ${index}`)
        }

        return parameter
    }

    getStatements(): Shape[] {
        return this.statements
    }

    getStatement(index: number): Shape {
        const statement = this.statements[index]
        if (!statement) {
            throw new Error(`Arrow function does not have a statement at index: ${index}`)
        }

        return statement
    }

    getReturnValue(): Shape | undefined {
        return this.returnValue
    }

    isImplicitReturn(): boolean {
        return this.implicitReturn
    }

    override getCode(): string {
        const params = `(${this.getParameters()
            .map((p) => p.getCode())
            .join(', ')})`

        // Use concise arrow function syntax when there are no statements and only a return value
        if (this.implicitReturn && this.returnValue) {
            if (this.returnValue instanceof ObjectShape) {
                return `${params} => (${this.returnValue.getCode()})`
            }
            return `${params} => ${this.returnValue.getCode()}`
        }

        // Block body syntax
        const parts = this.getStatements().map((s) => s.getCode())
        if (this.returnValue) {
            parts.push(`return ${this.returnValue.getCode()}`)
        }
        return `${params} => {\n    ${parts.join('\n    ')}\n}`
    }
}

export const ArrowFunctionPlugin = Plugin.create({
    name: 'ArrowFunctionPlugin',
    noTelemetry: true,
    nodes: [
        /**
         * Catch-all handler for ReturnStatement nodes encountered outside of arrow functions.
         * Return statements within arrow functions are handled inline by the ArrowFunction
         * node handler below, which extracts the return expression from the last statement.
         * Any ReturnStatement that reaches this handler is outside an arrow function context
         * and is not supported.
         */
        {
            node: 'ReturnStatement',
            fileTypes: ['fluent'],
            toShape(node, { diagnostics }) {
                diagnostics.error(node, 'Return statements are only supported inside arrow functions')
                return { success: false }
            },
        },
        {
            node: 'ArrowFunction',
            fileTypes: ['fluent'],
            async toShape(node, { transform, diagnostics }) {
                const parameters: IdentifierShape[] = []
                for (const param of node.getParameters()) {
                    const result = await transform.toShape(param)
                    if (result.success && result.value.isIdentifier()) {
                        parameters.push(result.value)
                    } else {
                        diagnostics.error(param, 'Unsupported parameter in arrow function')
                        return { success: false }
                    }
                }

                // Check if this is an expression-body arrow function (concise form) - () => 'value'
                const body = node.getBody()
                const isExpressionBody = body && !ts.Node.isBlock(body)

                if (isExpressionBody) {
                    // Expression body: () => expr - treat the body as the return value
                    const result = await transform.toShape(body)
                    if (!result.success) {
                        diagnostics.error(body, 'Unsupported expression in arrow function body')
                        return { success: false }
                    }
                    return {
                        success: true,
                        value: new ArrowFunctionShape({
                            source: node,
                            parameters,
                            statements: [],
                            returnValue: result.value,
                            implicitReturn: true,
                        }),
                    }
                }

                const nodeStatements = node.getStatements()
                const lastStatement = nodeStatements[nodeStatements.length - 1]
                const hasReturn = lastStatement && ts.Node.isReturnStatement(lastStatement)
                const bodyStatements = hasReturn ? nodeStatements.slice(0, -1) : nodeStatements

                const statements: Shape[] = []
                for (const statement of bodyStatements) {
                    const result = await transform.toShape(statement)
                    if (!result.success) {
                        diagnostics.error(statement, 'Unsupported statement in arrow function body')
                        return { success: false }
                    }
                    statements.push(result.value)
                }

                let returnValue: Shape | undefined
                if (hasReturn) {
                    const expr = lastStatement.getExpression()
                    if (expr) {
                        const result = await transform.toShape(expr)
                        if (!result.success) {
                            diagnostics.error(lastStatement, 'Unsupported expression in return statement')
                            return { success: false }
                        }
                        returnValue = result.value
                    }
                }

                const args = { source: node, parameters, statements }

                return {
                    success: true,
                    value: returnValue
                        ? new ArrowFunctionShape({ ...args, returnValue })
                        : new ArrowFunctionShape(args),
                }
            },
        },
    ],
    shapes: [
        {
            shape: ArrowFunctionShape,
            async commit(shape, target, { commit }) {
                const originalTargetWasArrow = ts.Node.isArrowFunction(target)
                const arrowTarget = originalTargetWasArrow
                    ? target
                    : target.replaceWithText('() => {}').asKindOrThrow(ts.SyntaxKind.ArrowFunction)

                const shapeStatements = shape.getStatements()
                const returnValue = shape.getReturnValue()
                const body = arrowTarget.getBody()

                // Preserve concise arrows when the shape is still an implicit return.
                if (
                    originalTargetWasArrow &&
                    !ts.Node.isBlock(body) &&
                    shapeStatements.length === 0 &&
                    returnValue &&
                    shape.isImplicitReturn()
                ) {
                    if (returnValue.is(ObjectShape)) {
                        body.replaceWithText(`(${returnValue.getCode()})`)
                    } else {
                        await commit(returnValue, body)
                    }
                    return { success: true }
                }

                if (!ts.Node.isBlock(body)) {
                    // ts-morph cannot grow expression-body arrows like `() => ({})` into
                    // block bodies with statement APIs, so convert the body to an empty block
                    // first and then continue through the normal commit flow below.
                    body.replaceWithText('{}')
                }

                const getNonReturnStatements = () =>
                    arrowTarget.getStatements().filter((statement) => !ts.Node.isReturnStatement(statement))

                const existingStatements = getNonReturnStatements()
                const excess = existingStatements.length - shapeStatements.length
                for (let i = 1; i <= excess; i++) {
                    existingStatements[existingStatements.length - i]?.remove()
                }

                for (const [i, statement] of shapeStatements.entries()) {
                    const currentStatements = getNonReturnStatements()
                    const targetStatement = currentStatements[i]
                    if (targetStatement) {
                        await commit(statement, targetStatement)
                    } else {
                        // Insert before `return` so return-only bodies can grow safely.
                        const returnIndex = arrowTarget
                            .getStatements()
                            .findIndex((node) => ts.Node.isReturnStatement(node))
                        const inserted =
                            returnIndex >= 0
                                ? arrowTarget.insertStatements(returnIndex, statement.getCode())[0]!
                                : arrowTarget.addStatements(statement.getCode())[0]!
                        await commit(statement, inserted)
                    }
                }

                let existingReturn = arrowTarget
                    .getStatements()
                    .find((statement) => ts.Node.isReturnStatement(statement))
                if (!returnValue) {
                    existingReturn?.remove()
                    return { success: true }
                }

                if (!existingReturn) {
                    const inserted = arrowTarget.addStatements(`return ${returnValue.getCode()}`)[0]!
                    existingReturn = ts.Node.isReturnStatement(inserted)
                        ? inserted
                        : arrowTarget.getStatements().find((statement) => ts.Node.isReturnStatement(statement))
                }

                if (existingReturn) {
                    const expr = existingReturn.getExpression()
                    if (expr) {
                        await commit(returnValue, expr)
                    } else {
                        existingReturn.replaceWithText(`return ${returnValue.getCode()}`)
                    }
                }
                return { success: true }
            },
        },
    ],
})