{-# 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.Text.Lazy.IO as Lt
import qualified Data.Map.Strict as Map
import qualified Data.Text.Lazy as Lt
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 (traceDocM)

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 . print . prettyTm . quote =<< 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 = Lt.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

              metas <- liftIO $ atomicModifyIORef' (unsolvedMetas env) (\x -> (mempty, x))
              unless (Map.null metas) $ do
                liftIO $ reportUnsolved (Just inp) metas

              liftIO $ writeIORef envvar env
              loop env envvar

    reload :: ElabEnv -> IORef ElabEnv -> InputT IO ()
    reload env@ElabEnv{loadedFiles} envref = do
      newe <- liftIO $ try $ mkrepl <$> checkFiles (reverse loadedFiles)
      case newe of
        Left e -> do
          liftIO $ do displayExceptions' ":r" (e :: SomeException)
          loop env envref
        Right newe -> do
          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 $ atomicModifyIORef' (unsolvedMetas env) (\x -> (mempty, x))
    unless (Map.null metas) $ do
      liftIO $ reportUnsolved Nothing 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) *> k
      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
  Lt.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
      Lt.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 -> Lt.pack "left"
          VI1 -> Lt.pack "right"
          _ -> render . prettyTm $ quoteWith False mempty ep
        left = render (prettyTm (quoteWith False mempty le))
        right = render (prettyTm (quoteWith False mempty ri))
      Lt.putStrLn . Lt.unlines $
        [ "\n\x1b[1;32mnote\x1b[0m: This path was expected to fill the diagram <<"
        , "\t " <> redact left <> " " <> Lt.replicate 7 (Lt.singleton '\x2500') <> " " <> redact right
        , "      >> in the direction " <> render (pretty dir) <> ", but the " <> endp <> " endpoint is incorrect"
        ]
  , Handler \(NotEqual ta tb) -> do
      Lt.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)
  , Handler \Elab.PathConPretype -> do
      putStrLn $ "Pretypes can not have path constructors, either change this definition so it lands in Type or remove it."
  ]

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

reportUnsolved :: Foldable t => Maybe Text -> Map.Map MV (t (Seq Projection, Value)) -> IO ()
reportUnsolved code metas = do
  for_ (Map.toList metas) \(m, p) ->  do
    case mvSpan m of
      Just (f, s, e) ->
        case code of
          Just code -> T.putStrLn $ squiggleUnder s e code
          Nothing -> T.putStrLn . squiggleUnder s e =<< T.readFile (T.unpack f)
      Nothing -> pure ()

    case null p of
      True -> do
        Lt.putStrLn . render $ "Unsolved metavariable" <+> prettyTm (Meta m) <+> pretty ':' <+> align (prettyVl (zonk (mvType m)) <> pretty '.')
      _ -> do
        Lt.putStrLn . render $
          "Unsolved metavariable" <+> prettyTm (Meta m) <+> pretty ':' <+> align (prettyVl (zonk (mvType m)) <+> "should satisfy:")
        for_ p \p ->
          Lt.putStrLn . render $ indent 2 $ prettyTm (quote (zonk (VNe (HMeta m) (fst p)))) <+> pretty '≡' <+> prettyTm (quote (snd p))
    
    when (mvInteraction m && not (Map.null (mvContext m))) do
      putStrLn "Context (first 10 entries):"
      for_ (take 10 (Map.toList (mvContext m))) \(x, t) -> unless (isIdkT t) do
        Lt.putStrLn . render $ indent 2 $ pretty x <+> pretty ':' <+> prettyVl (zonk t)

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
      traceDocM (viaShow t)
      dumpTokens