{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Main where

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.Text as T
import Data.Text ( Text )
import Data.Foldable

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

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

import System.Exit
import Syntax.Pretty
import Elab.WiredIn
import Options.Applicative
import Control.Monad
import Syntax
import Prettyprinter
import Control.Monad.IO.Class
import Control.Monad.Reader
import System.Console.Haskeline

main :: IO ()
main = do
  opts <- execParser parser
  case opts of
    Load files -> do
      env <- checkFiles files
      enterReplIn env
    Check files verbose -> do
      env <- checkFiles files
      when verbose $ dumpEnv (getEnv env)
    Repl -> enterReplIn emptyEnv

enterReplIn :: ElabEnv -> IO ()
enterReplIn env = runInputT defaultSettings (loop env') where
  env' = env { commHook = T.putStrLn . render . prettyTm . quote . zonk }

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

checkFiles :: [String] -> IO ElabEnv
checkFiles files = runElab (go files ask) emptyEnv where
  go [] k = k
  go (x:xs) k = do
    code <- liftIO $ Bsl.readFile x
    case runAlex code parseProg of
      Left e -> liftIO $ print e *> error (show e)
      Right prog -> do
        env <- ask
        liftIO $ runElab (checkProgram prog (go xs k)) env
                  `catch` \e -> displayAndDie (T.decodeUtf8 (Bsl.toStrict code)) (e :: SomeException)

dumpEnv :: Map.Map Name (NFType, Value) -> IO ()
dumpEnv env = for_ (Map.toList env) $ \(name, (nft, _)) ->
  T.putStrLn $ render (pretty name <+> colon <+> prettyTm (quote (zonk nft)))

parser :: ParserInfo Opts
parser = info (subparser (load <> check) <|> pure Repl <**> helper) (header "cubical - a cubical programming language")
  where
    load  = command "load" $ info (fmap Load (some (argument str (metavar "file..."))) <**> helper) (progDesc "Check and load a list of files in the REPL")
    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] }
  | 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 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))
      T.putStrLn . T.unlines $
        [ "\x1b[1;32mnote\x1b[0m: This path was expected to fill the diagram"
        , "\t " <> render (prettyTm (quote le)) <> " " <> T.replicate 7 (T.singleton '\x2500') <> " " <> render (prettyTm (quote ri))
        , "\x1b[1;32mnote\x1b[0m: the " <> endp <> " endpoint is incorrect"
        ]
  , Handler \(NotEqual ta tb) -> do
      putStrLn . unlines $
        [ "\x1b[1;31merror\x1b[0m: Mismatch between actual and expected types:"
        , "  * \x1b[1mActual\x1b[0m:   " ++ showValue (zonk ta)
        , "  * \x1b[1mExpected\x1b[0m: " ++ showValue (zonk tb)
        ]
  , Handler \(NoSuchPrimitive x) -> do
      putStrLn $ "Unknown primitive: " ++ T.unpack x
  , Handler \(NotInScope x) -> do
      putStrLn $ "Variable not in scope: " ++ show x
  ]

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.pack (show (Posn l c, Posn l' c'))