/**
 * Function return-type validator for Acurity Architect.
 *
 * ACU-FNC-010  functionReturnTypeMismatch
 *   Fires when a RETURN expression's inferred type does not match the function's
 *   declared return type. Only fires when the expression type can be determined
 *   from a literal or simple variable lookup — complex expressions are skipped
 *   to avoid false positives.
 *
 * ACU-FNC-011  functionMissingReturn
 *   Fires when a typed function contains no RETURN statement.
 *   Emitted on the ENDFUNCTION line.
 */

import { createDiagnostic } from "../diagnostics/diagnosticFactory.js";
import { buildDocumentAst } from "../ast/documentAst.js";
import { parseExpression } from "../ast/expressionParser.js";
import { NodeType } from "../ast/nodes.js";
import { ExprNodeType } from "../ast/expressionNodes.js";
import { walkStatements } from "../ast/astHelpers.js";

// ---------------------------------------------------------------------------
// AST helpers
// ---------------------------------------------------------------------------

/**
 * Build a map of function-declaration start-line → body statement array.
 * Recurses into nested block bodies so guards like `IF NOT SYMBOL(...)` are found.
 * Each FunctionDeclaration registers its own body; nested functions do the same
 * recursively via the visitor returning false only after registering.
 */
function collectFunctionBodies(statements, out) {
    walkStatements(statements, stmt => {
        if (stmt.type === NodeType.FunctionDeclaration) {
            out.set(stmt.line, stmt.body || []);
            // Register nested functions by recursing into the body explicitly.
            collectFunctionBodies(stmt.body || [], out);
            return false; // walkStatements would re-descend; we handled it above
        }
    });
}

/**
 * Collect all ReturnStatement nodes from a function body.
 * Does NOT descend into nested FunctionDeclaration or MacroDeclaration bodies
 * (their returns belong to a different scope).
 */
function collectReturnStatements(body) {
    const out = [];
    walkStatements(body, stmt => {
        if (stmt.type === NodeType.ReturnStatement) { out.push(stmt); }
        if (stmt.type === NodeType.FunctionDeclaration ||
            stmt.type === NodeType.MacroDeclaration) { return false; }
    });
    return out;
}

// ---------------------------------------------------------------------------
// Type inference
// ---------------------------------------------------------------------------

/**
 * Infers the return type from a ReturnStatement AST node using the expression parser.
 * Returns the uppercase type string, or null when the type cannot be determined
 * (complex expressions, function calls, etc.) — callers treat null as "unknown, skip".
 *
 * @param {object} returnStmt - ReturnStatement AST node
 * @param {Map<string,string>} params - Function parameter name → type map
 * @param {Map<string,string>} locals - Function local variable name → type map
 * @returns {string|null}
 */
function inferReturnExpressionType(returnStmt, params, locals) {
    if (!returnStmt.value) return null; // bare RETURN

    const expr = parseExpression(returnStmt.value);
    if (!expr) return null;

    // Literal — type already inferred by the expression parser
    if (expr.type === ExprNodeType.LITERAL) {
        return expr.inferredType ? expr.inferredType.toUpperCase() : null;
    }

    // Date / Time literals
    if (expr.type === ExprNodeType.DATE) return "DATE";
    if (expr.type === ExprNodeType.TIME) return "TIME";

    // Unary minus wrapping a literal: -42, -3.14
    if (expr.type === ExprNodeType.UNARY && expr.op === "-" &&
        expr.operand?.type === ExprNodeType.LITERAL) {
        return expr.operand.inferredType ? expr.operand.inferredType.toUpperCase() : null;
    }

    // Simple variable reference — look up in params then locals (case-insensitive)
    if (expr.type === ExprNodeType.IDENTIFIER) {
        const upper = expr.name.toUpperCase();
        for (const [k, v] of params) {
            if (k.toUpperCase() === upper) return v.toUpperCase();
        }
        for (const [k, v] of locals) {
            if (k.toUpperCase() === upper) return v.toUpperCase();
        }
    }

    return null;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/**
 * Validates return types for all user-defined functions in the document.
 * @param {object} document - VS Code TextDocument or PlainTextDocument
 * @param {Function} getRuleSeverity
 * @param {object} [options={}]
 * @param {Array} [options.functionScopes] - Pre-computed function scopes (from documentFacts)
 * @returns {object[]} Array of diagnostics
 */
function validateReturnTypes(document, getRuleSeverity, options = {}) {
    const { functionScopes = [] } = options;
    const diagnostics = [];

    const { ast } = buildDocumentAst(document);
    const funcBodyByLine = new Map();
    collectFunctionBodies(ast.statements, funcBodyByLine);

    for (const scope of functionScopes) {
        if (!scope.returnType) continue;

        const { name, startLine, endLine, params, locals, returnType } = scope;
        const expectedType = returnType.toUpperCase();

        const body = funcBodyByLine.get(startLine) || [];
        const returnStmts = collectReturnStatements(body);

        const hasReturn = returnStmts.length > 0;

        for (const returnStmt of returnStmts) {
            const actualType = inferReturnExpressionType(returnStmt, params, locals);

            if (actualType && actualType !== expectedType) {
                const lineText = document.lineAt(returnStmt.line).text;
                const col = lineText.search(/\bRETURN\b/i);
                const diag = createDiagnostic(
                    {
                        start: { line: returnStmt.line, character: col === -1 ? 0 : col },
                        end:   { line: returnStmt.line, character: lineText.length }
                    },
                    {
                        messageId: "ACU-FNC-010",
                        params: {
                            functionName: name,
                            expectedType,
                            expression: returnStmt.value
                                ? returnStmt.value.tokens.map(t => t.value).join("").trim()
                                : "",
                            actualType
                        }
                    },
                    getRuleSeverity
                );
                if (diag) diagnostics.push(diag);
            }
        }

        if (!hasReturn) {
            const lineText = document.lineAt(endLine).text;
            const col = lineText.search(/\bENDFUNCTION\b/i);
            const diag = createDiagnostic(
                {
                    start: { line: endLine, character: col === -1 ? 0 : col },
                    end:   { line: endLine, character: col === -1 ? lineText.length : col + "ENDFUNCTION".length }
                },
                {
                    messageId: "ACU-FNC-011",
                    params: { functionName: name, returnType: expectedType }
                },
                getRuleSeverity
            );
            if (diag) diagnostics.push(diag);
        }
    }

    return diagnostics;
}

export { validateReturnTypes, inferReturnExpressionType };
