// The Finch parser, written in Finch: tokens -> ast.Program.
import lexer
import ast

struct Parser {
    []lexer.Token toks
    int i
    int depth      // inside ( ) and [ ] newlines don't end anything
    int file
    str path
}

fn fail(ptr[Parser] p, lexer.Token t, str msg) {
    print(p.path + ":" + str(t.line) + ":" + str(t.col) + ": error: " + msg)
    exit(1)
}

fn cur(ptr[Parser] p) -> lexer.Token {
    return p.toks[p.i]
}

fn curKind(ptr[Parser] p) -> str {
    return p.toks[p.i].kind
}

fn peekKind(ptr[Parser] p, int off) -> str {
    k := p.i + off
    if k >= p.toks.len {
        k = p.toks.len - 1
    }
    return p.toks[k].kind
}

fn at(ptr[Parser] p, str k) -> bool {
    return p.toks[p.i].kind == k
}

fn next(ptr[Parser] p) -> lexer.Token {
    t := p.toks[p.i]
    p.i += 1
    return t
}

fn accept(ptr[Parser] p, str k) -> bool {
    if at(p, k) {
        p.i += 1
        return true
    }
    return false
}

fn describe(lexer.Token t) -> str {
    if t.kind == "ident" {
        return "'" + t.text + "'"
    }
    if t.kind == "eof" {
        return "end of file"
    }
    if t.kind == "int" || t.kind == "float" {
        return "a number"
    }
    if t.kind == "string" {
        return "a string"
    }
    if t.kind == "char" {
        return "a char"
    }
    return "'" + t.kind + "'"
}

fn expect(ptr[Parser] p, str k, str what) -> lexer.Token {
    if !at(p, k) {
        fail(p, cur(p), "expected " + what + ", found " + describe(cur(p)))
    }
    return next(p)
}

fn sameLine(ptr[Parser] p) -> bool {
    return p.depth > 0 || !p.toks[p.i].nl
}

fn node(ptr[Parser] p, str kind) -> ast.Node {
    t := p.toks[p.i]
    return ast.Node(kind: kind, line: t.line, col: t.col, file: p.file)
}

fn isTypeName(str n) -> bool {
    names := ["int", "float", "bool", "char", "str", "ptr", "i8", "i16", "i32", "i64", "u8", "u16", "u32", "u64", "f32", "f64"]
    return names.contains(n)
}

// int, str, ptr[T], []T, Name, module.Name
fn parseType(ptr[Parser] p) -> ast.Type {
    if accept(p, "[") {
        expect(p, "]", "']' (an array type looks like []int)")
        return ast.arrayOf(parseType(p))
    }
    if !at(p, "ident") {
        fail(p, cur(p), "unexpected " + describe(cur(p)) + " (expected a type like int, str, []int or a struct name)")
    }
    name := next(p).text
    if name == "map" && at(p, "[") {
        next(p)
        k := parseType(p)
        expect(p, "]", "']' (a map type looks like map[str]int)")
        return ast.mapOf(k, parseType(p))
    }
    if isTypeName(name) {
        if name == "ptr" && accept(p, "[") {
            t := ast.ptrTo(parseType(p))
            expect(p, "]", "']'")
            return t
        }
        if name == "i64" {
            return ast.simple("int")
        }
        if name == "f64" {
            return ast.simple("float")
        }
        return ast.simple(name)
    }
    t := ast.Type(kind: "named", name: name)
    if at(p, ".") && peekKind(p, 1) == "ident" {
        next(p)
        t.module = name
        t.name = next(p).text
    }
    // Stack[int], Pair[str, int]: the type arguments of a generic struct
    if at(p, "[") && !cur(p).nl && peekKind(p, 1) != "]" {
        next(p)
        while true {
            t.elem.push(parseType(p))
            if !accept(p, ",") {
                break
            }
        }
        expect(p, "]", "']' after the type arguments")
    }
    return t
}

// Stack[int] name: a generic type, then a name on the same line. (a[i] by itself is an index.)
fn genericDeclAhead(ptr[Parser] p) -> bool {
    if !at(p, "ident") || peekKind(p, 1) != "[" || peekKind(p, 2) == "]" {
        return false
    }
    depth := 0
    i := p.i + 1
    while i < p.toks.len {
        k := p.toks[i].kind
        if k == "[" {
            depth += 1
        } else if k == "]" {
            depth -= 1
            if depth == 0 {
                return i + 1 < p.toks.len && p.toks[i + 1].kind == "ident" && !p.toks[i + 1].nl
            }
        } else if k == "eof" || p.toks[i].nl && depth == 0 {
            return false
        }
        i += 1
    }
    return false
}

// The same, when the name was just taken and the parser stands on the `[`.
fn genericCallAheadFrom(ptr[Parser] p) -> bool {
    p.i -= 1
    r := genericCallAhead(p)
    p.i += 1
    return r
}

// Is the `[` after this name followed, once it closes, by `(`? Then it is  max[int](...)  and not an index.
fn genericCallAhead(ptr[Parser] p) -> bool {
    if peekKind(p, 1) != "[" || p.toks[p.i + 1].nl || peekKind(p, 2) == "]" {
        return false
    }
    depth := 0
    i := p.i + 1
    while i < p.toks.len {
        k := p.toks[i].kind
        if k == "[" {
            depth += 1
        } else if k == "]" {
            depth -= 1
            if depth == 0 {
                return i + 1 < p.toks.len && p.toks[i + 1].kind == "(" && !p.toks[i + 1].nl
            }
        } else if k == "eof" || p.toks[i].nl {
            return false
        }
        i += 1
    }
    return false
}

fn atDeclaration(ptr[Parser] p) -> bool {
    if genericDeclAhead(p) {
        return true
    }
    if at(p, "[") {
        return peekKind(p, 1) == "]"
    }
    if !at(p, "ident") {
        return false
    }
    n := peekKind(p, 1)
    if n == "ident" {
        return !p.toks[p.i + 1].nl
    }
    if (cur(p).text == "ptr" || cur(p).text == "map") && n == "[" {
        return true
    }
    return n == "." && peekKind(p, 2) == "ident" && peekKind(p, 3) == "ident" && !p.toks[p.i + 3].nl
}

fn endOfStatement(ptr[Parser] p) {
    if at(p, "}") || at(p, "eof") || cur(p).nl {
        return
    }
    fail(p, cur(p), "unexpected " + describe(cur(p)) + " (put each statement on its own line)")
}

// ---------- top level ----------

fn parseProgram(ptr[Parser] p) -> ast.Program {
    prog := ast.Program()
    while !at(p, "eof") {
        if at(p, "import") {
            t := next(p)
            if at(p, "string") {
                prog.cwhere.push(p.path + ":" + str(t.line) + ":" + str(t.col))
                prog.cimports.push(next(p).text)
            } else {
                if !at(p, "ident") {
                    fail(p, t, "expected a module name or a C header after 'import'")
                }
                prog.imports.push(next(p).text)
            }
        } else if at(p, "link") {
            next(p)
            prog.links.push(expect(p, "string", "a library or C file name, like  link \"glfw\"").text)
        } else if at(p, "struct") {
            prog.structs.push(parseStruct(p))
        } else if at(p, "enum") {
            prog.enums.push(parseEnum(p))
        } else {
            prog.fns.push(parseFn(p))
        }
    }
    return prog
}

// enum Color { Red, Green, Blue }      enum Status { Ok = 200, NotFound = 404 }
// The variants are numbered from 0 (or from the last number given) and are written Color.Red.
fn parseEnum(ptr[Parser] p) -> ast.Node {
    e := node(p, "enum")
    expect(p, "enum", "'enum'")
    e.name = expect(p, "ident", "an enum name").text
    expect(p, "{", "'{'")
    while !at(p, "}") {
        if at(p, "eof") {
            fail(p, cur(p), "an enum is never closed (missing '}')")
        }
        v := node(p, "variant")
        v.name = expect(p, "ident", "a variant name, like Red").text
        if accept(p, "=") {
            sign := ""
            if accept(p, "-") {
                sign = "-"
            }
            v.text = sign + expect(p, "int", "a whole number, like  Ok = 200").text
        }
        e.kids.push(v)
        if !accept(p, ",") {
            endOfStatement(p)
        }
    }
    next(p)
    if e.kids.len == 0 {
        fail(p, cur(p), "the enum " + e.name + " has no variants")
    }
    return e
}

// match value {  A, B -> ...   C -> ...   else -> ... }
// As a statement the arms are statements or { blocks }; as an expression they are expressions.
fn parseMatch(ptr[Parser] p, bool asExpr) -> ast.Node {
    m := node(p, "match")
    m.flag = asExpr
    expect(p, "match", "'match'")
    m.kids.push(parseExpr(p))
    expect(p, "{", "'{'")
    saved := p.depth
    p.depth = 0
    while !at(p, "}") {
        if at(p, "eof") {
            fail(p, cur(p), "a match is never closed (missing '}')")
        }
        arm := node(p, "arm")
        pats := node(p, "patterns")
        if accept(p, "else") {
            arm.flag = true
        } else {
            while true {
                pats.kids.push(parseExpr(p))
                if !accept(p, ",") {
                    break
                }
            }
        }
        arm.kids.push(pats)
        expect(p, "->", "'->' after the value to match, like   Red -> print(\"red\")")
        if at(p, "{") && !asExpr {
            arm.kids.push(parseBlock(p))
        } else if asExpr {
            arm.kids.push(parseExpr(p))
        } else {
            arm.kids.push(parseStatement(p))
        }
        m.kids.push(arm)
        if !accept(p, ",") {
            endOfStatement(p)
        }
    }
    p.depth = saved
    next(p)
    if m.kids.len == 1 {
        fail(p, cur(p), "a match needs at least one arm")
    }
    return m
}

// [T, U]
fn parseTParams(ptr[Parser] p) -> []str {
    []str names
    expect(p, "[", "'['")
    while true {
        names.push(expect(p, "ident", "a type parameter name, like T").text)
        if !accept(p, ",") {
            break
        }
    }
    expect(p, "]", "']' after the type parameters")
    return names
}

fn parseStruct(ptr[Parser] p) -> ast.Node {
    s := node(p, "struct")
    expect(p, "struct", "'struct'")
    s.name = expect(p, "ident", "a struct name").text
    if at(p, "[") {
        s.tparams = parseTParams(p)
    }
    expect(p, "{", "'{'")
    while !at(p, "}") {
        if at(p, "eof") {
            fail(p, cur(p), "a struct is never closed (missing '}')")
        }
        f := node(p, "field")
        if !atDeclaration(p) {
            fail(p, cur(p), "unexpected " + describe(cur(p)) + " (a field looks like: int x)")
        }
        f.ty = parseType(p)
        f.name = expect(p, "ident", "a field name").text
        if accept(p, "=") {
            f.kids.push(parseExpr(p))
        }
        s.kids.push(f)
        endOfStatement(p)
    }
    next(p)
    return s
}

fn parseFn(ptr[Parser] p) -> ast.Node {
    f := node(p, "fn")
    f.exported = accept(p, "export")
    expect(p, "fn", "'fn', 'struct' or 'import'")
    f.name = expect(p, "ident", "a function name").text
    if at(p, "[") {
        f.tparams = parseTParams(p)         // fn max[T](...)   or the struct's, in  fn Stack[T].push(...)
    }
    if accept(p, ".") {
        f.recv = f.name
        f.name = expect(p, "ident", "a method name").text
    }
    expect(p, "(", "'('")
    p.depth += 1
    if !at(p, ")") {
        while true {
            prm := node(p, "param")
            prm.ty = parseType(p)
            prm.name = expect(p, "ident", "a parameter name").text
            f.kids.push(prm)
            if !accept(p, ",") {
                break
            }
        }
    }
    p.depth -= 1
    expect(p, ")", "')'")
    f.ty = ast.simple("void")
    if accept(p, "->") {
        if !at(p, "!") {
            f.ty = parseType(p)
        }
        f.fallible = accept(p, "!")
        if !f.fallible && f.ty.kind == "void" {
            fail(p, cur(p), "unexpected " + describe(cur(p)) + " (after '->' comes a type, or ! for a function that can fail)")
        }
    }
    f.kids.push(parseBlock(p))
    return f
}

// ---------- statements ----------

fn parseBlock(ptr[Parser] p) -> ast.Node {
    blk := node(p, "block")
    expect(p, "{", "'{'")
    saved := p.depth
    p.depth = 0
    while !at(p, "}") {
        if at(p, "eof") {
            fail(p, cur(p), "a '{' is never closed (missing '}')")
        }
        blk.kids.push(parseStatement(p))
        endOfStatement(p)
    }
    p.depth = saved
    next(p)
    return blk
}

fn parseStatement(ptr[Parser] p) -> ast.Node {
    k := curKind(p)
    if k == "if" {
        return parseIf(p)
    }
    if k == "while" {
        s := node(p, "while")
        next(p)
        s.kids.push(parseExpr(p))
        s.kids.push(parseBlock(p))
        return s
    }
    if k == "for" {
        s := node(p, "for")
        next(p)
        s.name = expect(p, "ident", "a loop variable name").text
        if accept(p, ",") {
            s.text = expect(p, "ident", "a second loop variable name").text
        }
        expect(p, "in", "'in'")
        s.kids.push(parseExpr(p))
        if accept(p, "..") {
            if s.text.len > 0 {
                fail(p, cur(p), "a range of numbers has one loop variable: for i in 0..n")
            }
            s.kids.push(parseExpr(p))
        } else {
            s.kind = "each"
        }
        s.kids.push(parseBlock(p))
        return s
    }
    if k == "match" {
        return parseMatch(p, false)
    }
    if k == "return" {
        s := node(p, "return")
        next(p)
        if !at(p, "}") && !cur(p).nl {
            s.kids.push(parseExpr(p))
        }
        return s
    }
    if k == "break" || k == "continue" {
        s := node(p, k)
        next(p)
        return s
    }
    if k == "{" {
        return parseBlock(p)
    }
    if k == "defer" {
        s := node(p, "defer")
        next(p)
        if at(p, "defer") {
            fail(p, cur(p), "defer inside defer is not allowed")
        }
        s.kids.push(parseStatement(p))
        return s
    }
    if k == "fn" || k == "struct" || k == "import" || k == "link" {
        fail(p, cur(p), "'" + k + "' can't be used here in the self-hosted compiler")
    }

    if atDeclaration(p) {
        s := node(p, "decl")
        s.flag = true
        s.ty = parseType(p)
        s.name = expect(p, "ident", "a variable name").text
        if accept(p, "=") {
            s.kids.push(parseExpr(p))
        }
        return s
    }
    if at(p, "ident") && peekKind(p, 1) == ":=" {
        s := node(p, "decl")
        s.name = next(p).text
        next(p)
        s.kids.push(parseExpr(p))
        return s
    }

    start := cur(p)
    e := parseExpr(p)
    ops := ["=", "+=", "-=", "*=", "/=", "%="]
    if ops.contains(curKind(p)) && sameLine(p) {
        if e.kind != "var" && e.kind != "member" && e.kind != "index" {
            fail(p, start, "only a variable, a field, an element or p.value can be assigned to")
        }
        s := ast.Node(kind: "assign", line: start.line, col: start.col, file: p.file)
        s.name = next(p).kind
        s.kids.push(e)
        s.kids.push(parseExpr(p))
        return s
    }
    if e.kind != "call" && e.kind != "method" && e.kind != "try" && e.kind != "or" {
        fail(p, start, "this value is computed but never used")
    }
    s := ast.Node(kind: "expr", line: start.line, col: start.col, file: p.file)
    s.kids.push(e)
    return s
}

fn parseIf(ptr[Parser] p) -> ast.Node {
    s := node(p, "if")
    expect(p, "if", "'if'")
    s.kids.push(parseExpr(p))
    s.kids.push(parseBlock(p))
    if accept(p, "else") {
        if at(p, "if") {
            s.kids.push(parseIf(p))
        } else {
            s.kids.push(parseBlock(p))
        }
    }
    return s
}

// ---------- expressions ----------

fn binaryLevel(str op) -> int {
    if op == "||" {
        return 1
    }
    if op == "&&" {
        return 2
    }
    if op == "==" || op == "!=" || op == "<" || op == "<=" || op == ">" || op == ">=" {
        return 3
    }
    if op == "+" || op == "-" || op == "|" || op == "^" {
        return 4
    }
    if op == "*" || op == "/" || op == "%" || op == "<<" || op == ">>" || op == "&" {
        return 5
    }
    return 0
}

// f(x) or fallback  /  f(x) or { ... }: the lowest precedence, right to left
fn parseExpr(ptr[Parser] p) -> ast.Node {
    l := parseBinary(p, 1)
    if !sameLine(p) || !at(p, "or") {
        return l
    }
    e := node(p, "or")
    next(p)
    e.kids.push(l)
    if at(p, "{") {
        e.flag = true
        e.kids.push(parseBlock(p))
    } else {
        e.kids.push(parseExpr(p))
    }
    return e
}

// precedence climbing; all binary operators are left-associative
fn parseBinary(ptr[Parser] p, int level) -> ast.Node {
    if level > 5 {
        return parseUnary(p)
    }
    left := parseBinary(p, level + 1)
    while sameLine(p) && binaryLevel(curKind(p)) == level {
        e := node(p, "binary")
        e.name = next(p).kind
        e.kids.push(left)
        e.kids.push(parseBinary(p, level + 1))
        left = e
    }
    return left
}

fn parseUnary(ptr[Parser] p) -> ast.Node {
    if at(p, "try") {
        e := node(p, "try")
        next(p)
        e.kids.push(parsePostfix(p))
        return e
    }
    if at(p, "-") || at(p, "!") || at(p, "~") {
        e := node(p, "unary")
        e.name = next(p).kind
        e.kids.push(parseUnary(p))
        return e
    }
    return parsePostfix(p)
}

fn parseArgs(ptr[Parser] p, ptr[ast.Node] call) {
    expect(p, "(", "'('")
    p.depth += 1
    if !at(p, ")") {
        while !at(p, ")") {
            name := ""
            if at(p, "ident") && peekKind(p, 1) == ":" {
                name = next(p).text
                next(p)
            }
            call.names.push(name)
            call.kids.push(parseExpr(p))
            if !accept(p, ",") {
                break
            }
        }
    }
    p.depth -= 1
    expect(p, ")", "')'")
}

fn parsePostfix(ptr[Parser] p) -> ast.Node {
    e := parsePrimary(p)
    while sameLine(p) {
        if at(p, ".") {
            m := node(p, "member")
            next(p)
            m.name = expect(p, "ident", "a name after '.'").text
            m.kids.push(e)
            if at(p, "(") && !cur(p).nl {
                m.kind = "method"
                parseArgs(p, addr(m))
            }
            e = m
        } else if at(p, "[") {
            ix := node(p, "index")
            next(p)
            p.depth += 1
            ix.kids.push(e)
            ix.kids.push(parseExpr(p))
            p.depth -= 1
            expect(p, "]", "']'")
            e = ix
        } else {
            return e
        }
    }
    return e
}

fn parsePrimary(ptr[Parser] p) -> ast.Node {
    t := cur(p)
    e := node(p, "")
    k := t.kind
    if k == "int" || k == "float" || k == "char" {
        next(p)
        e.kind = k
        e.text = t.text
        return e
    }
    if k == "string" {
        next(p)
        e.kind = "str"
        e.text = t.text
        return e
    }
    if k == "true" || k == "false" {
        next(p)
        e.kind = "bool"
        e.text = k
        return e
    }
    if k == "null" {
        next(p)
        e.kind = "null"
        return e
    }
    if k == "(" {
        next(p)
        p.depth += 1
        inner := parseExpr(p)
        p.depth -= 1
        expect(p, ")", "')'")
        return inner
    }
    if k == "match" {
        return parseMatch(p, true)
    }
    if k == "[" {
        next(p)
        e.kind = "array"
        p.depth += 1
        if at(p, ":") && peekKind(p, 1) == "]" {       // [:], an empty map
            next(p)
            e.kind = "map"
        }
        while !at(p, "]") {
            e.kids.push(parseExpr(p))
            if e.kids.len == 1 && at(p, ":") {          // ["a": 1, ...]
                e.kind = "map"
            }
            if e.kind == "map" {
                expect(p, ":", "':' (a map entry looks like \"key\": value)")
                e.kids.push(parseExpr(p))
            }
            if !accept(p, ",") {
                break
            }
        }
        p.depth -= 1
        expect(p, "]", "']'")
        return e
    }
    if k == "ident" {
        next(p)
        e.name = t.text
        if at(p, "[") && p.toks[p.i - 1].kind == "ident" && genericCallAheadFrom(p) {
            // max[int](3, 4), Stack[int]()
            next(p)
            e.ty = ast.Type(kind: "targs")
            while true {
                e.ty.elem.push(parseType(p))
                if !accept(p, ",") {
                    break
                }
            }
            expect(p, "]", "']' after the type arguments")
            e.kind = "call"
            parseArgs(p, addr(e))
            return e
        }
        if at(p, "(") && !cur(p).nl {
            e.kind = "call"
            parseArgs(p, addr(e))
            return e
        }
        e.kind = "var"
        return e
    }
    fail(p, t, "unexpected " + describe(t) + " (expected a value)")
    return e
}

fn parse([]lexer.Token toks, int file, str path) -> ast.Program {
    parser := Parser(toks: toks, i: 0, depth: 0, file: file, path: path)
    return parseProgram(addr(parser))
}