import { Node, AbstractVisitor, Program, VariableDeclarations, VariableDeclarator, Monitor, ProcedureDeclaration, ConditionDeclaration, MainAgent, ProcedureParameter, Struct, VariableInitializer, IfStatement, BreakStatement, Case, ConditionStatement, ContinueStatement, DoStatement, ForInit, ForStatement, PrimitiveStatement, PrintStatement, ReturnStatement, SelectStatement, StatementBlock, StatementExpression, StatementWrapper, WhileStatement, AssignDestination, Expression, TypeNode } from "@pseuco/lang"
import { PseucoError } from "../PseucoError";
import { BlockEnvironment } from "./environments/BlockEnvironment";
import { ClassesEnvironment } from "./environments/ClassesEnvironment";
import { ProceduresEnvironment } from "./environments/ProceduresEnvironment";
import { RootEnvironment } from "./environments/RootEnvironment";
import { TypeEnvironment } from "./environments/TypeEnvironment";
import { TypeCheckCoordinator } from "./TypeCheckCoordinator";

export class TypeEnvironmentVisitor extends AbstractVisitor<void> {
    private constructor(readonly coordinator: TypeCheckCoordinator) {
        super();
    }


    static getEnvironmentForProgram(coordinator: TypeCheckCoordinator, program: Program): Map<Node, TypeEnvironment> {
        const visitor = new TypeEnvironmentVisitor(coordinator);
        try {
            visitor.go(program);
        } catch (error: any) {
            if (error instanceof PseucoError) {
                coordinator.submitAndThrowPseucoError(error);
            } else {
                coordinator.submitAndThrowTypeError(program, error.message ?? error.toString());
            }
        }
        

        return visitor.result;
    }

    private readonly result: Map<Node, TypeEnvironment> = new Map();

    private rootEnvironment: TypeEnvironment = new RootEnvironment();
    private currentEnvironment: TypeEnvironment = this.rootEnvironment;

    private go(node: Node) {
        if (this.result.has(node)) {
            throw Error("Internal Error: Node environment already computed!");
        }
        this.result.set(node, this.currentEnvironment);
        node.accept(this);
        // try {
        //     node.accept(this);
        // } catch (error: any) {
        //     if (error instanceof PseucoError) {
        //         this.coordinator.submitAndThrowPseucoError(error);
        //     } else {
        //         this.coordinator.submitAndThrowTypeError(node, error.message ?? error.toString());
        //     }
        // }
        
    }

    // Auxiliary method to track back environments (e.g., for blocks of curly braces { ... })
    private executeInNewBlockEnvironment(block: () => void, parentEnvironment: TypeEnvironment = this.currentEnvironment) {
        //this.currentEnvironment = new BlockEnvironment(parentEnvironment);
        this.currentEnvironment = (parentEnvironment instanceof BlockEnvironment) ? parentEnvironment : new BlockEnvironment(parentEnvironment);
        block();
        this.currentEnvironment = parentEnvironment;
    }

    
    // Some notes:
    // 1. Invariant: On calling this.go, this.currentEnvironment always is the environment that should be assigned to the visited node
    // 2. It is important to visit all children / to assign an env to every node in the tree
    // 3. Certain nodes must restore the environment, e.g., monitors


    /***************************** 
     **** DECLARATIONS ***********
     *****************************/

    override visitProgram(node: Program) {
        this.currentEnvironment = new ClassesEnvironment(this.currentEnvironment);
        node.childNodes.forEach(child => this.go(child));
    }

    override visitVariableDeclarations(node: VariableDeclarations) {
        node.childNodes.forEach(decl => this.go(decl));
    }

    override visitVariableDeclarator(node: VariableDeclarator) {
        if (node.initializer != null) {
            // The initializer needs the environment *without* the variable being declared
            this.go(node.initializer);
        }
        if (this.currentEnvironment.getDeclarationForVariableIdentifier(node.name) != undefined) {
            this.coordinator.submitAndThrowTypeError(node, `Variable ${node.name} is already declared and cannot be declared again.`);
        }
        this.currentEnvironment = this.currentEnvironment.addVariable(node.name, node);
    }

    override visitVariableInitializer(node: VariableInitializer) {
        node.childNodes.forEach(child => this.go(child));
    }

    override visitMonitor(node: Monitor) {
        const env = this.currentEnvironment.addClass(node.name, node);
        this.currentEnvironment = new ProceduresEnvironment(env, node);
        node.declarations.forEach(child => this.go(child));
        this.currentEnvironment = env;
    }

    override visitStruct(node: Struct) {
        const env = this.currentEnvironment.addClass(node.name, node);
        this.currentEnvironment = new ProceduresEnvironment(env, node);
        node.declarations.forEach(child => this.go(child));
        this.currentEnvironment = env;
    }

    override visitProcedureDeclaration(node: ProcedureDeclaration) {
        this.executeInNewBlockEnvironment(() => {
            node.childNodes.forEach(child => this.go(child));
        }, this.currentEnvironment.addProcedure(node.name, node));
    }

    override visitProcedureParameter(node: ProcedureParameter) {
        node.childNodes.forEach(child => this.go(child));
        this.currentEnvironment = this.currentEnvironment.addVariable(node.name, node);
    }

    override visitConditionDeclaration(node: ConditionDeclaration) {
        this.currentEnvironment = this.currentEnvironment.addCondition(node.name, node);
        node.childNodes.forEach(child => this.go(child));
    }

    override visitMainAgent(node: MainAgent) {
        this.executeInNewBlockEnvironment(() => {
            node.childNodes.forEach(child => this.go(child));
        }, this.currentEnvironment.setMainAgent(node));
    }

    




    /***************************** 
     **** STATEMENTS *************
     *****************************/

     override visitIfStatement(node: IfStatement) {
        // Consequence and Alternative do NOT need to be executed in separate blocks, because this is done by StatementBlock
        node.childNodes.forEach(child => this.go(child));
     }

     override visitBreakStatement(node: BreakStatement) {
         // Nothing to do
     }

     override visitCase(node: Case) {
        // Statement does NOT need to be executed in separate blocks, because this is done by StatementBlock
        node.childNodes.forEach(child => this.go(child));
     }

     override visitConditionStatement(node: ConditionStatement) {
        node.childNodes.forEach(child => this.go(child));
     }

     override visitContinueStatement(node: ContinueStatement) {
         // Nothing to do
     }

     override visitDoStatement(node: DoStatement) {
        node.childNodes.forEach(child => this.go(child));
     }

     override visitForInit(node: ForInit) {
        node.childNodes.forEach(child => this.go(child));
     }

     override visitForStatement(node: ForStatement) {
         this.executeInNewBlockEnvironment(() => {
            // Needs an extra block environment to enclose new variables in ForInit
            node.childNodes.forEach(child => this.go(child));
         });
     }

     override visitPrimitiveStatement(node: PrimitiveStatement) {
         this.go(node.expression);
     }

     override visitPrintStatement(node: PrintStatement) {
         node.childNodes.forEach(child => this.go(child));
     }

     override visitReturnStatement(node: ReturnStatement) {
         if (node.expression) {
             this.go(node.expression);
         }
     }

     override visitSelectStatement(node: SelectStatement) {
        node.childNodes.forEach(child => this.go(child));
     }

     override visitStatementBlock(node: StatementBlock) {
         this.executeInNewBlockEnvironment(() => node.childNodes.forEach(child => this.go(child)));
         // ToDo: There might be some empty environments now, because executeInNewBlockEnvironment is also called before go-calls on StatementBlocks
     }

     override visitStatementExpression(node: StatementExpression) {
         this.go(node.expression);
     }

     override visitStatementWrapper(node: StatementWrapper) {
        if (node.statement) {
            this.go(node.statement);
        }
     }

     override visitWhileStatement(node: WhileStatement) {
        node.childNodes.forEach(child => this.go(child));
        // loopBody is not necessarily a {}-block
     }




     /***************************** 
     **** EXPRESSIONS *************
     *****************************/

     override visitAssignDestination(node: AssignDestination) {
         node.childNodes.forEach(child => this.go(child));
     }

     private visitExpression(node: Expression) {
        node.childNodes.forEach(child => this.go(child));
     }

     visitAdditiveExpression = this.visitExpression;
     visitAndExpression = this.visitExpression;
     visitArrayExpression = this.visitExpression;
     visitAssignExpression = this.visitExpression;
     visitClassCall = this.visitExpression;
     visitConditionalExpression = this.visitExpression;
     visitEqualityExpression = this.visitExpression;
     visitIdentifierExpression = this.visitExpression;
     visitLiteralExpression = this.visitExpression;
     visitMultiplicativeExpression = this.visitExpression;
     visitOrExpression = this.visitExpression;
     visitPostfixExpression = this.visitExpression;
     visitProcedureCall = this.visitExpression;
     visitReceiveExpression = this.visitExpression;
     visitRelationalExpression = this.visitExpression;
     visitSendExpression = this.visitExpression;
     visitStartExpression = this.visitExpression;
     visitUnaryExpression = this.visitExpression;





     /***************************** 
     **** TYPE NODES **************
     *****************************/

     private visitTypeNode(node: TypeNode) {
        node.childNodes.forEach(child => this.go(child));
     }

     visitArrayTypeNode = this.visitTypeNode;
     visitChannelTypeNode = this.visitTypeNode;
     visitClassTypeNode = this.visitTypeNode;
     visitSimpleTypeNode = this.visitTypeNode;



}