/**
 * Verification logic for calculate blocks.
 */

import { Block } from './block.js';
import { Expression } from './expression.js';
import { Context } from './context.js';
import { VerificationError } from './errors.js';
import { simplify } from './simplifier.js';

/**
 * Verify a calculate block and return true if goal is resolved.
 */
export function verifyCalculate(
    step: Block,
    context: Context,
    axiomsAndTheorems: Map<string, Block>
): boolean {
    // 1. Validate that the goal is an equality
    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal for calculate`);
    }

    if (context.goal.name !== 'eq' || context.goal.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: Goal must be eq(a, b), got ${context.goal}`
        );
    }

    const initialExpr = context.goal.children[0];
    const finalExpr = context.goal.children[1];

    // 2. Validate calculate structure
    if (step.args.length !== 1) {
        throw new VerificationError(
            `Line ${step.lineNum}: 'calculate' must have exactly one argument (initial expression)`
        );
    }

    const calcInitial = step.args[0];
    if (!calcInitial.equals(initialExpr)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Calculate initial expression ${calcInitial} does not match goal LHS ${initialExpr}`
        );
    }

    if (step.children.length === 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: Calculate block must have at least one step`
        );
    }

    // 3. Process each calculation step
    let currentExpr = initialExpr;

    for (const child of step.children) {
        if (child.blockType !== '=') {
            throw new VerificationError(
                `Line ${child.lineNum}: Calculate children must be '=' blocks, got '${child.blockType}'`
            );
        }

        currentExpr = verifyCalcStep(child, currentExpr, context, axiomsAndTheorems);
    }

    // 4. Final check: current expression should equal goal's RHS
    if (!currentExpr.equals(finalExpr)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Calculate final expression ${currentExpr} does not match goal RHS ${finalExpr}`
        );
    }

    return true; // Goal resolved
}

/**
 * Verify a single calculation step and return the new expression.
 */
function verifyCalcStep(
    step: Block,
    prevExpr: Expression,
    context: Context,
    axiomsAndTheorems: Map<string, Block>
): Expression {
    if (step.args.length < 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: '=' must have at least 2 arguments (new_expr and justification)`
        );
    }

    const newExpr = step.args[0];
    const justificationType = step.args[1].name;

    if (step.args[1].children.length > 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: Justification type must be a simple identifier (by-lhs or by-rhs)`
        );
    }

    // Get the equation (lhs, rhs) based on justification
    const remainingArgs = step.args.slice(2);
    let lhs: Expression, rhs: Expression;

    if (justificationType === 'by-lhs') {
        [lhs, rhs] = getEquationForJustification(step, context, axiomsAndTheorems, remainingArgs);
        // by-lhs: verify that prevExpr can be transformed to newExpr by replacing lhs with rhs where needed
        if (!verifyTransformation(prevExpr, newExpr, lhs, rhs)) {
            throw new VerificationError(
                `Line ${step.lineNum}: Cannot transform ${prevExpr} to ${newExpr} using by-lhs ${lhs} = ${rhs}`
            );
        }
    } else if (justificationType === 'by-rhs') {
        [lhs, rhs] = getEquationForJustification(step, context, axiomsAndTheorems, remainingArgs);
        // by-rhs: verify that prevExpr can be transformed to newExpr by replacing rhs with lhs where needed
        if (!verifyTransformation(prevExpr, newExpr, rhs, lhs)) {
            throw new VerificationError(
                `Line ${step.lineNum}: Cannot transform ${prevExpr} to ${newExpr} using by-rhs ${lhs} = ${rhs}`
            );
        }
    } else {
        throw new VerificationError(
            `Line ${step.lineNum}: Unknown justification type '${justificationType}', expected 'by-lhs' or 'by-rhs'`
        );
    }

    return newExpr;
}

/**
 * Get the (lhs, rhs) equation from either a variable or axiom/theorem.
 */
function getEquationForJustification(
    step: Block,
    context: Context,
    axiomsAndTheorems: Map<string, Block>,
    args: Expression[]
): [Expression, Expression] {
    if (args.length === 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: Justification requires at least one argument`
        );
    }

    const firstArg = args[0];

    // First argument must be a simple identifier
    if (firstArg.children.length > 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: First justification argument must be a simple identifier (variable or axiom name)`
        );
    }

    const name = firstArg.name;

    // Check if it's a variable in context (with only one arg total)
    if (args.length === 1 && context.hasVariable(name)) {
        // This is a variable reference
        return getEquationFromVariable(step, name, context);
    }

    // Otherwise, treat it as an axiom/theorem
    const axiomArgs = args.slice(1); // Remaining args are the axiom arguments
    return getEquationFromAxiom(step, name, axiomArgs, context, axiomsAndTheorems);
}

/**
 * Extract (lhs, rhs) from a variable with type eq(lhs, rhs).
 */
function getEquationFromVariable(
    step: Block,
    varName: string,
    context: Context
): [Expression, Expression] {
    if (!context.hasVariable(varName)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Variable '${varName}' not in context`
        );
    }

    const varType = context.getVariableType(varName)!;

    if (varType.name !== 'eq' || varType.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: Variable '${varName}' must have type eq(a, b), got ${varType}`
        );
    }

    return [varType.children[0], varType.children[1]];
}

/**
 * Extract (lhs, rhs) from an axiom/theorem after substituting arguments.
 */
function getEquationFromAxiom(
    step: Block,
    axiomName: string,
    args: Expression[],
    context: Context,
    axiomsAndTheorems: Map<string, Block>
): [Expression, Expression] {
    if (!axiomsAndTheorems.has(axiomName)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Axiom or theorem '${axiomName}' not found`
        );
    }

    const axiomBlock = axiomsAndTheorems.get(axiomName)!;

    // Extract suppose blocks
    const supposeBlocks: Block[] = [];
    let concludeBlock: Block | null = null;

    for (const child of axiomBlock.children) {
        if (child.blockType === 'suppose') {
            supposeBlocks.push(child);
        } else if (child.blockType === 'conclude') {
            concludeBlock = child;
        }
    }

    if (concludeBlock === null) {
        throw new VerificationError(
            `Line ${step.lineNum}: Axiom '${axiomName}' has no conclude block`
        );
    }

    // Verify number of arguments matches
    if (args.length !== supposeBlocks.length) {
        throw new VerificationError(
            `Line ${step.lineNum}: Axiom '${axiomName}' expects ${supposeBlocks.length} arguments, got ${args.length}`
        );
    }

    // Build substitution map
    const substitutions = new Map<string, Expression>();

    for (let i = 0; i < supposeBlocks.length; i++) {
        const suppose = supposeBlocks[i];
        const arg = args[i];

        // Parse suppose: suppose name : type
        if (suppose.args.length !== 1) {
            throw new VerificationError(
                `Line ${step.lineNum}: Invalid suppose block in axiom '${axiomName}'`
            );
        }

        const supposeArg = suppose.args[0];
        if (supposeArg.name !== 'var' || supposeArg.children.length !== 2) {
            throw new VerificationError(
                `Line ${step.lineNum}: Invalid suppose format in axiom '${axiomName}'`
            );
        }

        const varName = supposeArg.children[0].name;
        substitutions.set(varName, arg);
    }

    // Get the conclusion and verify it's an equation
    let conclusion = concludeBlock.args[0];

    // Apply substitutions to the conclusion
    for (const [varName, argExpr] of substitutions) {
        conclusion = conclusion.substitute(varName, argExpr);
    }

    if (conclusion.name !== 'eq' || conclusion.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: Axiom '${axiomName}' conclusion must be eq(a, b), got ${conclusion}`
        );
    }

    return [conclusion.children[0], conclusion.children[1]];
}

/**
 * Verify that prevExpr can be transformed to targetExpr by replacing fromExpr with toExpr where needed.
 */
function verifyTransformation(
    prevExpr: Expression,
    targetExpr: Expression,
    fromExpr: Expression,
    toExpr: Expression
): boolean {
    // If they're already equal, no substitution needed
    if (prevExpr.equals(targetExpr)) {
        return true;
    }

    // Try substituting at the root level
    if (prevExpr.equals(fromExpr)) {
        return toExpr.equals(targetExpr);
    }

    // If names don't match or different number of children, can't transform
    if (prevExpr.name !== targetExpr.name || prevExpr.children.length !== targetExpr.children.length) {
        return false;
    }

    // Recursively check all children
    return prevExpr.children.every((pc, i) =>
        verifyTransformation(pc, targetExpr.children[i], fromExpr, toExpr)
    );
}