local colorStringsTree
local breakerChars = {}
local captureChars = {}

local nDefaultColor = colors.white
local nNumberColor = colors.orange
local nStringColor = colors.lightGray
local nCommentColor = colors.lime
local nBackgroundColor = colors.black

--RESPONSE PROTOTYPE--
local tokenResponsePrototype = {}
tokenResponsePrototype.__index = tokenResponsePrototype

setmetatable(tokenResponsePrototype, {
    __call = function (cls, ...)
    return cls.new(...)
    end,
})

--we check each token added to see if it is actually a number
function tokenResponsePrototype:add(args)
    self[#self+1] = args
    --if the text is a number but is not a hyphen (because tonumber will consider "-" a number) then set its type to number
    if tonumber(args.text) and args.text ~= "-" then
        args.type = "number"
    end
end

function tokenResponsePrototype.new()
    local self = setmetatable({}, tokenResponsePrototype)
    return self
end
----------------------

--COLORED STRING PROTOTYPE--
local coloredStringProtoype = {}
coloredStringProtoype.__index = coloredStringProtoype

setmetatable(coloredStringProtoype, {
  __call = function (cls, ...)
    return cls.new(...)
  end,
})

function coloredStringProtoype.new()
  local self = setmetatable({}, coloredStringProtoype)
  self.strings = {}
  return self
end

function coloredStringProtoype:addSubString(args)
  if args.text ~= nil and args.text ~= "" then
    self.strings[#self.strings+1] = {}
    for key,value in pairs(args) do
        self.strings[#self.strings][key] = value
    end
    return self.strings[#self.strings]
  end
end

function coloredStringProtoype:clearHighlights()
    for i = 1, #self.strings do
        self.strings[i]["highlights"] = nil
    end
end
-------------------------

--TREE SEARCH PROTOTYPE--
local searchBranchProtoype = {}
searchBranchProtoype.__index = searchBranchProtoype

setmetatable(searchBranchProtoype, {
  __call = function (cls, ...)
    return cls.new(...)
  end,
})

function searchBranchProtoype.new(char)
  local self = setmetatable({}, searchBranchProtoype)
  self.char = char
  self.branches = {}
  return self
end

function searchBranchProtoype:addBranch(args) -- char
  if args.char ~= nil and args.char ~= "" then
    self.branches[#self.branches+1] = args
    return self.branches[#self.branches]
  end
end

function searchBranchProtoype:setEndBranch(args) -- color, highlight
  if args.color ~= nil and args.color ~= "" then
    self.color = args.color
    self.endpoint = true
    if args.highlight ~= nil and args.highlight ~= "" then
        self.highlight = args.highlight
    end
  end
end
-------------------------

--TREE BUILDING AND SEARCHING--
function buildTree(listOfStrings)
    local tree = searchBranchProtoype("root")
    for key,value in pairs(listOfStrings) do
        addBranchtoTree(tree, key, value)
    end
    return tree
end

function addBranchtoTree(tree, stringToAdd, finalColor)
    local found = false
    if tree.branches ~= nil then
        for i = 1, #tree.branches do --foreach branch in the tree already there
            if string.sub(stringToAdd,1,1) == tree.branches[i].char and #stringToAdd > 1 then --if the first letter of our string matches one of the branches, and it has additional letters
                addBranchtoTree(tree.branches[i], string.sub(stringToAdd, 2), finalColor) --then check inside to see if another matches
                found = true
            elseif string.sub(stringToAdd,1,1) == tree.branches[i].char and #stringToAdd == 1 then --first letter matches and its also the last letter of the string.. even though a branch continues... set an endpoint here
                tree.branches[i]:setEndBranch{["color"] = finalColor}
            end
        end
    end
    if found == false then --didnt find a match.. add branch
        local temp = searchBranchProtoype(string.sub(stringToAdd, 1, 1))
        local newbranch = tree:addBranch(temp) -- the char should be equal to the first char of the stringToAdd
        if #stringToAdd == 1 then --if its the last char in the string
            newbranch:setEndBranch{["color"] = finalColor}
        else --if its not hte last char, go into the next one
            addBranchtoTree(newbranch, string.sub(stringToAdd, 2, #stringToAdd), finalColor) --then
        end
    end
end

function searchTree(branch, stringToSearch) --0 is no match, 1 is partial match, 2 is complete match
    if branch.branches ~= nil then
        for i = 1, #branch.branches do --foreach branch in the tree already there
            if string.sub(stringToSearch,1,1) == branch.branches[i].char then --if the first letter of our string matches one of the branches
                if #stringToSearch ~= 1 then
                    if #branch.branches[i].branches > 0 then --if the next branch is not an endpoint
                        return searchTree(branch.branches[i], string.sub(stringToSearch, 2)) --then check inside to see if another matches
                    else --if the next branch is an endpoint but the string continues
                        return {["result"] = 0}
                    end
                elseif #stringToSearch == 1 and #branch.branches[i].branches == 0 then --if this is the last char and the branch it matches is an endpoint
                    return {["result"] = 2, ["color"] = branch.branches[i].color}
                else --string has 1 char but branch is not an endpoint so this is a partial match
                    return {["result"] = 1}
                end
            end
        end
        return {["result"] = 0}
    end
end
-------------------------

--TOP LEVEL FUNCTIONS----

--return a table that is indexed
--entries contain the keys "text" and "type"
--"text" is the text inside the token
--type is either "token", "capture" (a string), "comment" (a comment string), "number" (a number)
function tokenizeString(stringToT)
    --where we store the tokens
    response = tokenResponsePrototype()
    --used to check the next character after a hyphen (for comment detection)
    local waitForHyphen = false
    --used to wait for a specific character and capture all characters inbetween. For double quote and single quotes
    local currentCapture = ""
    --the current token we're building
    local token = ""
    --loop through each character in the string we were asked to tokenize
    for i = 1, #stringToT do
        --get the character
        local char = string.sub(stringToT, i, i)
        --if we're waiting for a hyphen
        if waitForHyphen then
            waitForHyphen = false
            --if we found the hyphen, then set the character we should wait for to "EOL", representing End Of Line, and add the hyphen to the current token
            if char == "-" then
                token = token..char
                currentCapture = "EOL"
            --if the character is not a hyphen then add a hyphen to the response list to make up for the first hyphen we withheld
            else
                response:add{["text"] = "-", ["type"] = "token"}
            end
        end
        --if its the last character
        if i == #stringToT then
            --if either the character is not a break character, or we're capturing anything
            if breakerChars[char] == nil or currentCapture ~= "" then
                --if we're capturing for a comment then add the token and the character to the response with a 'type' of comment
                if currentCapture == "EOL" then
                    response:add{["text"] = token..char, ["type"] = "comment"}
                --if we're not capturing for a comment then add token and character to response with 'type' of token
                else
                    response:add{["text"] = token..char, ["type"] = "token"}
                end
            --if char is a break character and we're not capturing anything then add token and char to response, but seperatley
            else
                response:add{["text"] = token, ["type"] = "token"}
                response:add{["text"] = char, ["type"] = "token"}
            end
        --if char is a capture character
        elseif captureChars[char] ~= nil then
            --if char is the capture character we're looking for to end the current capture, and the previous character isnt a "\"
            --then add token..char to response, empty token, and stop the capture
            if char == currentCapture and string.sub(token, #token, #token) ~= "\\" then
                response:add{["text"] = token..char, ["type"] = "capture"}
                token = ""
                currentCapture = ""
            --if there is no current capture then start a new one, and add the current token to response
            elseif currentCapture == "" then
                currentCapture = char
                if token ~= "" then
                    response:add{["text"] = token, ["type"] = "token"}
                end
                token = char
            --if the current character is a capture character but not the one we're looking for then add it to the token
            else
                token = token..char
            end
        --if the character is a break character and we're not in the middle of a capture
        elseif breakerChars[char] ~= nil and currentCapture == "" then
            --add token to response
            if token ~= "" then
                response:add{["text"] = token, ["type"] = "token"}
                token = ""
            end
            --if the character is a hyphen then withold doing anything and set waitForHyphen in order to check the next character
            if char == "-" then
                waitForHyphen = true
            --if the character is a normal break character then add it to the response
            else
                response:add{["text"] = char, ["type"] = "token"}
            end
        --if the character is not a break character, or capture character then add it to the token
        else
            token = token..char
        end
    end
    return response
end

--converts regular string into a colored string
function stringToColorString(stringToWrite)
    local response = coloredStringProtoype()
    response["text"] = stringToWrite
    local tokens = tokenizeString(stringToWrite)
    for i = 1, #tokens do
        local color
        if tokens[i]["type"] == "comment" then
            color = nCommentColor
        elseif tokens[i]["type"] == "capture" then
            color = nStringColor
        elseif tokens[i]["type"] == "number" then
            color = nNumberColor
        elseif tokens[i]["type"] == "token" then
            color = searchTree(colorStringsTree, tokens[i].text)["color"] or nDefaultColor
        end
        response:addSubString{["text"] = tokens[i].text, ["color"] = color}
    end
    return response
end

function writeColorString(coloredString)
    for i=1, #coloredString.strings do
        if coloredString.strings[i]["highlight"] ~= nil then
            term.setBackgroundColor(coloredString.strings[i]["highlight"])
        else
            term.setBackgroundColor(nBackgroundColor)
        end
        term.setTextColor(coloredString.strings[i]["color"])
        term.write(coloredString.strings[i]["text"])
    end
    term.setTextColor(nDefaultColor)
    term.setBackgroundColor(nBackgroundColor)
end

---------------------------

--initalize everything
function initialize()

    local colorStrings = {}

    local keywordColor = colors.blue
    local osAPIcolor = colors.red

    colorStrings["local"] = keywordColor
    colorStrings["function"] = keywordColor
    colorStrings["if"] = keywordColor
    colorStrings["os.unloadAPI"] = osAPIcolor
    colorStrings["os.pullEvent"] = osAPIcolor

    colorStringsTree = buildTree(colorStrings)

    ----------------------------

    breakerChars[" "] = nDefaultColor
    breakerChars["-"] = nDefaultColor
    breakerChars["+"] = nDefaultColor
    breakerChars["/"] = nDefaultColor
    breakerChars["*"] = nDefaultColor
    breakerChars["#"] = nDefaultColor
    breakerChars["("] = nDefaultColor
    breakerChars[")"] = nDefaultColor
    breakerChars["{"] = nDefaultColor
    breakerChars["}"] = nDefaultColor
    breakerChars["["] = nDefaultColor
    breakerChars["]"] = nDefaultColor
    breakerChars["~"] = nDefaultColor
    breakerChars["+"] = nDefaultColor
    breakerChars[">"] = nDefaultColor
    breakerChars["<"] = nDefaultColor
    breakerChars[","] = nDefaultColor

    captureChars["\""] = colors.gray
    captureChars["'"] = colors.gray
    captureChars["EOL"] = colors.green
end

initialize()
--tokenize the tests string
local coloredTokens = stringToColorString("local something = \"test\\\" string\" --test comment")

if term.isColor() then
    writeColorString(coloredTokens)
else
    local file = fs.open("tableSaved","w")
    file.write(textutils.serialize(coloredTokens))
    file.close()
end