{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE NamedFieldPuns #-}
module Main where

import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Exception

import qualified Data.ByteString.Lazy as Bsl
import qualified Data.Text.Encoding as T
import qualified Data.Map.Strict as Map
import qualified Data.Text.IO as T
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.Sequence (Seq)
import Data.Text ( Text )
import Data.Foldable
import Data.Maybe
import Data.IORef

import Debug.Trace

import Elab.Monad hiding (switch)
import Elab.WiredIn
import Elab.Eval
import Elab

import Options.Applicative

import Presyntax.Presyntax (Posn(Posn))
import Presyntax.Parser
import Presyntax.Tokens
import Presyntax.Lexer

import Prettyprinter

import Syntax.Pretty
import Syntax

import System.Console.Haskeline
import System.Exit

main :: IO ()
main = do
  opts <- execParser parser
  case opts of
    Load files exp -> do
      env <- checkFiles files
      case exp of
        [] -> enterReplIn env
        xs -> traverse_ (evalArgExpr env) xs
    Check files verbose -> do
      env <- checkFiles files
      when verbose $ dumpEnv env
    Repl -> enterReplIn =<< emptyEnv

evalArgExpr :: ElabEnv -> String -> IO ()
evalArgExpr env str =
  case runAlex (Bsl.fromStrict (T.encodeUtf8 inp)) parseExp of
    Right e ->
      flip runElab env (do
        (e, _) <- infer e
        liftIO . putStrLn . show . prettyTm . quote . zonk =<< Elab.Eval.eval e)
      `catch` \e -> do
          displayExceptions' inp (e :: SomeException)
    Left e -> liftIO $ print e
  where
    inp = T.pack str

enterReplIn :: ElabEnv -> IO ()
enterReplIn env =
  do
    let env' = mkrepl env
    envref <- newIORef env'
    runInputT (setComplete (complete envref) defaultSettings) (loop env' envref)
  where
    mkrepl env = env { commHook = T.putStrLn . render }

    complete :: IORef ElabEnv -> (String, String) -> IO (String, [Completion])
    complete c = completeWord Nothing " \n\t\r" go where
      go w = do
        env <- readIORef c
        let
          w' = T.pack w
          words = Set.toList $ Set.filter ((w' `T.isPrefixOf`) . getNameText) (definedNames env)
        pure (map (simpleCompletion . T.unpack . getNameText) words)

    loop :: ElabEnv -> IORef ElabEnv -> InputT IO ()
    loop env envvar = do
      inp <- fmap T.pack <$> getInputLine "% "
      case inp of
        Nothing -> pure ()
        Just inp | ":r" `T.isPrefixOf` inp -> reload env envvar
        Just inp ->
          case runAlex (Bsl.fromStrict (T.encodeUtf8 inp)) parseRepl of
            Left e -> do
              liftIO $ print e
              loop env envvar
            Right st -> do
              env <- liftIO $
                runElab (checkStatement st ask) env
                  `catch` \e -> do
                    displayExceptions' inp (e :: SomeException)
                    pure env
              liftIO $ writeIORef envvar env
              loop env envvar

    reload :: ElabEnv -> IORef ElabEnv -> InputT IO ()
    reload ElabEnv{loadedFiles} envref = do
      newe <- liftIO $ mkrepl <$> checkFiles (reverse loadedFiles)
      liftIO $ writeIORef envref newe
      loop newe envref

checkFiles :: [String] -> IO ElabEnv
checkFiles files = runElab (go 1 files ask) =<< emptyEnv where
  size = length files
  sl = length (show size)

  pad s
    | length s < sl = replicate (sl - length s) ' ' ++ s
    | otherwise = s

  go _ [] k = do
    env <- ask
    for_ (Map.toList (nameMap env)) \case
      (_, v@Defined{})
        | Set.member v (definedNames env) -> pure ()
        | otherwise ->
          let
            pos = fromJust (Map.lookup v (whereBound env))
           in withSpan (fst pos) (snd pos) $ throwElab $ Elab.DeclaredUndefined v
      _ -> pure ()

    metas <- liftIO $ readIORef (unsolvedMetas env)
    unless (Map.null metas) $ do
      liftIO $ reportUnsolved metas
    k

  go n (x:xs) k = do
    liftIO . putStrLn $ "[" ++ pad (show n) ++ "/" ++ show size ++ "] Loading " ++ x
    code <- liftIO $ Bsl.readFile x
    case runAlex (code <> Bsl.singleton 10) parseProg of
      Left e -> liftIO $ print e *> error (show e)
      Right prog ->
        local (\e -> e { currentFile = Just (T.pack x), loadedFiles = x:loadedFiles e }) (checkProgram prog (go (n + 1 :: Int) xs k))
          `catchElab` \e -> liftIO $ displayAndDie (T.decodeUtf8 (Bsl.toStrict code)) (e :: SomeException)

dumpEnv :: ElabEnv -> IO ()
dumpEnv env = for_ (Map.toList (nameMap env)) $ \(_, name) ->
  let nft = fst $ getEnv env Map.! name in
  T.putStrLn $ render (pretty name <+> nest (negate 1) (colon <+> align (prettyTm (quote (zonk nft)))))

parser :: ParserInfo Opts
parser = info (subparser (load <> check <> repl) <|> pure Repl <**> helper) (header "cubical - a cubical programming language")
  where
    load  = command "load" $
      info ((Load <$> (some (argument str (metavar "file...")))
                  <*> (many (strOption (long "eval" <> short 'e' <> help "Also evaluate this expression"))))
       <**> helper) (progDesc "Check and load a list of files in the REPL")

    repl  = command "load" $
      info ((Load <$> (many (argument str (metavar "file...")))
                  <*> (many (strOption (long "eval" <> short 'e' <> help "Also evaluate this expression"))))
       <**> helper) (progDesc "Enter the REPL, optionally with loaded files")

    check = command "check" $
      info ((Check <$> some (argument str (metavar "file..."))
                   <*> switch ( long "verbose"
                             <> short 'v'
                             <> help "Print the types of all declared/defined values"))
       <**> helper)
           (progDesc "Check a list of files")

data Opts
  = Load  { files :: [String], eval :: [String] }
  | Check { files :: [String], verbose :: Bool }
  | Repl
  deriving (Eq, Show, Ord)

displayAndDie :: Exception e => Text -> e -> IO a
displayAndDie lines e = do
  () <- throwIO e `catches` displayExceptions lines
  exitFailure

displayExceptions :: Text -> [Handler ()]
displayExceptions lines =
  [ Handler \(WhileChecking a b e) -> do
      T.putStrLn $ squiggleUnder a b lines
      displayExceptions' lines e
  , Handler \(SeeAlso a b e) -> do
      displayExceptions' lines e
      T.putStrLn $ squiggleUnder a b lines
  , Handler \(AttachedNote n e) -> do
      displayExceptions' lines e
      T.putStrLn $ "\x1b[1;32mnote\x1b[0m: " <> render n
  , Handler \(WhenCheckingEndpoint dir le ri ep e) -> do
      displayExceptions' lines e
      let
        endp = case ep of
          VI0 -> T.pack "left"
          VI1 -> T.pack "right"
          _ -> T.pack $ show (prettyTm (quote ep))
        left = render (prettyTm (quote le))
        right = render (prettyTm (quote ri))
      T.putStrLn . T.unlines $
        [ "\n\x1b[1;32mnote\x1b[0m: This path was expected to fill the diagram <<"
        , "\t " <> redact left <> " " <> T.replicate 7 (T.singleton '\x2500') <> " " <> redact right
        , "      >> in the direction " <> render (pretty dir) <> ", but the " <> endp <> " endpoint is incorrect"
        ]
  , Handler \(NotEqual ta tb) -> do
      T.putStrLn . render . vsep $
        [ "\x1b[1;31merror\x1b[0m: Mismatch between actual and expected types:"
        , indent 2 $ "* \x1b[1mActual\x1b[0m:" <> softline <> align (prettyVl (zonk ta))
        , indent 2 $ "* \x1b[1mExpected\x1b[0m:" <> softline <> align (prettyVl (zonk tb))
        ]
  , Handler \(NoSuchPrimitive x) -> do
      putStrLn $ "Unknown primitive: " ++ T.unpack x
  , Handler \(NotInScope x) -> do
      putStrLn $ "Variable not in scope: " ++ show (pretty x)
  , Handler \(Elab.DeclaredUndefined n) -> do
      putStrLn $ "Name declared but not defined: " ++ show (pretty n)
  ]

redact :: Text -> Text
redact x
  | length (T.lines x) >= 2 = head (T.lines x) <> T.pack "\x1b[1;32m[...]\x1b[0m"
  | otherwise = x

reportUnsolved :: Foldable t => Map.Map MV (t (Seq Projection, Value)) -> IO ()
reportUnsolved metas = do
  for_ (Map.toList metas) \(m, p) ->  do
    case mvSpan m of
      Just (f, s, e) -> T.putStrLn . squiggleUnder s e =<< T.readFile (T.unpack f)
      Nothing -> pure ()
    T.putStrLn . render $
      "Unsolved metavariable" <+> prettyTm (Meta m) <+> pretty ':' <+> prettyTm (quote (mvType m)) <+> "should satisfy:"
    for_ p \p ->
      T.putStrLn . render $ indent 2 $ prettyTm (quote (VNe (HMeta m) (fst p))) <+> pretty '≡' <+> prettyTm (quote (snd p))


displayExceptions' :: Exception e => Text -> e -> IO ()
displayExceptions' lines e = displayAndDie lines e `catch` \(_ :: ExitCode) -> pure ()

squiggleUnder :: Posn -> Posn -> Text -> Text
squiggleUnder (Posn l c) (Posn l' c') file
  | l == l' =
    let
      line     = T.pack (show l) <> " | " <> T.lines file !! (l - 1)
      padding  = T.replicate (length (show l)) (T.singleton ' ') <> " |"
      squiggle = T.replicate c " " <> T.pack "\x1b[1;31m" <> T.replicate (c' - c) "~" <> T.pack "\x1b[0m"

     in T.unlines [ padding, line, padding <> squiggle ]
  | otherwise = T.unlines (take (l' - l) (drop l (T.lines file)))

dumpTokens :: Alex ()
dumpTokens = do
  t <- alexMonadScan
  case tokenClass t of
    TokEof -> pure ()
    _ -> do
      traceM (show t)
      dumpTokens