import std.stdio;
import std.string;
import std.conv;
import std.algorithm : canFind;

import parser;
import chunk;
import dbg;

struct Local {
    Symbol sym;
    int depth;
}

struct TC {
    bool exact;
    bool maybe;
    OpCode op;
}

struct CompileError {
    string message;
    int line;
}

string printableCompilerError(CompileError err) {
    return format("[line %d] %s", err.line, err.message);
}

class Compiler {

    Function func;
    Parser* parser;
    Local[] locals;
    int localCount;
    int scopeDepth;
    Form current;
    Form previous;

    int currentLine;

    CompileError[] errors;

    Form outer;
    int outerIdx;

    void advance() {
        this.previous = this.current;
        /*
        if (outer.type != FormType.EOF && outerIdx < outer) {
            current = parser.parseForm();
        } else {
            
        }
        */
        this.current = this.parser.parseForm();
    }

    void error(string message, int line) {
        CompileError err = { message, line };
        this.errors ~= err;
    }

    void compileNil(Form form) {
        form.compile(this.func);
        this.advance();
    }

    TC typeCheck(ValueType actual, ValueType expecting) {
        TC ret = { false, false, OpCode.OP_NIL };
        if (actual == expecting) {
            //writeln("good types");
            ret.exact = true;
            return ret;
        } else if (actual == ValueType.ANY) {
            OpCode op;
            switch (expecting) {
                case ValueType.NUMBER:
                    op = OpCode.OP_TYPE_CHECK_NUMBER;
                    break;
                case ValueType.BOOLEAN:
                    op = OpCode.OP_TYPE_CHECK_BOOLEAN;
                    break;
                default:
                    writeln("not sure what type to check :(");
                    break;
            }
            //func.chunk.writeOp(to!ubyte(op), atom.line);
            //tc = { false, true, op };
            ret.maybe = true;
            ret.op = op;
            return ret;
        } else {
            return ret;
        }
        /*
        writeln("in typecheck");
        ValueType[] types = [actual, expecting];
        foreach (ValueType t ; types) {
            switch (t) {
                case ValueType.NUMBER:
                    writeln("number");
                    break;
                case ValueType.BOOLEAN:
                    writeln("boolean");
                    break;
                case ValueType.ANY:
                    writeln("any");
                    break;
                default:
                    writeln("something else");
                    break;
            }
        }
        */
    }

    void compileString(Form form, ValueType expecting) {
        LiteralString ls = cast(LiteralString)form;

        this.func.chunk.writeOp(OpCode.OP_CONSTANT, ls.line);
        String str = new String(ls.value);
        int idx = this.func.chunk.addConstant(makeSeqValue(str));
        this.func.chunk.writeOp(to!ubyte(idx), ls.line);

        this.advance();
    }

    void compileAtom(Form form, const ValueType expecting) {
        Atom atom = cast(Atom)form;
        /*
        ValueType[] types = [expecting, atom.value.type];
        foreach (ValueType t ; types) {
            switch (t) {
                case ValueType.NUMBER:
                    writeln("number");
                    break;
                case ValueType.BOOLEAN:
                    writeln("boolean");
                    break;
                case ValueType.ANY:
                    writeln("any");
                    break;
                default:
                    writeln("something else");
                    break;
            }
        }
        */
        TC tc = this.typeCheck(atom.value.type, expecting);
        form.compile(this.func);
        //if (tc.exact) {
        //} else if (tc.maybe) {

        /*
        if (tc.maybe) {
            form.compile(func);
            func.chunk.writeOp(to!ubyte(tc.op), atom.line);
        } else if (!tc.exact) {
            writefln("COMPILE ERROR: typecheck on line %d", atom.line);
            //form.compile(func);  // don't do this later
        }
        */

        this.advance();





        /*
        ValueType typ = atom.value.type;
        switch (typ) {
            case ValueType.ANY:
                writeln("IN ANY");
                OpCode op;
                switch (expecting) {
                    case ValueType.NUMBER:
                        op = OpCode.OP_TYPE_CHECK_NUMBER;
                        break;
                    case ValueType.BOOLEAN:
                        op = OpCode.OP_TYPE_CHECK_BOOLEAN;
                        break;
                    default:
                        writeln("not sure what type to check :(");
                        break;
                }
                func.chunk.writeOp(to!ubyte(op), atom.line);
                break;
            case expecting:
                form.compile(func);
                break;
            default:
                writefln("COMPILE ERROR: typecheck on line %d", atom.line);
                form.compile(func);
                break;
        }
        */
        /*
        if (atom.value.type == expecting) {
            writefln("COMPILE ERROR: typecheck on line %d", atom.line);
        } else if (atom.value
        */
        //form.compile(func);
        //advance();
    }

    void compileAdd(Form[] args, ValueType expecting) {
        int line = args[0].line;
        this.resolve(args[0], expecting);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, line);

        // (+ n) always returns n
        if (args.length == 1) {
            return;
        }

        for (int i = 1; i < args.length; i++) {
            this.resolve(args[i], expecting);
            this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, line);
            this.func.chunk.writeOp(OpCode.OP_ADD, line);
        }
    }

    void compileNegate(Form arg) {
        this.resolve(arg);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, arg.line);
        this.func.chunk.writeOp(OpCode.OP_NEGATE, arg.line);
    }

    ValueType compileSubtract(Form[] args, ValueType expected) {
        ValueType vt;
        vt = this.resolve(args[0], expected);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[0].line);
        for (int i = 1; i < args.length; i++) {
            vt = this.resolve(args[i], expected);
            this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[i].line);
            this.func.chunk.writeOp(OpCode.OP_SUBTRACT, args[i].line);
        }
        return ValueType.NUMBER;
    }

    //ValueType compileLess(Form[] args, ValueType expected = ValueType.ANY) {
    ValueType compileLess(Form[] args, ValueType expected) {
        if (args.length != 2) {
            /*
            writeln("'<' requires 2 arguments");
            */
            this.error("'<' requires 2 arguments", -1);
            this.advance();
            return ValueType.NIL;
        }
        ValueType vt1 = this.resolve(args[0], ValueType.NUMBER);
        /*
        TC tc1 = typeCheck(vt1, expected);
        if (tc1.maybe) {
            func.chunk.writeOp(to!ubyte(tc1.op), currentLine);
        } else if (!tc1.exact) {
            writeln("COMPILE ERROR: bad type (less)");
        }
        */
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[0].line);

        ValueType vt2 = this.resolve(args[1], ValueType.NUMBER);
        /*
        TC tc2 = typeCheck(vt2, expected);
        if (tc2.maybe) {
            func.chunk.writeOp(to!ubyte(tc2.op), currentLine);
        } else if (!tc2.exact) {
            writeln("COMPILE ERROR: bad type (less)");
        }
        */
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[1].line);


        //ValueType vt1 = resolve(args[0], expected);
        //func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, currentLine);
        //ValueType vt2 = resolve(args[1], expected);
        //func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, currentLine);
        this.func.chunk.writeOp(OpCode.OP_LESS, args[1].line);

        // this should probably return a ValueType
        return ValueType.BOOLEAN;
    }

    ValueType compileGreater(Form[] args) {
        if (args.length != 2) {
            this.error("'>' requires 2 arguments", -1);
            this.advance();
            /*
            writeln("'>' requires 2 arguments");
            */
            return ValueType.NIL;
        }
        ValueType vt1 = this.resolve(args[0], ValueType.NUMBER);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[0].line);

        ValueType vt2 = this.resolve(args[1], ValueType.NUMBER);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_NUMBER, args[1].line);

        this.func.chunk.writeOp(OpCode.OP_GREATER, args[1].line);

        // this should probably return a ValueType
        return ValueType.BOOLEAN;
    }


    //int parseVariable(Def def) {
    int parseVariable(Symbol sym) {

        this.declareVariable(sym);
        if (this.scopeDepth > 0) {
            return 0;
        }

        int addr = this.func.chunk.addConstant(makeStringValue(sym.name));
        return addr;
    }

    void defineVariable(int addr) {
        if (this.scopeDepth > 0) {
            return;
        }
        this.func.chunk.writeOp(OpCode.OP_DEF_GLOBAL, this.current.line);
        this.func.chunk.writeOp(to!ubyte(addr), this.current.line);
    }

    void declareVariable(Symbol sym) {

        if (this.scopeDepth == 0) {
            return;
        }

        for (int i = this.localCount - 1; i >= 0; i--) {
            Local local = this.locals[i];
            if (local.depth != -1 && local.depth < this.scopeDepth) {
                break;
            }

            if (sym.name == local.sym.name) {
                this.error(format("already a variable named '%s' in this scope", sym.name), sym.line);
                //writefln("ERROR: already a variable named '%s' in this scope", sym.name);
                return;
            }
        }

        Local loc = Local(sym, this.scopeDepth);
        if (this.localCount == this.locals.length) {
            this.locals ~= loc;
        } else {
            this.locals[this.localCount] = loc;
        }
        this.localCount++;
    }

    void compileDef(Form form, ValueType expecting) {
        Def def = cast(Def)form;

        // resolve the value
        this.resolve(def.val, expecting);

        // add the variable name to the chunk (if applicable)
        int addr = this.parseVariable(def.name);

        // are we setting a local?
        if (this.scopeDepth > 0) {
            this.func.chunk.writeOp(OpCode.OP_DEF_LOCAL, def.line);
        }

        // define the variable
        this.defineVariable(addr);
    }

    /*
    struct LV {
        int i;
        bool needsTypeCheck;
    }
    */

    int resolveLocal(Symbol sym) { 
        for (int i = this.localCount - 1; i >= 0; i--) {
            Local local = this.locals[i];
            if (local.sym.name == sym.name) {
                return i;
            }
        }

        return -1;
    }

    void compileSymbol(Form form, ValueType expecting = ValueType.ANY) {
        Symbol sym = cast(Symbol)form;

        int arg = this.resolveLocal(sym);
        if (arg != -1) {
            this.func.chunk.writeOp(OpCode.OP_GET_LOCAL, sym.line);
        } else {
            arg = this.func.chunk.addConstant(makeStringValue(sym.name));
            this.func.chunk.writeOp(OpCode.OP_GET_GLOBAL, sym.line);
        }

        // get the variable
        this.func.chunk.writeOp(to!ubyte(arg), sym.line);

        this.advance();
    }

    int jump(OpCode type) {
        this.func.chunk.writeOp(to!ubyte(type), this.current.line);
        this.func.chunk.writeOp(0xff, this.current.line);
        this.func.chunk.writeOp(0xff, this.current.line);
        return to!int(this.func.chunk.code.length) - 2;
    }

    void patchJump(int offset) {
        // additional -2 to account for the 2 byte jump itself
        int jmp = to!int(this.func.chunk.code.length) - offset - 2;

        // TODO check to make sure we didn't jump too far?

        this.func.chunk.code[offset] = (jmp >> 8) & 0xff;
        this.func.chunk.code[offset + 1] = jmp & 0xff;
    }

    void compileIf(Form form) {
        If if_ = cast(If)form;

        this.resolve(if_.cond);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, if_.line);

        int thenJump = this.jump(OpCode.OP_JUMP_IF_FALSE);
        this.func.chunk.writeOp(OpCode.OP_POP, if_.line);

        this.resolve(if_.thenForm);

        int elseJump = this.jump(OpCode.OP_JUMP);

        this.patchJump(thenJump);
        this.func.chunk.writeOp(OpCode.OP_POP, if_.line);

        if (if_.hasElse) {
            this.resolve(if_.elseForm);
        }

        this.patchJump(elseJump);

    }

    void compileAnd(Form form) {
        And and_ = cast(And)form;

        this.resolve(and_.clauses[0]);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, and_.clauses[0].line);

        int[] jumps;
        jumps ~= this.jump(OpCode.OP_JUMP_IF_FALSE);
        int count = 1;
        while (count < and_.clauses.length) {
            this.func.chunk.writeOp(OpCode.OP_POP, and_.clauses[count].line);
            this.resolve(and_.clauses[count]);
            this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, and_.clauses[count].line);
            jumps ~= this.jump(OpCode.OP_JUMP_IF_FALSE);
            count++;
        }

        // patch all the jumps
        foreach (int jmp; jumps) {
            this.patchJump(jmp);
        }
    }

    void compileOr(Form form) {
        Or or_ = cast(Or)form;

        this.resolve(or_.clauses[0]);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, or_.clauses[0].line);
        int[] jumps;
        jumps ~= this.jump(OpCode.OP_JUMP_IF_TRUE);
        int count = 1;
        while (count < or_.clauses.length) {
            this.func.chunk.writeOp(OpCode.OP_POP, or_.clauses[count].line);
            this.resolve(or_.clauses[count]);
            this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, or_.clauses[count].line);
            jumps ~= this.jump(OpCode.OP_JUMP_IF_TRUE);
            count++;
        }

        // patch all the jumps
        foreach (int jmp; jumps) {
            this.patchJump(jmp);
        }
    }

    ValueType compileNot(Form[] args) {
        ValueType vt = this.resolve(args[0], ValueType.BOOLEAN);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_BOOLEAN, args[0].line);
        this.func.chunk.writeOp(OpCode.OP_NOT, args[0].line);
        return vt;
    }

    ValueType compileConcat(Form[] args) {
        ValueType vt = this.resolve(args[0], ValueType.STRING);
        for (int i = 1; i < args.length; i++) {
            vt = this.resolve(args[i], ValueType.STRING);
            this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_SEQ, args[i].line); // TODO wrong
            this.func.chunk.writeOp(OpCode.OP_CONCAT, args[i].line);
        }
        return ValueType.STRING;
    }

    void compileBlock(Form form) {
        Block block = cast(Block)form;
        this.beginScope();
        foreach (Form inner; block.blockBody) {
            this.resolve(inner);
        }
        this.endScope();
    }

    void beginScope() {
        this.scopeDepth++;
    }

    void endScope() {

        this.scopeDepth--;

        while (this.localCount > 0 &&
                this.locals[this.localCount - 1].depth > this.scopeDepth) {
            this.func.chunk.writeOp(OpCode.OP_POPB, -1);
            this.localCount--;
        }
    }

    void call(Form[] args) {
        //ubyte argCount = argumentList();
        foreach (Form f ; args) {
            this.resolve(f);
        }
        this.func.chunk.writeOp(OpCode.OP_CALL, -1);
        //func.chunk.writeOp(argCount, -1);
        this.func.chunk.writeOp(to!ubyte(args.length), -1);
    }

    void compileList(Form[] args) {
        int length = to!int(args.length);
        foreach (Form arg ; args) {
            resolve(arg, ValueType.ANY); // resolve everything onto the stack
        }
        this.func.chunk.writeOp(OpCode.OP_LIST, args[0].line);
        this.func.chunk.writeOp(to!ubyte(length), args[0].line);
    }

    void compileLength(Form[] args) {
        // TODO how do we identify/propagate errors?
        if (args.length != 1) {
            this.error("'length': expected [1] argument, received 0", -1);
            //writeln("COMPILE ERROR: 'first' expects exactly one argument");
            this.advance();
            return;
        }
        ValueType vt = this.resolve(args[0], ValueType.SEQ); // TODO need a new type
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_SEQ, args[0].line);

        this.func.chunk.writeOp(OpCode.OP_LENGTH, args[0].line);
    }

    void compileAppend(Form[] args) {
        if (args.length != 2) {
            this.error(format("'append': expected [2] arguments, received %d", args.length), -1);
            this.advance();
            return;
        }

        ValueType vt1 = this.resolve(args[0], ValueType.SEQ);
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_SEQ, args[0].line);

        // TODO if this is used for both string/list, should be able to typecheck second arg here
        ValueType vt2 = this.resolve(args[1], ValueType.ANY);

        this.func.chunk.writeOp(OpCode.OP_CONCAT, args[0].line);
    }

    void compileFirst(Form[] args) {
        // TODO how do we identify/propagate errors?
        if (args.length != 1) {
            this.error("'first': expected [1] argument, received 0", -1);
            this.advance();
            //writeln("COMPILE ERROR: 'first' expects exactly one argument");
            return;
        }
        ValueType vt = this.resolve(args[0], ValueType.SEQ); // TODO need a new type
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_SEQ, args[0].line);

        // pop the value off the stack
        // get the address of the first item in the list
        // push that item onto the stack

        // or, use a special instruction
        this.func.chunk.writeOp(OpCode.OP_FIRST, args[0].line);
    }

    void compileRest(Form[] args) {
        // TODO how do we identify/propagate errors?
        if (args.length != 1) {
            this.error("'rest': expected [1] argument, received ?", -1);
            this.advance();
            //writeln("COMPILE ERROR: 'first' expects exactly one argument");
            return;
        }
        ValueType vt = this.resolve(args[0], ValueType.OBJ); // TODO need a new type
        this.func.chunk.writeOp(OpCode.OP_TYPE_CHECK_SEQ, args[0].line);

        this.func.chunk.writeOp(OpCode.OP_REST, args[0].line);

        // there's probably a nicer way to copy
        //List lst = new List();

        // create a copy of the list
    }

    void compileLambda(Form form) {
        Lambda lamb = cast(Lambda)form;
        Compiler compiler = new Compiler(ObjType.FUNCTION, this.parser, "<lambda>");
        compiler.beginScope();

        if (lamb.args.type != FormType.NIL) {
            Symbol sym = cast(Symbol)lamb.args.head;
            int constant = compiler.parseVariable(sym);
            compiler.defineVariable(constant);

            foreach (Form inner ; lamb.args.tail) {
                sym = cast(Symbol)inner;
                constant = compiler.parseVariable(sym);
                compiler.defineVariable(constant);
            }
        }

        advance();  // ??

        Block b = new Block(lamb.line);
        b.blockBody = lamb.lambdaBody;
        compiler.compileBlock(b);

        // write the function as a Value
        Function outFunc = compiler.finish();
        this.func.chunk.writeOp(OpCode.OP_CONSTANT, lamb.line);
        int funcAddr = this.func.chunk.addConstant(makeObjValue(outFunc));
        this.func.chunk.writeOp(to!ubyte(funcAddr), lamb.line);

        // capture the errors
        foreach (CompileError err ; compiler.errors) {
            this.errors ~= err;
        }
    }

    void compileFunc(Form form) {
        Func f = cast(Func)form;

        // name the function
        int global = this.parseVariable(f.name);

        Compiler compiler = new Compiler(ObjType.FUNCTION, this.parser, f.name.name);
        compiler.beginScope();
        
        if (f.args.type != FormType.NIL) {
            Symbol sym = cast(Symbol)f.args.head;
            int constant = compiler.parseVariable(sym);
            compiler.defineVariable(constant);

            foreach (Form inner ; f.args.tail) {
                sym = cast(Symbol)inner;
                constant = compiler.parseVariable(sym);
                compiler.defineVariable(constant);
            }
        }

        advance();  // ??

        Block b = new Block(f.line);
        b.blockBody = f.funcBody;
        compiler.compileBlock(b);

        // write the function as a Value
        Function outFunc = compiler.finish();
        this.func.chunk.writeOp(OpCode.OP_CONSTANT, f.line);
        int funcAddr = this.func.chunk.addConstant(makeObjValue(outFunc));
        this.func.chunk.writeOp(to!ubyte(funcAddr), f.line);

        // capture the errors
        foreach (CompileError err ; compiler.errors) {
            this.errors ~= err;
        }

        // define the global variable
        this.defineVariable(global);
    }

    ValueType compileCons(Form form, ValueType expected) {
        Cons cons = cast(Cons)form;
        Form head = cons.head;

        if (head.type == FormType.LAMBDA) {
            this.resolve(head);
            this.call(cons.tail);
            this.advance();
            return ValueType.NIL;
        }

        if (head.type != FormType.SYMBOL) {
            this.error("can't evaluate without a symbol or lambda", head.line);
            this.advance();
            return ValueType.NIL;
        }

        Symbol sym = cast(Symbol)head;
        switch (sym.name) {
            case "+":
                this.compileAdd(cons.tail, ValueType.NUMBER);
                break;
            case "-":
                if (cons.tail.length == 1) {
                    this.compileNegate(cons.tail[0]);
                } else {
                    return this.compileSubtract(cons.tail, ValueType.NUMBER);
                }
                break;
            case "<":
                return this.compileLess(cons.tail, ValueType.NUMBER);
                break;
            case "<=":
                ValueType vt = this.compileGreater(cons.tail);
                this.func.chunk.writeOp(OpCode.OP_NOT, cons.line);
                return vt;
            case ">":
                return this.compileGreater(cons.tail);
                break;
            case ">=":
                ValueType vt = this.compileLess(cons.tail, ValueType.NUMBER);
                this.func.chunk.writeOp(OpCode.OP_NOT, cons.line);
                return vt;
            case "not":
                return this.compileNot(cons.tail);

            // STRINGS
            case "concat":
                return this.compileConcat(cons.tail);

            // LISTS
            case "append":
                this.compileAppend(cons.tail);
                break;
            case "first":
                this.compileFirst(cons.tail);
                break;
            case "length":
                this.compileLength(cons.tail);
                break;
            case "list":
                //return compileList(cons.tail);
                this.compileList(cons.tail);
                break;
            case "rest":
                this.compileRest(cons.tail);
                break;
            default:
                /*
                writefln("unsure how to compile %s", sym.name);
                advance();
                */
                this.resolve(head);
                this.call(cons.tail);
                this.advance();
                break;
        }
        return ValueType.NIL;

    }

    ValueType resolve(Form form, const ValueType expecting = ValueType.ANY) {
        //printForm(form);
        this.currentLine = form.line;
        switch(form.type) {
            case FormType.ATOM:
                this.compileAtom(form, expecting);
                break;
            case FormType.LITERALSTRING:
                this.compileString(form, expecting);
                break;
            case FormType.CONS:
                return this.compileCons(form, expecting);
                break;
            case FormType.NIL:
                this.compileNil(form);
                break;
            case FormType.DEF:
                this.compileDef(form, expecting);
                break;
            case FormType.SYMBOL:
                this.compileSymbol(form, expecting);
                break;
            case FormType.BLOCK:
                this.compileBlock(form);
                break;
            case FormType.IF:
                this.compileIf(form);
                break;
            case FormType.AND:
                this.compileAnd(form);
                break;
            case FormType.OR:
                this.compileOr(form);
                break;
            case FormType.FUNC:
                this.compileFunc(form);
                break;
            case FormType.LAMBDA:
                this.compileLambda(form);
                break;
            default:
                write("not sure how to resolve: ");
                printForm(form);
                this.advance();
                break;
        }
        return ValueType.NIL;
    }

    Function compile() {
        this.advance();
        while(current.type != FormType.EOF) {
            this.resolve(current);
        }
        
        return this.finish();
    }

    Function finish() {
        this.func.chunk.writeOp(OpCode.OP_RETURN, this.current.line);
        if (DEBUG) {
            disassembleChunk(this.func.chunk, to!string(func));
        }
        return this.func;
    }

    this(ObjType type, Parser* parser, string name = "") {
        this.parser = parser;
        this.func = new Function(type, name);
        //localCount = 0;

        this.locals ~= Local(new Symbol("", -1), 0);
        this.localCount++;
    }
}