{-# 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 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 =<< checkFiles ["./test.tt"]

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 = 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 $ Elab.DeclaredUndefined v
      _ -> pure ()

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

    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 ->
        local (\e -> e { currentFile = Just (T.pack x) }) (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 ((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")
    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 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 \(Elab.DeclaredUndefined n) -> do
      putStrLn $ "Name declared but not defined: " ++ show (pretty n)
  ]

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