/**
 * Proof verification logic.
 */

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';
import { verifyCalculate } from './calculateVerifier.js';

/**
 * Verify all axioms and theorems in a file.
 */
export function verifyFile(blocks: Block[], verbose: boolean = false): void {
    // First pass: collect all axioms and theorems
    const axiomsAndTheorems = new Map<string, Block>();

    for (const block of blocks) {
        if (block.blockType === 'axiom' || block.blockType === 'theorem') {
            if (block.args.length !== 1) {
                throw new VerificationError(
                    `Line ${block.lineNum}: ${block.blockType} must have exactly one argument (the name)`
                );
            }

            const name = block.args[0].name;
            if (block.args[0].children.length > 0) {
                throw new VerificationError(
                    `Line ${block.lineNum}: ${block.blockType} name must be a simple identifier`
                );
            }

            axiomsAndTheorems.set(name, block);

            if (verbose) {
                console.log(`Found ${block.blockType}: ${name}`);
            }
        } else {
            throw new VerificationError(
                `Line ${block.lineNum}: Top-level blocks must be 'axiom' or 'theorem', got '${block.blockType}'`
            );
        }
    }

    // Second pass: verify theorems
    for (const block of blocks) {
        if (block.blockType === 'theorem') {
            const name = block.args[0].name;
            if (verbose) {
                console.log(`\nVerifying theorem: ${name}`);
            }
            verifyTheorem(block, axiomsAndTheorems, verbose);
            if (verbose) {
                console.log(`✓ Theorem ${name} verified`);
            }
        }
    }
}

/**
 * Verify a single theorem.
 */
function verifyTheorem(theoremBlock: Block, axiomsAndTheorems: Map<string, Block>, verbose: boolean): void {
    // Create initial context from suppose blocks
    const context = new Context();

    const supposeBlocks: Block[] = [];
    let concludeBlock: Block | null = null;
    let proofBlock: Block | null = null;

    // Parse structure: suppose* conclude proof
    for (const child of theoremBlock.children) {
        if (child.blockType === 'suppose') {
            supposeBlocks.push(child);
        } else if (child.blockType === 'conclude') {
            if (concludeBlock !== null) {
                throw new VerificationError(`Line ${child.lineNum}: Multiple 'conclude' blocks in theorem`);
            }
            concludeBlock = child;
        } else if (child.blockType === 'proof') {
            if (proofBlock !== null) {
                throw new VerificationError(`Line ${child.lineNum}: Multiple 'proof' blocks in theorem`);
            }
            proofBlock = child;
        } else {
            throw new VerificationError(
                `Line ${child.lineNum}: Invalid block type '${child.blockType}' in theorem`
            );
        }
    }

    if (concludeBlock === null) {
        throw new VerificationError(`Line ${theoremBlock.lineNum}: Theorem missing 'conclude' block`);
    }

    if (proofBlock === null) {
        throw new VerificationError(`Line ${theoremBlock.lineNum}: Theorem missing 'proof' block`);
    }

    // Process suppose blocks
    for (const suppose of supposeBlocks) {
        processSuppose(suppose, context);
    }

    // Process conclude block
    processConclude(concludeBlock, context);

    // Verify the proof
    const goalResolved = verifyProof(proofBlock, context, axiomsAndTheorems, verbose);

    if (!goalResolved) {
        throw new VerificationError(`Line ${proofBlock.lineNum}: Proof does not resolve the goal`);
    }
}

/**
 * Process a suppose block: suppose name : type
 */
function processSuppose(supposeBlock: Block, context: Context): void {
    if (supposeBlock.args.length !== 1) {
        throw new VerificationError(`Line ${supposeBlock.lineNum}: 'suppose' must have exactly one argument`);
    }

    const arg = supposeBlock.args[0];
    if (arg.name !== 'var' || arg.children.length !== 2) {
        throw new VerificationError(
            `Line ${supposeBlock.lineNum}: 'suppose' argument must be of form 'name : type'`
        );
    }

    const varName = arg.children[0].name;
    const varType = arg.children[1];

    if (arg.children[0].children.length > 0) {
        throw new VerificationError(`Line ${supposeBlock.lineNum}: Variable name must be a simple identifier`);
    }

    context.addVariable(varName, varType);
}

/**
 * Process a conclude block: conclude type
 */
function processConclude(concludeBlock: Block, context: Context): void {
    if (concludeBlock.args.length !== 1) {
        throw new VerificationError(`Line ${concludeBlock.lineNum}: 'conclude' must have exactly one argument`);
    }

    context.goal = concludeBlock.args[0];
}

/**
 * Verify a proof block and return whether the goal was resolved.
 */
function verifyProof(proofBlock: Block, context: Context, axiomsAndTheorems: Map<string, Block>, verbose: boolean): boolean {
    let goalResolved = false;

    for (const child of proofBlock.children) {
        const result = verifyProofStep(child, context, axiomsAndTheorems, verbose);
        if (result) {
            goalResolved = true;
        }
    }

    return goalResolved;
}

/**
 * Verify a single proof step. Returns true if the goal was resolved.
 */
function verifyProofStep(step: Block, context: Context, axiomsAndTheorems: Map<string, Block>, verbose: boolean): boolean {
    switch (step.blockType) {
        case 'unpack-and':
            return verifyUnpackAnd(step, context, axiomsAndTheorems, verbose);
        case 'cases':
            return verifyCases(step, context, axiomsAndTheorems, verbose);
        case 'witness':
            return verifyWitness(step, context);
        case 'assert-goal':
            return verifyAssertGoal(step, context);
        case 'exact':
            return verifyExact(step, context);
        case 'calculate':
            return verifyCalculate(step, context, axiomsAndTheorems);
        case 'assert':
            return verifyAssert(step, context);
        case 'define':
            return verifyDefine(step, context);
        case 'consider':
            return verifyConsider(step, context);
        case 'forall-apply':
            return verifyForallApply(step, context);
        case 'deconstruct-exists':
            return verifyDeconstructExists(step, context);
        case 'we-have':
            return verifyWeHave(step, context, axiomsAndTheorems, verbose);
        case 'focus-or':
            return verifyFocusOr(step, context);
        default:
            throw new VerificationError(`Line ${step.lineNum}: Unknown proof step type '${step.blockType}'`);
    }
}

/**
 * Verify unpack-and: goal must be and(a, b), requires two goal children.
 */
function verifyUnpackAnd(step: Block, context: Context, axiomsAndTheorems: Map<string, Block>, verbose: boolean): boolean {
    if (step.args.length !== 0) {
        throw new VerificationError(`Line ${step.lineNum}: 'unpack-and' takes no arguments`);
    }

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal to unpack`);
    }

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

    if (step.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'unpack-and' requires exactly 2 children`);
    }

    // Verify both children are goal blocks
    for (let i = 0; i < step.children.length; i++) {
        const child = step.children[i];
        if (child.blockType !== 'goal') {
            throw new VerificationError(`Line ${child.lineNum}: 'unpack-and' children must be 'goal' blocks`);
        }

        if (child.args.length !== 1) {
            throw new VerificationError(`Line ${child.lineNum}: 'goal' must have exactly one argument`);
        }

        const expectedGoal = context.goal.children[i];
        if (!child.args[0].equals(expectedGoal)) {
            throw new VerificationError(
                `Line ${child.lineNum}: Expected goal ${expectedGoal}, got ${child.args[0]}`
            );
        }

        // Each goal must have one proof child
        if (child.children.length !== 1 || child.children[0].blockType !== 'proof') {
            throw new VerificationError(`Line ${child.lineNum}: 'goal' must have exactly one 'proof' child`);
        }

        // Verify the proof with the subgoal
        const subContext = context.copy();
        subContext.goal = expectedGoal;
        const resolved = verifyProof(child.children[0], subContext, axiomsAndTheorems, verbose);

        if (!resolved) {
            throw new VerificationError(`Line ${child.lineNum}: Proof does not resolve goal ${expectedGoal}`);
        }
    }

    return true; // Goal resolved
}

/**
 * Verify cases: variable must have type or(a, b), requires two case children.
 */
function verifyCases(step: Block, context: Context, axiomsAndTheorems: Map<string, Block>, verbose: boolean): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'cases' takes exactly one argument (variable name)`);
    }

    const varName = step.args[0].name;
    if (step.args[0].children.length > 0) {
        throw new VerificationError(`Line ${step.lineNum}: 'cases' argument must be a simple variable name`);
    }

    if (!context.hasVariable(varName)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${varName}' not in context`);
    }

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

    if (step.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'cases' requires exactly 2 'case' children`);
    }

    // Verify both case children
    for (let i = 0; i < step.children.length; i++) {
        const child = step.children[i];
        if (child.blockType !== 'case') {
            throw new VerificationError(`Line ${child.lineNum}: 'cases' children must be 'case' blocks`);
        }

        if (child.args.length !== 1) {
            throw new VerificationError(`Line ${child.lineNum}: 'case' must have exactly one argument`);
        }

        const arg = child.args[0];
        if (arg.name !== 'var' || arg.children.length !== 2) {
            throw new VerificationError(`Line ${child.lineNum}: 'case' argument must be of form 'name : type'`);
        }

        const caseVarName = arg.children[0].name;
        const caseType = arg.children[1];

        const expectedType = varType.children[i];
        if (!caseType.equals(expectedType)) {
            throw new VerificationError(
                `Line ${child.lineNum}: Expected case type ${expectedType}, got ${caseType}`
            );
        }

        // Verify proof with added case variable
        const caseContext = context.copy();
        caseContext.addVariable(caseVarName, caseType);

        // Each case needs a proof
        if (child.children.length !== 1 || child.children[0].blockType !== 'proof') {
            throw new VerificationError(`Line ${child.lineNum}: 'case' must have exactly one 'proof' child`);
        }

        const resolved = verifyProof(child.children[0], caseContext, axiomsAndTheorems, verbose);
        if (!resolved) {
            throw new VerificationError(`Line ${child.lineNum}: Proof in case does not resolve goal`);
        }
    }

    return true; // Goal resolved
}

/**
 * Verify witness: goal must be exists(var(name, type), body), modifies goal.
 */
function verifyWitness(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'witness' takes exactly one argument`);
    }

    const witnessName = step.args[0].name;
    if (step.args[0].children.length > 0) {
        throw new VerificationError(`Line ${step.lineNum}: 'witness' argument must be a simple identifier`);
    }

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal for witness`);
    }

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

    const varExpr = context.goal.children[0];
    if (varExpr.name !== 'var' || varExpr.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: exists must have var(name, type) as first argument`
        );
    }

    const boundVarName = varExpr.children[0].name;
    const body = context.goal.children[1];

    // Substitute witness into body
    context.goal = body.substitute(boundVarName, new Expression(witnessName));

    return false; // Does not resolve goal
}

/**
 * Verify assert-goal: expression must match current goal.
 */
function verifyAssertGoal(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'assert-goal' takes exactly one argument`);
    }

    const expectedGoal = step.args[0];

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No current goal`);
    }

    if (!context.goal.equals(expectedGoal)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Expected goal ${expectedGoal}, but current goal is ${context.goal}`
        );
    }

    return false; // Does not resolve goal
}

/**
 * Verify exact: two variants - variable name or expression after simplification.
 */
function verifyExact(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'exact' takes exactly one argument`);
    }

    const arg = step.args[0];

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal to resolve`);
    }

    // Variant 1: Simple variable name
    if (arg.children.length === 0 && context.hasVariable(arg.name)) {
        const varType = context.getVariableType(arg.name)!;
        if (!varType.equals(context.goal)) {
            throw new VerificationError(
                `Line ${step.lineNum}: Variable '${arg.name}' has type ${varType}, but goal is ${context.goal}`
            );
        }
        return true; // Goal resolved
    }

    // Variant 2: Expression after simplification
    const simplified = simplify(arg, context);
    if (!simplified.equals(context.goal)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Expression ${arg} simplifies to ${simplified}, but goal is ${context.goal}`
        );
    }

    return true; // Goal resolved
}

/**
 * Verify assert: assert name : type.
 */
function verifyAssert(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'assert' must have exactly one argument`);
    }

    const arg = step.args[0];
    if (arg.name !== 'var' || arg.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'assert' argument must be of form 'name : type'`);
    }

    const varName = arg.children[0].name;
    const expectedType = arg.children[1];

    if (!context.hasVariable(varName)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${varName}' not in context`);
    }

    const actualType = context.getVariableType(varName)!;
    if (!actualType.equals(expectedType)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Variable '${varName}' has type ${actualType}, expected ${expectedType}`
        );
    }

    return false; // Does not resolve goal
}

/**
 * Verify define: define name : eq(lhs, rhs).
 */
function verifyDefine(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'define' must have exactly one argument`);
    }

    const arg = step.args[0];
    if (arg.name !== 'var' || arg.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'define' argument must be of form 'name : type'`);
    }

    const varName = arg.children[0].name;
    const varType = arg.children[1];

    if (context.hasVariable(varName)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${varName}' already in context`);
    }

    // Check that type is an eq expression
    if (varType.name !== 'eq' || varType.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'define' type must be eq(lhs, rhs)`);
    }

    context.addVariable(varName, varType);

    return false; // Does not resolve goal
}

/**
 * Verify consider: goal must be forall(var(name, type), body).
 */
function verifyConsider(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'consider' takes exactly one argument`);
    }

    const varName = step.args[0].name;
    if (step.args[0].children.length > 0) {
        throw new VerificationError(`Line ${step.lineNum}: 'consider' argument must be a simple identifier`);
    }

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal for consider`);
    }

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

    const varExpr = context.goal.children[0];
    if (varExpr.name !== 'var' || varExpr.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: forall must have var(name, type) as first argument`
        );
    }

    const boundVarName = varExpr.children[0].name;
    const varType = varExpr.children[1];
    const body = context.goal.children[1];

    if (boundVarName !== varName) {
        throw new VerificationError(
            `Line ${step.lineNum}: Expected variable '${boundVarName}', got '${varName}'`
        );
    }

    context.addVariable(varName, varType);
    context.goal = body;

    return false; // Does not resolve goal
}

/**
 * Verify forall-apply: apply forall to an argument.
 */
function verifyForallApply(step: Block, context: Context): boolean {
    if (step.args.length !== 3) {
        throw new VerificationError(`Line ${step.lineNum}: 'forall-apply' takes exactly 3 arguments`);
    }

    const forallVar = step.args[0].name;
    const argVar = step.args[1].name;
    const resultVar = step.args[2].name;

    if (step.args[0].children.length > 0 || step.args[1].children.length > 0 || step.args[2].children.length > 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: 'forall-apply' arguments must be simple identifiers`
        );
    }

    if (!context.hasVariable(forallVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${forallVar}' not in context`);
    }

    if (!context.hasVariable(argVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${argVar}' not in context`);
    }

    if (context.hasVariable(resultVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${resultVar}' already in context`);
    }

    const forallType = context.getVariableType(forallVar)!;
    if (forallType.name !== 'forall' || forallType.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: Variable '${forallVar}' must have type forall(var(name, type), body)`
        );
    }

    const varExpr = forallType.children[0];
    if (varExpr.name !== 'var' || varExpr.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: forall must have var(name, type) as first argument`
        );
    }

    const boundVarName = varExpr.children[0].name;
    const expectedArgType = varExpr.children[1];
    const body = forallType.children[1];

    const argType = context.getVariableType(argVar)!;
    if (!argType.equals(expectedArgType)) {
        throw new VerificationError(
            `Line ${step.lineNum}: Argument '${argVar}' has type ${argType}, expected ${expectedArgType}`
        );
    }

    // Substitute argVar for boundVarName in body
    const resultType = body.substitute(boundVarName, new Expression(argVar));
    context.addVariable(resultVar, resultType);

    return false; // Does not resolve goal
}

/**
 * Verify deconstruct-exists: extract witness and hypothesis from exists.
 */
function verifyDeconstructExists(step: Block, context: Context): boolean {
    if (step.args.length !== 3) {
        throw new VerificationError(`Line ${step.lineNum}: 'deconstruct-exists' takes exactly 3 arguments`);
    }

    const existsVar = step.args[0].name;
    const witnessVar = step.args[1].name;
    const hypVar = step.args[2].name;

    if (step.args[0].children.length > 0 || step.args[1].children.length > 0 || step.args[2].children.length > 0) {
        throw new VerificationError(
            `Line ${step.lineNum}: 'deconstruct-exists' arguments must be simple identifiers`
        );
    }

    if (!context.hasVariable(existsVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${existsVar}' not in context`);
    }

    if (context.hasVariable(witnessVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${witnessVar}' already in context`);
    }

    if (context.hasVariable(hypVar)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${hypVar}' already in context`);
    }

    const existsType = context.getVariableType(existsVar)!;
    if (existsType.name !== 'exists' || existsType.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: Variable '${existsVar}' must have type exists(var(name, type), body)`
        );
    }

    const varExpr = existsType.children[0];
    if (varExpr.name !== 'var' || varExpr.children.length !== 2) {
        throw new VerificationError(
            `Line ${step.lineNum}: exists must have var(name, type) as first argument`
        );
    }

    const boundVarName = varExpr.children[0].name;
    const witnessType = varExpr.children[1];
    const body = existsType.children[1];

    // Add witness variable and hypothesis variable
    context.addVariable(witnessVar, witnessType);
    const hypType = body.substitute(boundVarName, new Expression(witnessVar));
    context.addVariable(hypVar, hypType);

    return false; // Does not resolve goal
}

/**
 * Verify we-have: prove intermediate result.
 */
function verifyWeHave(step: Block, context: Context, axiomsAndTheorems: Map<string, Block>, verbose: boolean): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'we-have' must have exactly one argument`);
    }

    const arg = step.args[0];
    if (arg.name !== 'var' || arg.children.length !== 2) {
        throw new VerificationError(`Line ${step.lineNum}: 'we-have' argument must be of form 'name : type'`);
    }

    const varName = arg.children[0].name;
    const varType = arg.children[1];

    if (context.hasVariable(varName)) {
        throw new VerificationError(`Line ${step.lineNum}: Variable '${varName}' already in context`);
    }

    if (step.children.length !== 1 || step.children[0].blockType !== 'proof') {
        throw new VerificationError(`Line ${step.lineNum}: 'we-have' must have exactly one 'proof' child`);
    }

    // Verify the proof with the new goal
    const subContext = context.copy();
    subContext.goal = varType;
    const resolved = verifyProof(step.children[0], subContext, axiomsAndTheorems, verbose);

    if (!resolved) {
        throw new VerificationError(`Line ${step.lineNum}: Proof does not establish ${varType}`);
    }

    context.addVariable(varName, varType);

    return false; // Does not resolve goal
}

/**
 * Verify focus-or: goal must be or(a, b), focus on left or right.
 */
function verifyFocusOr(step: Block, context: Context): boolean {
    if (step.args.length !== 1) {
        throw new VerificationError(`Line ${step.lineNum}: 'focus-or' takes exactly one argument`);
    }

    const direction = step.args[0].name;
    if (step.args[0].children.length > 0 || (direction !== 'left' && direction !== 'right')) {
        throw new VerificationError(`Line ${step.lineNum}: 'focus-or' argument must be 'left' or 'right'`);
    }

    if (context.goal === null) {
        throw new VerificationError(`Line ${step.lineNum}: No goal for focus-or`);
    }

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

    if (direction === 'left') {
        context.goal = context.goal.children[0];
    } else {
        context.goal = context.goal.children[1];
    }

    return false; // Does not resolve goal
}