module Datalog.LSP.Highlight where

import Control.Lens (Lens', (^.))
import Data.Either (fromRight)
import Data.Maybe (fromJust)
import Datalog.Parser (SrcLoc (..))
import Datalog.Syntax (Atom' (Atom), ConId (..), Program' (..), RelId (RelId), Rule' (..), Term' (..))
import Language.LSP.Protocol.Lens (HasParams, HasTextDocument, HasUri, params, textDocument, uri)
import Language.LSP.Protocol.Message
import Language.LSP.Protocol.Types
import Language.LSP.Server
import Text.Megaparsec
import Language.LSP.VFS (virtualFileText)
import Datalog.Parser (parseProgram)

tokenHandler :: Handlers (LspM ())
tokenHandler = requestHandler SMethod_TextDocumentSemanticTokensFull $ \req responder -> do
    c <- virtualFileText . fromJust <$> getVirtualFile (toNormalizedUri $ req ^. docUri)
    let
        parsedProgram = fromRight (error "") $ runParser parseProgram "" c
        absTokens = highlightProg parsedProgram
        tokens = fromRight (error "")$ makeSemanticTokens defaultSemanticTokensLegend absTokens

    responder (Right $ InL tokens)

docUri :: (HasParams s s1, HasTextDocument s1 s2, HasUri s2 a) => Lens' s a
docUri = params . textDocument . uri

highlightProg :: Program' SrcLoc -> [SemanticTokenAbsolute]
highlightProg (Program _ rs) = rs >>= highlightRule

highlightRule :: Rule' SrcLoc -> [SemanticTokenAbsolute]
highlightRule (Rule _ a as) = a : as >>= highlightAtom

highlightAtom :: Atom' SrcLoc -> [SemanticTokenAbsolute]
highlightAtom (Atom loc (RelId relId) ts) = highlightRel ++ (ts >>= highlightTerm) where
    highlightRel = []

getConLoc :: Term' SrcLoc -> SrcLoc
getConLoc (Con loc _) = loc
getConLoc (Var loc _) = loc

highlightTerm :: Term' SrcLoc -> [SemanticTokenAbsolute]
highlightTerm t =
    [ SemanticTokenAbsolute
        { _line = unPos' startLine
        , _startChar = unPos' startCol
        , _length = fromIntegral $ unPos stopCol - unPos startCol
        , _tokenType = case t of
            Con _ _ -> SemanticTokenTypes_Number
            Var _ _ -> SemanticTokenTypes_Variable
        , _tokenModifiers = []
        }
    ]
  where
    loc = getConLoc t
    unPos' = fromIntegral . unPos
    startLine = sourceLine . start $ loc
    startCol = sourceColumn . start $ loc
    stopCol = sourceColumn . end $ loc