(* FSKalc, a simple calculator implemented as a tree-walker interpreter *)

open FParsec
open System

(* Lexer *)

let manyCharsBetween popen pclose pchar = popen >>? manyCharsTill pchar pclose
let anyStringBetween popen pclose = manyCharsBetween popen pclose anyChar

let pcomment: Parser<unit, unit> = (skipString "(*" |> anyStringBetween <| skipString "*)") |>> ignore
//let pwhitespace: Parser<unit, unit> = (pcomment <|> spaces) |> many1 |>> ignore
let pwhitespace = spaces

let numberFormat = 
    NumberLiteralOptions.AllowBinary |||
    NumberLiteralOptions.AllowExponent |||
    NumberLiteralOptions.AllowFraction |||
    NumberLiteralOptions.AllowFractionWOIntegerPart |||
    NumberLiteralOptions.AllowHexadecimal |||
    NumberLiteralOptions.AllowOctal
    
let pdouble = numberLiteral numberFormat "number" |>> (fun nl -> double nl.String)
let pnumber: Parser<double, unit> = pdouble .>> pwhitespace
let pidentifier: Parser<string, unit> =
    let isIdentChar c = isLetter c || c = '_'

    many1SatisfyL isIdentChar "identifier" .>> pwhitespace

type Symbol =
    | Caret
    | Minus
    | Plus
    | Slash
    | Star
    | SlashSlash
    | Percent

let pleftparen = skipString "(" .>> pwhitespace
let prightparen = skipString ")" .>> pwhitespace
let pequal = skipString "=" .>> pwhitespace
let pcomma = skipString "," .>> pwhitespace
let pbang = skipString "!" .>> pwhitespace
let ppower = stringReturn "^" <| Caret .>> pwhitespace
let pminus = stringReturn "-" <| Minus .>> pwhitespace
let pplus = stringReturn "+" <| Plus .>> pwhitespace
let pstar = stringReturn "*" <| Star .>> pwhitespace
let pslashslash = skipString "/" >>? (stringReturn "/" <| SlashSlash .>> pwhitespace)
let pslash = stringReturn "/" <| Slash .>> pwhitespace
let ppercent = stringReturn "%" <| Percent .>> pwhitespace


(* Parser 

    - expressions prefixed by 'e'
    - statements prefixed by 's'

*)

// EXPRESSIONS

type Expression =
    | Number of double
    | VarGet of string
    | Grouping of Expression
    | FuncCall of Expression * list<Expression>
    | Power of Expression * list<Expression>
    | Negate of Expression
    | Factor of Expression * list<Symbol * Expression>
    | Term of Expression * list<Symbol * Expression>

let expr, exprref = createParserForwardedToRef()

// literals and function call

let enumber = pnumber |>> Number
let evarget = pidentifier |>> VarGet

// parentheses

let egrouping = pleftparen >>. (expr) .>> prightparen |>> Grouping

let efunccall = ((pidentifier |>> VarGet) .>>? pbang) .>>. (sepBy expr pcomma) |>> FuncCall

let primary = efunccall <|> enumber <|> evarget <|> egrouping

// power

let epower = choice [
    primary .>>.? many1 (ppower >>. primary) |>> Power
    primary
]

// unary

let enegate = choice [
    pminus >>? epower |>> Negate
    epower
]

// factor

let pfactorop = pstar <|> pslashslash <|> pslash <|> ppercent

let pimplicitmul: Parser<Symbol, unit> = stringReturn "" <| Star

let efactor = choice [
    enegate .>>.? many1 ((pfactorop .>>. enegate) <|> (pimplicitmul .>>. epower)) |>> Factor
    enegate
]


let ptermop = pplus <|> pminus

let eterm = choice [
    efactor .>>.? many1 (ptermop .>>. efactor) |>> Term
    efactor
]

do exprref := eterm

// STATEMENTS

type Statement =
    | ExprStatement of Expression
    | VarSet of string * Expression
    | FuncDef of string * List<string> * Expression

let sexprstmt = expr |>> ExprStatement

let sfuncdef = pidentifier .>>.? (((sepBy1 pidentifier pcomma)) .>>.? (pequal >>. expr))  |>> (fun (a, (b, c)) -> FuncDef (a, b, c))

let svarset = pidentifier .>>.? (pequal >>. expr) |>> VarSet

let stmt = choice [
    sfuncdef // first, so it falls through if multiple args not provided
    svarset // second so it falls through if no =
    sexprstmt
]

// REPL oriented => only 1 statement at once
let program = pwhitespace >>. stmt .>> eof

type Value =
    | Number of double
    | Native of string
    | Function of List<string> * Expression
    | Nil
    | Fail of string



// Store global variables
let mutable globals = new Collections.Generic.Dictionary<string, Value>()

// natives
let mutable natives = new Collections.Generic.Dictionary<string, (List<Value> -> Value)>()

// Natives

let wrap1ArgNative name func =
    let native (args: List<Value>) =
        match args.Length with
        | 1 ->
            let arg = args[0]
            match arg with
            | Number x -> Number (func x)
            | Fail _ -> arg
            | _ -> Fail $"Invalid argument to '{name}'."
        | _ -> Fail $"Invalid argument count to '{name}', expect 1, got {args.Length}."
    globals[name] <- Native name
    natives[name] <- native

wrap1ArgNative "sin" Math.Sin
wrap1ArgNative "cos" Math.Cos
wrap1ArgNative "tan" Math.Tan
wrap1ArgNative "arcsin" Math.Asin
wrap1ArgNative "arccos" Math.Acos
wrap1ArgNative "arctan" Math.Atan
wrap1ArgNative "ln" Math.Log
wrap1ArgNative "log" Math.Log10
wrap1ArgNative "exp" Math.Exp
wrap1ArgNative "sqrt" Math.Sqrt

let wrap2ArgNative name func =
    let native (args: List<Value>) =
        match args.Length with
        | 2 ->
            let arg = args[0]
            let arg2 = args[1]
            match arg, arg2 with
            | Number x, Number y -> Number (func x y)
            | Fail _, _ -> arg
            | _, Fail _ -> arg2
            | _ -> Fail $"Invalid argument to '{name}'."
        | _ -> Fail $"Invalid argument count to '{name}', expect 2, got {args.Length}."
    globals[name] <- Native name
    natives[name] <- native

let logx logBase logArg =
    Math.Log2(logArg) / Math.Log2(logBase)

let rtx rootBase rootArg =
    rootArg ** (1.0/rootBase)

wrap2ArgNative "logx" logx
wrap2ArgNative "rtx" rtx

// native values

globals["pi"] <- Number Math.PI
globals["tau"] <- Number Math.Tau
globals["e"] <- Number Math.E
globals["inf"] <- Number infinity
globals["nan"] <- Number nan
globals["ans"] <- Number 0

// exec

let varGet key =
    let suc, value = globals.TryGetValue(key)
    match suc with
    | true -> value
    | false -> $"Attempt to get the value of undefined global variable {key}." |> Fail

let applyOp op left right = 
    match op, left, right with
    | Plus, Number x, Number y -> Number (x + y)
    | Minus, Number x, Number y -> Number (x - y)
    | Star, Number x, Number y -> Number (x * y)
    | Slash, Number x, Number y -> Number (x / y)
    | Percent, Number x, Number y -> Number (x % y)
    | SlashSlash, Number x, Number y -> Number (floor (x / y))
    | Caret, Number x, Number y -> Number (x ** y)
    | _, Fail _, _ -> left
    | _, _, Fail _ -> right
    | _ -> Fail $"Invalid operation {op} on {left} and {right}."

let negate value =
    match value with
    | Number x -> Number -x
    | Fail _ -> value
    | _ -> Fail $"Attempt to negate invalid type {value}."

let rec execExpr expr =
    match expr with
    | Expression.Number num -> Number num
    | VarGet key -> varGet key
    | Factor (left, rights)
    | Term (left, rights) ->
        // left associative ops
        let mutable res = execExpr left
        for (op, right) in rights do
            res <- execExpr right |> applyOp op res
        res
    | Negate expr -> negate <| execExpr expr
    | Power (left, rights) ->
        // right associative ops
        // parser only creates a power node if rights contains at least 1 element
        let mutable res = execExpr rights[rights.Length - 1]
        let mutable x = rights.Length - 2
        while x >= 0 do
            res <- applyOp Caret (execExpr rights[x]) res
            x <- x - 1
        applyOp Caret (execExpr left) res
    | Grouping expr -> execExpr expr
    | FuncCall (expr, args) ->
        let func = execExpr expr
        let args = args |> List.map execExpr

        match func with
        | Function (parameters, f) ->
            let argLen = args.Length
            let paramLen = parameters.Length
            if paramLen <> argLen then
                Fail $"Invalid number of arguments, expected {paramLen}, got {argLen}"
            else
                List.iter2 (fun k v -> globals[k] <- v) parameters args
                execExpr f
        | Native nat ->
            natives[nat] args
        | Fail _ -> func
        | _ -> Fail $"Attempt to call an invalid type {func}"

let printValue value =
    match value with
    | Number x -> printfn "%A" x
    | _ -> printfn "%A" value

let execStmt stmt =
    match stmt with
    | ExprStatement expr -> 
        let res = execExpr expr
        match res with
        | Fail _ -> ()
        | _ -> globals["ans"] <- res
        printValue res
    | VarSet (name, expr) ->
        let res = execExpr expr
        match res with
        | Fail _ -> ()
        | _ -> globals[name] <- res
        printValue res
    | FuncDef (name, parameters, expr) -> globals[name] <- Function (parameters, expr)

printfn "FSKalc"
while true do
    let source = Console.ReadLine ()
    match run program source with
    | Success(result, _, _) ->
        execStmt result
    | Failure(errorMsg, _, _) -> printfn "Parser error: %s" errorMsg