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

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

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.WiredIn
import Elab.Eval
import Elab

import Options.Applicative

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

import Prettyprinter

import Syntax.Pretty
import Syntax

import System.Console.Haskeline
import System.Exit
import qualified Data.Set as Set
import Data.Maybe
import Presyntax.Tokens
import Debug.Trace

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 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 = 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 $ DeclaredUndefined v
      _ -> pure ()
    k
  go (x:xs) k = do
    code <- liftIO $ Bsl.readFile x
    case runAlex (code <> Bsl.singleton 10) parseProg of
      Left e -> liftIO $ print e *> error (show e)
      Right prog ->
        checkProgram prog (go 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 <+> align (nest (negate 1) (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 (pretty x)
  , Handler \(DeclaredUndefined n) -> do
      putStrLn $ "Name declared but not defined: " ++ show (pretty n)
  ]

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)))

newtype DeclaredUndefined = DeclaredUndefined { declaredUndefName :: Name }
  deriving (Eq, Show, Exception)

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