module Parse

open FParsec.Primitives
open FParsec.CharParsers
open FParsec.OperatorPrecedenceParser

let commentLine = parse {
  do! skipString "//"
  do! skipManyCharsTill anyChar newline }

let commentBlock = parse {
  do! skipString "/*"
  do! skipManyCharsTill anyChar (skipString "*/") }

// white space or comment
let ws = (many (choice [spaces1; commentLine; commentBlock]) |>> ignore) <?> "whitespace"

let ch c = skipChar c >>. ws
let str s = pstring s .>> ws

// An ident must start by a letter or underscore
let ident =
  let nonDigit = asciiLetter <|> pchar '_'
  let p = pipe2 nonDigit (manyChars (nonDigit <|> digit)) (fun c s -> c.ToString() + s)
  (p .>> ws) <?> "identifier"

// The operator parser from FParsec takes care of precedence issues
let opp = new OperatorPrecedenceParser<_,_>()

// exprNoComma is very important in GLSL and C parsers
// "int a=2,b;"   is different of   "int a=(2,b);"
let exprNoComma = opp.ExpressionParser
let expr = sepBy1 exprNoComma (ch ',') |>> (List.reduce (fun acc e -> Ast.FunCall(Ast.Var ",", [acc;e])))
let parenExp = between (ch '(') (ch ')') expr

// Primitives

// Parse a number (pfloat from FParsec doesn't handle every GLSL float)
let number =
  let r = @"(\d+\.?\d*|\.\d+)([eE][-+]?[0-9]+)?"
  let conv (s: string) =
    let ok, res = System.Int32.TryParse(s)
    if ok then Ast.Int res
    else Ast.Float (float s)
  regex r .>> ws |>> conv

let prim = parenExp <|> (ident |>> Ast.Var) <|> number

// Very high priority (parenthesis, function call, field access)
let argList = sepBy exprNoComma (ch ',')
let fcall = between (ch '(') (ch ')') argList |>>
            (fun args fct -> Ast.FunCall(fct, args))
let subscript = between (ch '[') (ch ']') expr |>>
                (fun ind arr -> Ast.Subscript(arr, ind))
let dot = ch '.' >>. ident |>> (fun field r -> Ast.Dot(r, field))
let post = dot <|> subscript <|> fcall

let simpleExpr = pipe2 prim (many post)
                  (fun prim posts -> List.fold (fun acc elt -> elt acc) prim posts)
opp.TermParser <- simpleExpr


// Operators

let precedence1 = [
    ["*"; "/"; "%"], Assoc.Left
    ["+"; "-"], Assoc.Left
    ["<<"; ">>"], Assoc.Left
    ["<"; ">"; "<="; ">="], Assoc.Left
    ["=="; "!="], Assoc.Left
    ["&"], Assoc.Left
    ["^"], Assoc.Left
    ["|"], Assoc.Left
    ["&&"], Assoc.Left
    ["^^"], Assoc.Left
    ["||"], Assoc.Left
]
// precedence of ?: is between precedence1 and precedence2
let precedence2 = [
    ["="; "+="; "-="; "*="; "/="; "%="; "<<="; ">>="; "&="; "^="; "|="], Assoc.Right
]

// Add all the operators in the OperatorParser
let makeOperator =
  // we start with operators with highest priority, then we decrement the counter.
  let precCounter = ref 20 //(we have at most 20 different priorities)
  let addInfix li =
    for ops, assoc in li do
      decr precCounter
      for op in ops do
        opp.AddOperator(InfixOp(op, ws, !precCounter, assoc, fun x y -> Ast.FunCall(Ast.Var op, [x; y])))

  let addPrefix() =
    decr precCounter
    for op in ["++"; "--"; "+"; "-"; "~"; "!"] do
      opp.AddOperator(PrefixOp(op, ws, !precCounter, true, fun x -> Ast.FunCall(Ast.Var op, [x])))

  let addPostfix() =
    decr precCounter
    for op in ["++"; "--"] do
      // the '$' prefix in the name is a trick to distinguish ++a from a++ in the AST
      opp.AddOperator(PostfixOp(op, ws, !precCounter, true, fun x -> Ast.FunCall(Ast.Var ("$"+op), [x])))

  addPostfix()
  addPrefix()
  addInfix precedence1
  decr precCounter
  opp.AddOperator(TernaryOp("?", ws, ":", ws, !precCounter, Assoc.Right, fun x y z -> Ast.FunCall(Ast.Var "?:", [x; y; z])))
  addInfix precedence2

// a statement made of a single expression
let simpleStatement = opt expr |>> (function Some exp -> Ast.Expr exp | None -> Ast.Block [])
let statement, stmtRef = createParserForwardedToRef()

let keyword s = attempt (pstring s .>> notFollowedBy letter .>> notFollowedBy digit) .>> ws

let typeQualifier = ["const"; "attribute"; "varying"; "uniform"]
                  |> List.map keyword |> choice

let paramQualifier = ["in"; "out"; "inout"]
                  |> List.map keyword |> choice

// eg. "const out int", "uniform float"
let specifiedType = pipe3 (opt typeQualifier) (opt paramQualifier) ident (fun tyQ paQ name -> (name, tyQ, paQ) : Ast.Type)

// eg. "int foo[] = exp, bar = 3"
let declaration =
  let bracket = between (ch '[') (ch ']') (opt expr) |>> (fun size -> defaultArg size (Ast.Int 0))
  let init = ch '=' >>. exprNoComma
  let var = tuple3 ident (opt bracket) (opt init)
  let list = sepBy1 var (ch ',')
  tuple2 specifiedType list

// eg. int foo[]   used for function arguments
let singleDeclaration =
  let bracket = between (ch '[') (ch ']') (opt expr) |>> (fun size -> defaultArg size (Ast.Int 0))
  pipe3 specifiedType ident (opt bracket) (fun ty id brack -> (ty, [id, brack, None]):Ast.Decl)

let forLoop =
  let init1 = declaration |>> (fun decl e2 e3 body -> Ast.ForD(decl, e2, e3, body))
  let init2 = opt expr |>> (fun e1 e2 e3 body -> Ast.ForE(e1, e2, e3, body))
  let init = attempt init1 <|> init2 .>> ch ';'
  let cond = opt expr .>> ch ';'
  let inc = opt expr .>> ch ')'
  pipe4 (keyword "for" >>. ch '(' >>. init) cond inc statement (fun f e2 e3 body -> f e2 e3 body)

let whileLoop =
    pipe2 (keyword "while" >>. parenExp) statement
      (fun cond stmt -> Ast.While(cond, stmt))

let doWhileLoop =
    pipe2 (keyword "do" >>. statement) (str "while" >>. parenExp)
      (fun stmt cond -> Ast.DoWhile(cond, stmt))

let ifStatement =
    pipe3 (keyword "if" >>. parenExp) statement (opt (str "else" >>. statement))
      (fun cond stmt1 stmt2 -> Ast.If(cond, stmt1, stmt2))

let block =
  let list = many statement |>> Ast.Block
  between (ch '{') (ch '}') list

let macro =
  let line = manyCharsTill anyChar newline
  pchar '#' >>. line .>> ws |>> Ast.Macro

let special =
  let key =
   choice [
    keyword "break"
    keyword "continue"
    keyword "discard"
   ] |>> (fun k -> Ast.Keyword(k, None))

  let ret = pipe2 (keyword "return") (opt expr) (fun k e -> Ast.Keyword(k, e))
  (key <|> ret) .>> ch ';'

// A statement
stmtRef := choice [
  block
  special
  forLoop
  ifStatement
  whileLoop
  doWhileLoop
  attempt ((declaration .>> ch ';') |>> Ast.Decl)
  simpleStatement .>> ch ';']

// eg. "int foo(float a[], out int b)"
let functionHeader =
  let header = tuple2 specifiedType (ident .>> (ch '('))
  let argList = sepBy singleDeclaration (ch ',')
  pipe2 header (argList .>> ch ')') (fun (ty, id) args -> ty, id, args)

let pfunction =
  pipe2 functionHeader block (fun head body -> Ast.Function(head, body))

let precord =
  keyword "struct" |>> (fun x -> failwith "User-defined records are not yet supported.")

let toplevel =
  let decl = declaration .>> ch ';'
  let item = choice [
              macro
              precord
              attempt decl |>> Ast.TLDecl
              pfunction
  ]
  many item

let parse = ws >>. toplevel .>> eof


// Test

let runParser file =
  let res = runParserOnFile parse () file System.Text.Encoding.Default
  match res with
  | Success(r,_,_) -> printfn "%A" r
  | Failure(str, exn, _) -> printfn "Failure: %s" str

let () =
  // arguments list
  let args = System.Environment.GetCommandLineArgs().[1..]

  if Array.length args = 0 then
      printfn "Usage"
      printfn "  %s [shader_file]" (System.IO.Path.GetFileName (System.Environment.GetCommandLineArgs().[0]))
  else
      Array.iter runParser args

  System.Console.ReadLine() |> ignore