{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.Monad.IRTree qualified as IRTree
import ReduceC qualified

import Colog
import Control.Applicative
import Control.Monad
import Control.Monad.State
import Data.Foldable
import Data.Function
import Data.Functor
import Data.IORef (atomicModifyIORef', atomicWriteIORef, newIORef)
import Data.List.NonEmpty qualified as NE
import Data.RPath qualified as RPath
import Data.Text qualified as Text
import Data.Time (getCurrentTime)
import Data.Time qualified as Time
import Data.Time.Format.ISO8601 (iso8601Show)
import GHC.Stack
import Language.C qualified as C
import Options.Applicative
import System.Directory
import System.Exit
import System.FilePath
import System.IO
import System.IO.Error (tryIOError)
import System.IO.Temp
import System.Process.Typed
import Text.Pretty.Simple
import Text.PrettyPrint qualified as P
import Prelude hiding (log)

main :: IO ()
main =
  join
    . execParser
    $ info (run <**> helper)
      . fold
    $ []

process :: (WithLog env (Msg sev) m, MonadIO m) => sev -> Text.Text -> m a -> m a
process sev p ma = do
  msg "Started "
  a <- ma
  msg "Done "
  pure a
 where
  s = withFrozenCallStack callStack
  msg t = Colog.logMsg (Msg sev s (t <> p))

data Mode
  = Lin
  | Exp
  deriving (Show, Read, Eq, Ord, Enum)

run :: (HasCallStack) => Parser (IO ())
run = do
  mode <-
    option auto $
      fold
        [ long "mode"
        , help "search mode (Lin, Exp)"
        , value Exp
        ]

  verbose <-
    flag False True $
      fold
        [ long "verbose"
        , short 'v'
        , help "log more"
        ]

  checkmode <-
    flag False True $
      fold
        [ long "dry-run"
        , short 'n'
        , help "don't do any reduction"
        ]

  keywords <-
    many . option auto $
      fold
        [ long "keyword"
        , short 'k'
        , help "special versions of the code"
        ]

  validity <-
    optional $
      option auto $
        fold
          [ long "validity"
          , short 'V'
          , help "check the output for the validity code and output debug information"
          ]

  debug <-
    flag False True $
      fold
        [ long "debug"
        , short 'D'
        , help "keep debug versions of all copies"
        ]

  fixpoint <-
    flag False True $
      fold
        [ long "fixpoint"
        , short 'F'
        , help "run the computation to a fixpoint"
        ]

  pedandic <-
    flag False True $
      fold
        [ long "pedandic"
        , short 'P'
        , help "when checking for validity, throw error if command fails"
        ]

  cmd <- strArgument $ fold [metavar "COMMAND"]

  file <- strArgument $ fold [metavar "FILE"]

  pure do
    time <- getCurrentTime
    let
      fmt m = do
        t' <- getCurrentTime
        pure
          ( Text.pack
              ( Time.formatTime
                  Time.defaultTimeLocale
                  "%_3m:%04ES "
                  (t' `Time.diffUTCTime` time)
              )
              <> fmtMessage m
          )

    let logger =
          filterBySeverity (if verbose then Debug else Info) msgSeverity $
            cmapM fmt logTextStderr

    usingLoggerT logger do
      counter <- liftIO (newIORef (0 :: Int))

      let
        test f = process D ("test " <> Text.pack f) do
          err <- liftIO $ runProcess (proc cmd [])
          log D (Text.pack $ show err)
          pure err

        output f c = process D ("writing " <> Text.pack f) do
          let x = P.render (C.pretty (c $> C.undefNode))
          liftIO $ writeFile f x

        cleanup f = do
          logError "Cleaning up"
          liftIO do
            copyFile (f <.> "bak") f
            removeFile (f <.> "bak")
          liftIO exitFailure

        check ::
          FilePath ->
          [RPath.AnnotatedChoice (String, C.Position)] ->
          C.CTranslationUnit C.NodeInfo ->
          LoggerT (Msg Severity) IO Bool
        check f l c = process D "Checking predictate" do
          let xs = NE.nonEmpty (filter ((RPath.Yes ==) . RPath.choice) l)
          let decided = length l - RPath.numberOfUndecided l

          count <- liftIO (atomicModifyIORef' counter (\a -> (a + 1, a)))
          logInfo
            ( "Checking (num: "
                <> Text.pack (show count)
                <> ") D="
                <> Text.pack (show decided <> "/" <> show (length l))
                <> ": "
                <> Text.pack (maybe "-" ((\(r, p) -> r <> " at " <> show p) . RPath.label . NE.last) xs)
            )
          output f c
          res <- test f
          when debug $ liftIO do
            copyFile f ("case-" <> show count <.> "c")

          ret <- case res of
            ExitSuccess -> do
              liftIO $ copyFile f (f <.> "last")
              pure True
            ExitFailure x
              | Just x == validity -> do
                  folder <- liftIO do
                    temp <- getCanonicalTemporaryDirectory
                    createTempDirectory temp "rtree-c"
                  logWarning $ "Produced invalid code, see " <> Text.pack folder <> " for details"
                  liftIO do
                    withFile (folder </> "reduction.path") WriteMode \h -> do
                      forM (takeWhile ((/= RPath.Undecided) . RPath.choice) l) \(RPath.AnnotatedChoice c' (msg, pos)) -> do
                        hPutStr h [RPath.debugShowChoice c']
                        hPutStr h " "
                        hPutStr h msg
                        hPutStr h " at "
                        hPrint h pos
                    copyFile cmd (folder </> takeFileName cmd)
                    copyFile f (folder </> takeFileName f)
                    tryIOError do
                      copyFile (f <.> "last") (folder </> takeFileName (f <.> "last"))
                    pure folder
                  when pedandic do
                    logError "Pedandic, stoping now."
                    liftIO exitFailure
                  pure False
              | otherwise -> pure False
          logInfo ("Predicate was " <> Text.pack (show ret))
          pure ret

      let bak = file <.> "bak"

      process D "backing up file" do
        liftIO $ copyFile file bak

      process I "check predicate" do
        t <- test file
        unless (t == ExitSuccess) do
          logError "did not satisify predicate"
          liftIO exitFailure

      c <- process D "parsing file" do
        parseCFile file

      output file c

      process I "sanity checks" do
        c' <- parseCFile file
        unless (void c' == void c) do
          liftIO do
            withFile "error.1.hs" WriteMode (`pHPrint` void c)
            withFile "error.2.hs" WriteMode (`pHPrint` void c')
          logError "Outputted a different file than I read... Please report original file and error.{1,2}.hs"
          liftIO exitFailure

        let rtree = ReduceC.defaultReduceCWithKeywords keywords c
        t <- check file [] (IRTree.extract rtree)

        unless t do
          logError "file did not satisfy predicate"
          liftIO exitFailure

      when checkmode do
        logInfo "In checking mode, stopping now"
        liftIO exitSuccess

      l <-
        c & fix \rec prevc -> do
          let rtree = ReduceC.defaultReduceCWithKeywords keywords prevc
          c' <- case mode of
            Lin -> IRTree.reduce (check file) rtree
            Exp -> IRTree.reduceAdaptive (check file) rtree

          if fixpoint && (c' $> ()) /= (prevc $> ())
            then do
              logInfo "Running again until fixpoint"
              rec c'
            else pure c'

      output file l
      logInfo "Everything went well"
 where
  parseCFile file = do
    res <- liftIO $ C.parseCFilePre file
    case res of
      Right c -> pure c
      Left err -> do
        logError (Text.pack (show err))
        liftIO exitFailure