package io.trygvis.ld;

import io.trygvis.ld.antlr.GnuLdBaseVisitor;
import io.trygvis.ld.antlr.GnuLdLexer;
import io.trygvis.ld.antlr.GnuLdParser;
import org.antlr.v4.runtime.ANTLRFileStream;
import org.antlr.v4.runtime.BufferedTokenStream;
import org.antlr.v4.runtime.tree.ParseTreeProperty;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class LdScript {

    public enum MemoryAttribute {
        R, W, X, A, I,
    }

    public static class MemoryArea {
        private final BigInteger MEGA = BigInteger.valueOf(1024 * 1024);
        private final BigInteger KILO = BigInteger.valueOf(1024);

        String name;
        BigInteger origin;
        BigInteger length;
        Set<MemoryAttribute> attributes;

        @Override
        public String toString() {
            return "MemoryArea{" +
                    "name='" + name + '\'' +
                    ", origin=" + origin +
                    ", length=" + length +
                    ", attributes='" + attributes + '\'' +
                    '}';
        }

        public String prettyOrigin() {
            if (origin == null) {
                return "";
            }
            String hex = origin.toString(16);

            return "0x" + ("00000000").substring(0, 8 - hex.length()) + hex;
        }

        public String prettyLength() {
            if (length == null) {
                return "";
            }

            if (length.compareTo(MEGA) >= 0 && length.mod(MEGA).equals(BigInteger.ZERO)) {
                return length.divide(MEGA) + "M";
            }

            if (length.compareTo(KILO) >= 0 && length.mod(KILO).equals(BigInteger.ZERO)) {
                return length.divide(KILO) + "k";
            }

            return String.valueOf(length);
        }

        public String prettyAttributes() {
            String s = "";
            for (MemoryAttribute attribute : attributes) {
                s += attribute.name().toLowerCase();
            }
            return s;
        }
    }

    public static class Section {
        String name;
    }

    public static void main(String[] args) throws Exception {
        try {
            work(args);
        } catch (Exception e) {
            e.printStackTrace(System.out);
        }
    }

    private static void work(String[] args) throws Exception {
        args = new String[]{"lds/d2000.ld"};

        GnuLdLexer lexer = new GnuLdLexer(new ANTLRFileStream(args[0]));
        GnuLdParser parser = new GnuLdParser(new BufferedTokenStream(lexer));

//        ElfInfoGnuLdBaseListener listener = new ElfInfoGnuLdBaseListener();
//        parser.addParseListener(listener);
//
//        parser.file();
//
//        for (Section section : listener.sections) {
//            System.out.println("section.name = " + section.name);
//        }

        ElfinfoGnuLdBaseVisitor visitor = new ElfinfoGnuLdBaseVisitor();
        visitor.visit(parser.file());

        System.out.println("--------------------------------------------------------");
        for (MemoryArea area : visitor.memoryAreas) {
            System.out.println("  " + area.name + "(" + area.prettyAttributes() + ") : ORIGIN = " + area.prettyOrigin() + ", LENGTH=" + area.prettyLength());
        }
    }

    private static class ElfinfoExprGnuLdBaseVisitor extends GnuLdBaseVisitor<ElfinfoExprGnuLdBaseVisitor> {
        private BigInteger value;
        private ParseTreeProperty<BigInteger> es = new ParseTreeProperty<>();

        public static BigInteger evaluate(GnuLdParser.Mustbe_expContext ctx) {
            ElfinfoExprGnuLdBaseVisitor v = new ElfinfoExprGnuLdBaseVisitor();
            v.visitMustbe_exp(ctx);
            return v.value();
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLt(GnuLdParser.ExpLtContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpBlock(GnuLdParser.ExpBlockContext ctx) {
            return visitChildren(ctx);

        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAbsolute(GnuLdParser.ExpAbsoluteContext ctx) {
            return visitChildren(ctx);

        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLog2ceil(GnuLdParser.ExpLog2ceilContext ctx) {
            return visitChildren(ctx);

        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAddr(GnuLdParser.ExpAddrContext ctx) {
            return visitChildren(ctx);

        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpDataSegmentEnd(GnuLdParser.ExpDataSegmentEndContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpSub(GnuLdParser.ExpSubContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitExpSub");

            ElfinfoExprGnuLdBaseVisitor ret = visitChildren(ctx);
            BigInteger a = es.removeFrom(ctx.exp(0));
            BigInteger b = es.removeFrom(ctx.exp(1));

            BigInteger x = a.subtract(b);

            es.put(ctx, x);
            return ret;
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpDefined(GnuLdParser.ExpDefinedContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpMod(GnuLdParser.ExpModContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpInvert(GnuLdParser.ExpInvertContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAlign(GnuLdParser.ExpAlignContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpMul(GnuLdParser.ExpMulContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAnd(GnuLdParser.ExpAndContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpXor(GnuLdParser.ExpXorContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpParen(GnuLdParser.ExpParenContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpMinus(GnuLdParser.ExpMinusContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpDiv(GnuLdParser.ExpDivContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpGe(GnuLdParser.ExpGeContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpMin(GnuLdParser.ExpMinContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAlign2(GnuLdParser.ExpAlign2Context ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpNegate(GnuLdParser.ExpNegateContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpName(GnuLdParser.ExpNameContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpOr(GnuLdParser.ExpOrContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpTrinary(GnuLdParser.ExpTrinaryContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpOror(GnuLdParser.ExpOrorContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpDataSegmentAlign(GnuLdParser.ExpDataSegmentAlignContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLengthExp(GnuLdParser.ExpLengthExpContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAdd(GnuLdParser.ExpAddContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLoadaddr(GnuLdParser.ExpLoadaddrContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpGt(GnuLdParser.ExpGtContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpOrigin(GnuLdParser.ExpOriginContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpEq(GnuLdParser.ExpEqContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpMax(GnuLdParser.ExpMaxContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLshift(GnuLdParser.ExpLshiftContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpSizeofHeaders(GnuLdParser.ExpSizeofHeadersContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpLe(GnuLdParser.ExpLeContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpNe(GnuLdParser.ExpNeContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAndand(GnuLdParser.ExpAndandContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpRshift(GnuLdParser.ExpRshiftContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpNextParen(GnuLdParser.ExpNextParenContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAlignof(GnuLdParser.ExpAlignofContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpSegmentStart(GnuLdParser.ExpSegmentStartContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpInt(GnuLdParser.ExpIntContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitExpInt");

            String str = ctx.INT().getText().toLowerCase();
            int base = 10;
            if (str.startsWith("0x")) {
                base = 16;
                str = str.substring(2);
            }

            int factor = 1;
            if (str.endsWith("k")) {
                factor = 1024;
                str = str.substring(0, str.length() - 1);
            } else if (str.endsWith("k")) {
                factor = 1024 * 1024;
                str = str.substring(0, str.length() - 1);
            }

            BigInteger i = new BigInteger(str, base);
            if (factor > 1) {
                i = i.multiply(BigInteger.valueOf(factor));
            }
            System.out.println("ctx = " + ctx);
            System.out.println("i = " + i);
            es.put(ctx, i);

            return this;
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpConstant(GnuLdParser.ExpConstantContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpPlus(GnuLdParser.ExpPlusContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpSizeof(GnuLdParser.ExpSizeofContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpDataSegmentRelRoEnd(GnuLdParser.ExpDataSegmentRelRoEndContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitExpAssert(GnuLdParser.ExpAssertContext ctx) {
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoExprGnuLdBaseVisitor visitMustbe_exp(GnuLdParser.Mustbe_expContext ctx) {
            ElfinfoExprGnuLdBaseVisitor ret = visitChildren(ctx);

            value = es.removeFrom(ctx.exp());
            es.put(ctx, value);

            return ret;
        }

        public BigInteger value() {
            if (value != null) {
                return value;
            }

            throw new RuntimeException("Something bad happened, probably an unevaluated expression part.");
        }
    }

    private static class ElfinfoGnuLdBaseVisitor extends GnuLdBaseVisitor<ElfinfoGnuLdBaseVisitor> {
        public List<MemoryArea> memoryAreas = new ArrayList<>();

        public List<Section> sections = new ArrayList<>();

        ParseTreeProperty<BigInteger> es = new ParseTreeProperty<>();

        @Override
        public ElfinfoGnuLdBaseVisitor visitMemory_spec(GnuLdParser.Memory_specContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitMemory_spec");
            visitChildren(ctx);

            if (attributesInverted) {
                // not quite sure how this is supposed to work yet
                throw new RuntimeException("Inverted memory region attributes is not implemented yet.");
            }

            MemoryArea ma = new MemoryArea();
            ma.name = ctx.NAME().getText();
            ma.attributes = attributes;
            ma.origin = ElfinfoExprGnuLdBaseVisitor.evaluate(ctx.origin_spec().mustbe_exp());
            ma.length = ElfinfoExprGnuLdBaseVisitor.evaluate(ctx.length_spec().mustbe_exp());
            System.out.println(ma);
            memoryAreas.add(ma);
            return this;
        }

        MemoryAttribute attribute;
        boolean attributesInverted;
        Set<MemoryAttribute> attributes;

        @Override
        public ElfinfoGnuLdBaseVisitor visitAttributes_opt(GnuLdParser.Attributes_optContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitAttributes_opt");
            attributes = new HashSet<>();
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoGnuLdBaseVisitor visitAttributeNormal(GnuLdParser.AttributeNormalContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitAttributeNormal");
            String name = ctx.NAME().getText();
            System.out.println("ctx.ATTRIBUTE().getText() = " + name);
            for (int i = 0; i < name.length(); i++) {
                attribute = MemoryAttribute.valueOf(String.valueOf(Character.toUpperCase(name.charAt(i))));
                attributes.add(attribute);
            }
            attributesInverted = false;
            return visitChildren(ctx);
        }

        @Override
        public ElfinfoGnuLdBaseVisitor visitAttributeInverted(GnuLdParser.AttributeInvertedContext ctx) {
            System.out.println("ElfinfoGnuLdBaseVisitor.visitAttributeInverted");
            if (!attributes.isEmpty()) {
                throw new RuntimeException("Attributes for memory areas can only be attributesInverted (with '!') as the first character in a specification; foo(!rw), not foo(x!rw).");
            }
            String name = ctx.NAME().getText();
            System.out.println("ctx.ATTRIBUTE().getText() = " + name);

            attributesInverted = true;
            return visitChildren(ctx);
        }
    }
}