module Tessera.Serialize where

import Data.Char (isAlphaNum, isDigit)
import Data.List (isPrefixOf)

import Tessera.Check (checkProgram)
import Tessera.Env (Env, lookupConstructor)
import Tessera.Matrix (Key (..))
import Tessera.Parser (parseProgram)
import Tessera.Syntax
import Tessera.Tree

data Artifact = Artifact
  { artifactFunction :: String
  , artifactStrategy :: Strategy
  , artifactTree :: Tree
  , artifactSource :: String
  , artifactProgram :: Program
  , artifactEnv :: Env
  }

formatVersion :: String
formatVersion = "1"

encodeArtifact :: String -> String -> Strategy -> Tree -> String
encodeArtifact function source strategy tree = unlines
  [ "tessera-tree " ++ formatVersion
  , "function " ++ function
  , "strategy " ++ strategyName strategy
  , "tree " ++ encodeTree tree
  , "source-begin"
  , source
  , "source-end"
  ]

encodeTree :: Tree -> String
encodeTree tree = case tree of
  Success index -> "L" ++ show index
  Failure -> "X"
  Bind (Occurrence path) name inner -> "B[" ++ joinIntegers path ++ "|" ++ name ++ "]" ++ encodeTree inner
  Switch (Occurrence path) branches fallback ->
    "W[" ++ joinIntegers path ++ "]{"
      ++ concatMap encodeBranch branches
      ++ "}[" ++ maybe "" encodeTree fallback ++ "]"

encodeBranch :: (Key, Tree) -> String
encodeBranch (key, tree) = encodeKey key ++ ":" ++ encodeTree tree ++ ";"

encodeKey :: Key -> String
encodeKey key = case key of
  KeyCon name -> "c" ++ name
  KeyInt n -> "i" ++ show n
  KeyBool b -> if b then "bT" else "bF"
  KeyTuple n -> "t" ++ show n

joinIntegers :: [Int] -> String
joinIntegers [] = ""
joinIntegers [x] = show x
joinIntegers (x : rest) = show x ++ "," ++ joinIntegers rest

decodeArtifact :: String -> Either String Artifact
decodeArtifact input = do
  let numbered = zip [1 ..] (lines input)
  version <- lookupHeader "tessera-tree" numbered
  if version /= formatVersion then Left ("unsupported tree format version " ++ version) else Right ()
  function <- lookupHeader "function" numbered
  strategyText <- lookupHeader "strategy" numbered
  strategy <- case strategyText of
    "leftmost" -> Right Leftmost
    "heuristic" -> Right Heuristic
    _ -> Left ("unknown strategy " ++ strategyText)
  treeText <- lookupHeader "tree" numbered
  tree <- decodeTree treeText
  source <- extractSource numbered
  program <- either (Left . renderDiagnostic) Right (parseProgram source)
  env <- either (Left . renderDiagnostic) Right (checkProgram program)
  if function `elem` [name | MatchDeclaration name _ _ _ <- program]
    then Right ()
    else Left ("function " ++ function ++ " is not declared in the embedded source")
  validateTree env tree
  return (Artifact function strategy tree source program env)

lookupHeader :: String -> [(Int, String)] -> Either String String
lookupHeader key numbered =
  case [drop (length key + 1) line | (_, line) <- numbered, (key ++ " ") `isPrefixOf` line] of
    (value : _) -> Right value
    [] -> Left ("missing " ++ key ++ " header")

extractSource :: [(Int, String)] -> Either String String
extractSource numbered = case break (\(_, line) -> line == "source-begin") numbered of
  (_, _ : rest) -> case break (\(_, line) -> line == "source-end") rest of
    (body, _ : _) -> Right (unlines (map snd body))
    _ -> Left "missing source-end marker"
  _ -> Left "missing source-begin marker"

renderDiagnostic :: Diagnostic -> String
renderDiagnostic (Diagnostic span' message) =
  "line " ++ show (spanLine span') ++ " column " ++ show (spanColumn span') ++ ": " ++ message

decodeTree :: String -> Either String Tree
decodeTree input = do
  (tree, rest) <- parseTree input
  if null rest then Right tree else Left ("trailing characters in tree encoding " ++ take 20 rest)

parseTree :: String -> Either String (Tree, String)
parseTree ('L' : rest) =
  let (digits, remaining) = span isDigit rest
  in if null digits then Left "malformed success node" else Right (Success (read digits), remaining)
parseTree ('X' : rest) = Right (Failure, rest)
parseTree ('B' : '[' : rest) = do
  let (inside, afterInside) = break (== ']') rest
  (path, name) <- case break (== '|') inside of
    (p, '|' : n) -> Right (p, n)
    _ -> Left "malformed bind node"
  (inner, remaining) <- parseTree (drop 1 afterInside)
  Right (Bind (Occurrence (parseIntegers path)) name inner, remaining)
parseTree ('W' : '[' : rest) = do
  let (path, afterPath) = break (== ']') rest
  afterPath' <- expect ']' afterPath
  afterBrace <- expect '{' afterPath'
  (branches, afterBranches) <- parseBranches afterBrace
  afterBranches' <- expect '}' afterBranches
  afterBracket <- expect '[' afterBranches'
  let (fallbackText, afterFallback) = break (== ']') afterBracket
  fallback <- if null fallbackText then Right Nothing else do
    (tree, _) <- parseTree fallbackText
    Right (Just tree)
  afterFallback' <- expect ']' afterFallback
  Right (Switch (Occurrence (parseIntegers path)) branches fallback, afterFallback')
parseTree _ = Left "malformed tree encoding"

expect :: Char -> String -> Either String String
expect c (x : rest) | x == c = Right rest
expect c _ = Left ("expected " ++ [c])

parseBranches :: String -> Either String ([(Key, Tree)], String)
parseBranches input = go input []
  where
    go ('}' : rest) acc = Right (reverse acc, '}' : rest)
    go (c : rest) acc = do
      (key, rest1) <- parseKey c rest
      rest2 <- expect ':' rest1
      (tree, rest3) <- parseTree rest2
      rest4 <- expect ';' rest3
      go rest4 ((key, tree) : acc)
    go [] _ = Left "unterminated branch list"

parseKey :: Char -> String -> Either String (Key, String)
parseKey 'c' rest = let (name, remaining) = span (\c -> isAlphaNum c || c == '_') rest in Right (KeyCon name, remaining)
parseKey 'i' rest = let (digits, remaining) = span isDigit rest in Right (KeyInt (read digits), remaining)
parseKey 'b' rest = case rest of
  ('T' : remaining) -> Right (KeyBool True, remaining)
  ('F' : remaining) -> Right (KeyBool False, remaining)
  _ -> Left "malformed boolean key"
parseKey 't' rest = let (digits, remaining) = span isDigit rest in Right (KeyTuple (read digits), remaining)
parseKey _ _ = Left "malformed branch key"

parseIntegers :: String -> [Int]
parseIntegers [] = []
parseIntegers text = map read (splitOn ',' text)

splitOn :: Char -> String -> [String]
splitOn _ [] = []
splitOn c text = case break (== c) text of
  (piece, []) -> [piece]
  (piece, _ : rest) -> piece : splitOn c rest

validateTree :: Env -> Tree -> Either String ()
validateTree env = go
  where
    go tree = case tree of
      Success _ -> Right ()
      Failure -> Right ()
      Bind (Occurrence path) _ inner -> validatePath path >> go inner
      Switch (Occurrence path) branches fallback -> do
        validatePath path
        mapM_ (validateKey env) (map fst branches)
        mapM_ (go . snd) branches
        mapM_ go fallback

    validatePath path = if all (>= 0) path then Right () else Left "negative occurrence path"

validateKey :: Env -> Key -> Either String ()
validateKey env key = case key of
  KeyCon name -> case lookupConstructor env name of
    Just _ -> Right ()
    Nothing -> Left ("unknown constructor in tree " ++ name)
  KeyInt _ -> Right ()
  KeyBool _ -> Right ()
  KeyTuple _ -> Right ()