//
//  main.swift
//  TokenizerWithRegex
//
//  Created by Guy Lapalme on 2024-08-02.
//  adapted from : https://docs.python.org/3/library/re.html#writing-a-tokenizer

import Foundation
import RegexBuilder

enum kinds {
    case NUMBER (Double)
    case ASSIGN
    case END
    case ID (Substring)
    case KEYWORD (Substring)
    case OP (Substring)
    case SKIP
    case NEWLINE
    case MISMATCH
}

struct Token:CustomStringConvertible {
    let kind: kinds
    let line,column: Int
    public var description:String {
        "\(kind) [\(line),\(column)]"
    }
}

// find the position as an integer of a substring within its original/base string
// columns are numbered from 1
func colPos(_ s:Substring) -> Int {
    let base = s.base
    return base.distance(from: base.startIndex, to: s.startIndex)+1
}

// according to https://docs.swift.org/swift-book/documentation/the-swift-programming-language/closures/#Escaping-Closures
// A closure is said to escape a function when the closure is passed as an argument
// to the function, but is called after the function returns.
func pat(_ regex:Regex<Substring>, kindClosure: @escaping (Substring)->kinds)
           -> Capture<(Substring,Token)>{
    return Capture {regex} transform:{
        Token(kind: kindClosure($0), line:lineNumber, column: colPos($0))
    }
}

let keywords:Set = ["IF", "THEN", "ENDIF", "FOR", "NEXT", "GOSUB", "RETURN"]
var lineNumber = 0

// tokenSpecifications: ChoiceOf<(Substring, Token?, Token?, Token?, Token?, Token?, Token?, Token?, Token?)>
let tokenSpecifications =  ChoiceOf {
    pat(/\d+(?:\.\d*)?/){s in .NUMBER(Double(s)!)}  //  .1  number
    pat(/:=/)           {_ in .ASSIGN}              //  .2  assignment
    pat(/;/)            {_ in .END}                 //  .3  end of statement
    pat(/[A-Za-z]\w*/)  {s in                       //  .4  identifier or keyword
                           return keywords.contains(String(s)) ? .KEYWORD(s) : .ID(s)}
    pat(/[+\-*\/]/)     {s in .OP(s)}               //  .5  arithmetic operator
    // ignored values
    pat(/[ \t]+/)       {_ in .SKIP}                //  .6  spaces
    pat(/$/)            {_ in .NEWLINE}             //  .7  end of line
    pat(/./)            {_ in .MISMATCH}            //  .8  error
}

struct Tokenizer:IteratorProtocol {
    let program:String
    var lines:[(Int,Substring)]
    var line:Substring
    
    init(_ program: String){
        self.program = program
        self.lines = Array(zip(1...,program.split(separator:"\n")))
        (lineNumber,line) = Tokenizer.nextLine(&self.lines)
    }
    
    mutating func next()->Token? {
        while line.isEmpty {
            if !lines.isEmpty {
                (lineNumber,line) = Tokenizer.nextLine(&self.lines)
            } else {
                return nil
            }
        }
        if let m = line.prefixMatch(of: tokenSpecifications) {
            let out = m.output
            line = line.dropFirst(out.0.count) // remove matched substring
            if let _ = out.8 { // error message for mismatch and search next token
                print("\(out.0) unexpected at line \(lineNumber), column: \(colPos(out.0))")
                return next()
            }
            if let _ = out.6 ?? out.7 { // ignore newline, skip
                return next()
            }
            if let t = out.1 ?? out.2 ?? out.3 ?? out.4 ?? out.5 ?? out.6 {
                return t
            }
        }
        return nil // should never happen
    }

    static func nextLine( _ lines:inout [(Int,Substring)])->(Int,Substring) {
        // create new string for setting column numbers relative to this line
        let (number,substr) = lines.removeFirst()
        let str = String(substr)
        return (number, str[str.startIndex..<str.endIndex])
    }
}


let statements = #"""
IF quantity THEN
    total := total + price * quantity;
    tax := price * 0.05; #
ENDIF;
"""#

func runTokenizer (_ program:String = statements){
    var tokenizer = Tokenizer(program)
    while let t=tokenizer.next() {
        print(t)
    }
}

/*   output
 KEYWORD("IF") [1,1]
 ID("quantity") [1,4]
 KEYWORD("THEN") [1,13]
 ID("total") [2,5]
 ASSIGN [2,11]
 ID("total") [2,14]
 OP("+") [2,20]
 ID("price") [2,22]
 OP("*") [2,28]
 ID("quantity") [2,30]
 END [2,38]
 ID("tax") [3,5]
 ASSIGN [3,9]
 ID("price") [3,12]
 OP("*") [3,18]
 NUMBER(0.05) [3,20]
 END [3,24]
 # unexpected at line 3, column: 26
 KEYWORD("ENDIF") [4,1]
 END [4,6]
*/