module REPL where

import Eval
import FileEval
import Lexer
import Parser
import Research
import ContentStore

import Control.Concurrent (forkIO, threadDelay, killThread, ThreadId)
import Control.Monad (forever, void, when, forM, forM_, foldM, unless)
import Data.ByteString (ByteString)
import Data.Maybe (isNothing, isJust, fromJust, catMaybes)
import Database.SQLite.Simple (Connection, Only(..), query, query_, execute, execute_, open)
import System.Directory (doesFileExist, createDirectoryIfMissing)
import System.FSNotify
import System.FilePath (takeDirectory, (</>))
import Text.Read (readMaybe)

import Control.Exception         (IOException, SomeException, catch
                                 , displayException)
import Control.Monad             (forM_)
import Control.Monad.Catch       (handle, MonadCatch)
import Control.Monad.IO.Class    (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT)
import Data.Char                 (isSpace, isUpper)
import Data.List                 ((\\), dropWhile, dropWhileEnd, isPrefixOf, nub, sortBy, groupBy, intercalate, find)
import Data.Version (showVersion)
import Paths_tricu (version)
import System.Console.Haskeline
import System.Console.ANSI (setSGR, SGR(..), ConsoleLayer(..), ColorIntensity(..), 
                           Color(..), ConsoleIntensity(..), clearFromCursorToLineEnd)

import qualified Data.Map as Map
import qualified Data.Text as T
import qualified Data.Text.IO as T

import Control.Concurrent (forkIO, threadDelay)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Time (UTCTime, getCurrentTime, diffUTCTime)
import Control.Concurrent.MVar (MVar, newMVar, putMVar, takeMVar)

import Data.Time.Format (formatTime, defaultTimeLocale)
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)

data REPLState = REPLState
  { replForm :: EvaluatedForm
  , replContentStore :: Maybe Connection
  , replWatchedFile :: Maybe FilePath
  , replSelectedVersions :: Map.Map String T.Text
  , replWatcherThread :: Maybe ThreadId
  }

repl :: IO ()
repl = do
  conn <- ContentStore.initContentStore 
  runInputT settings (withInterrupt (loop (REPLState Decode (Just conn) Nothing Map.empty Nothing)))
  where
    settings :: Settings IO
    settings = Settings
      { complete = completeWord Nothing " \t" completeCommands
      , historyFile = Just "~/.local/state/tricu/history"
      , autoAddHistory = True
      }

    completeCommands :: String -> IO [Completion]
    completeCommands str = return $ map simpleCompletion $
      filter (str `isPrefixOf`) commands
      where
        commands = [ "!exit"
                   , "!output"
                   , "!import"
                   , "!clear"
                   , "!reset"
                   , "!help"
                   , "!definitions"
                   , "!watch"
                   , "!refresh"
                   , "!versions"
                   , "!select"
                   , "!tag"
                   ]

    loop :: REPLState -> InputT IO ()
    loop state = handle (\Interrupt -> interruptHandler state Interrupt) $ do
      minput <- getInputLine "tricu < "
      case minput of
        Nothing -> return ()
        Just s
          | strip s == "" -> loop state
          | strip s == "!exit" -> outputStrLn "Exiting tricu"
          | strip s == "!clear" -> do
              liftIO $ putStr "\ESC[2J\ESC[H"
              loop state
          | strip s == "!reset" -> do
              outputStrLn "Selected versions reset"
              loop state { replSelectedVersions = Map.empty }
          | strip s == "!help" -> do
              outputStrLn $ "tricu version " ++ showVersion version
              outputStrLn "Available commands:"
              outputStrLn "  !exit        - Exit the REPL"
              outputStrLn "  !clear       - Clear the screen"
              outputStrLn "  !reset       - Reset preferences for selected versions"
              outputStrLn "  !help        - Show tricu version and available commands"
              outputStrLn "  !output      - Change output format (tree|fsl|ast|ternary|ascii|decode)"
              outputStrLn "  !definitions - List all defined terms in the content store"
              outputStrLn "  !import      - Import definitions from file to the content store"
              outputStrLn "  !watch       - Watch a file for changes, evaluate terms, and store them"
              outputStrLn "  !versions    - Show all versions of a term by name"
              outputStrLn "  !select      - Select a specific version of a term for subsequent lookups"
              outputStrLn "  !tag         - Add or update a tag for a term by hash or name"
              loop state
          | strip s == "!output" -> handleOutput state
          | strip s == "!definitions" -> handleDefinitions state
          | "!import" `isPrefixOf` strip s -> handleImport state
          | "!watch" `isPrefixOf` strip s -> handleWatch state
          | strip s == "!refresh" -> handleRefresh state
          | "!versions" `isPrefixOf` strip s -> handleVersions state
          | "!select" `isPrefixOf` strip s -> handleSelect state
          | "!tag" `isPrefixOf` strip s -> handleTag state
          | take 2 s == "--" -> loop state
          | otherwise -> do
              result <- liftIO $ catch
                (processInput state s)
                (errorHandler state)
              loop result

    handleOutput :: REPLState -> InputT IO ()
    handleOutput state = do
      let formats = [Decode, TreeCalculus, FSL, AST, Ternary, Ascii]
      outputStrLn "Available output formats:"
      mapM_ (\(i, f) -> outputStrLn $ show i ++ ". " ++ show f)
        (zip [1..] formats)

      result <- runMaybeT $ do
        input <- MaybeT $ getInputLine "Select output format (1-6) < "
        case reads input of
          [(n, "")] | n >= 1 && n <= 6 ->
            return $ formats !! (n-1)
          _ -> MaybeT $ return Nothing

      case result of
        Nothing -> do
          outputStrLn "Invalid selection. Keeping current output format."
          loop state
        Just newForm -> do
          outputStrLn $ "Output format changed to: " ++ show newForm
          loop state { replForm = newForm }

    handleDefinitions :: REPLState -> InputT IO ()
    handleDefinitions state = case replContentStore state of
      Nothing -> do
        liftIO $ printError "Content store not initialized"
        loop state
      Just conn -> do
        terms <- liftIO $ ContentStore.listStoredTerms conn
        
        if null terms
          then do
            liftIO $ printWarning "No terms in content store."
            loop state
          else do
            liftIO $ do
              printSuccess $ "Content store contains " ++ show (length terms) ++ " terms:"

            let maxNameWidth = maximum $ map (length . T.unpack . termNames) terms

            forM_ terms $ \term -> do
              let namesStr = T.unpack (termNames term)
                  hash = termHash term
                  padding = replicate (maxNameWidth - length namesStr) ' '
                  
              liftIO $ do
                putStr "  "
                printVariable namesStr
                putStr padding
                putStr " [hash: "
                displayColoredHash hash
                putStrLn "]"
                
                tags <- ContentStore.termToTags conn hash 
                unless (null tags) $ displayTags tags

            loop state

    handleImport :: REPLState -> InputT IO ()
    handleImport state = do
      let fset = setComplete completeFilename defaultSettings
      filename <- runInputT fset $ getInputLineWithInitial "File to import: " ("", "")
      case filename of
        Nothing -> loop state
        Just f -> do
          let cleanFilename = strip f
          exists <- liftIO $ doesFileExist cleanFilename
          if not exists
            then do
              liftIO $ printError $ "File not found: " ++ cleanFilename
              loop state
            else importFile state cleanFilename

    importFile :: REPLState -> String -> InputT IO ()
    importFile state cleanFilename = do
      code <- liftIO $ readFile cleanFilename
      case replContentStore state of
        Nothing -> do
          liftIO $ printError "Content store not initialized"
          loop state
        Just conn -> do
          env <- liftIO $ evaluateFile cleanFilename
          
          liftIO $ do
            printSuccess $ "Importing file: " ++ cleanFilename
            let defs = Map.toList $ Map.delete "!result" env
            
            importedCount <- foldM (\count (name, term) -> do
              hash <- ContentStore.storeTerm conn [name] term
              printSuccess $ "Stored definition: " ++ name ++ " with hash " ++ T.unpack hash
              return (count + 1)
              ) 0 defs
            
            printSuccess $ "Imported " ++ show importedCount ++ " definitions successfully"
          
          loop state

    handleWatch :: REPLState -> InputT IO ()
    handleWatch state = do
      dbPath <- liftIO ContentStore.getContentStorePath
      let filepath = takeDirectory dbPath </> "scratch.tri"
      let dirPath = takeDirectory filepath

      liftIO $ createDirectoryIfMissing True dirPath

      fileExists <- liftIO $ doesFileExist filepath
      unless fileExists $ liftIO $ writeFile filepath "-- tricu scratch file\n\n"

      outputStrLn $ "Using scratch file: " ++ filepath

      when (isJust (replWatcherThread state)) $ do
        outputStrLn "Stopping previous file watch"
        liftIO $ killThread (fromJust $ replWatcherThread state)

      outputStrLn $ "Starting to watch file: " ++ filepath
      outputStrLn "Press Ctrl+C to stop watching and return to REPL"

      liftIO $ processWatchedFile filepath (replContentStore state) (replSelectedVersions state) (replForm state)

      lastProcessedRef <- liftIO $ newIORef =<< getCurrentTime

      watcherId <- liftIO $ forkIO $ withManager $ \mgr -> do
        stopAction <- watchDir mgr dirPath (\event -> eventPath event == filepath) $ \event -> do
          now <- getCurrentTime
          lastProcessed <- readIORef lastProcessedRef
          when (diffUTCTime now lastProcessed > 0.5) $ do
            putStrLn $ "\nFile changed: " ++ filepath
            processWatchedFile filepath (replContentStore state) (replSelectedVersions state) (replForm state)
            writeIORef lastProcessedRef now
        forever $ threadDelay 1000000

      watchLoop state { replWatchedFile = Just filepath, replWatcherThread = Just watcherId }

    handleUnwatch :: REPLState -> InputT IO ()
    handleUnwatch state = case replWatchedFile state of
      Nothing -> do
        outputStrLn "No file is currently being watched"
        loop state
      Just path -> do
        outputStrLn $ "Stopped watching " ++ path
        when (isJust (replWatcherThread state)) $ do
          liftIO $ killThread (fromJust $ replWatcherThread state)
        loop state { replWatchedFile = Nothing, replWatcherThread = Nothing }

    handleRefresh :: REPLState -> InputT IO ()
    handleRefresh state = case replContentStore state of
      Nothing -> do
        outputStrLn "Content store not initialized"
        loop state
      Just conn -> do
        outputStrLn "Environment refreshed from content store (definitions are live)"
        loop state

    handleVersions :: REPLState -> InputT IO ()
    handleVersions state = case replContentStore state of
      Nothing -> do
        liftIO $ printError "Content store not initialized"
        loop state
      Just conn -> do
        liftIO $ printPrompt "Term name: "
        nameInput <- getInputLine ""
        case nameInput of
          Nothing -> loop state
          Just n -> do
            let termName = strip n
            versions <- liftIO $ ContentStore.termVersions conn termName
            if null versions
              then liftIO $ printError $ "No versions found for term: " ++ termName
              else do
                liftIO $ do
                  printKeyword "Versions of "
                  printVariable termName
                  putStrLn ":"
                  
                  forM_ (zip [1..] versions) $ \(i, (hash, _, ts)) -> do
                    tags <- ContentStore.termToTags conn hash
                    putStr $ show (i :: Int) ++ ". "
                    displayColoredHash hash
                    putStr $ " (" ++ formatTimestamp ts ++ ")"
                    unless (null tags) $ do
                      putStr " ["
                      printKeyword "Tags: "
                      forM_ (zip [0..] tags) $ \(j, tag) -> do
                        printTag (T.unpack tag)
                        when (j < length tags - 1) $ putStr ", "
                      putStr "]"
                    putStrLn ""
        loop state

    handleSelect :: REPLState -> InputT IO ()
    handleSelect state = case replContentStore state of
      Nothing -> do
        liftIO $ printError "Content store not initialized"
        loop state
      Just conn -> do
        liftIO $ printPrompt "Term name: "
        nameInput <- getInputLine ""
        case nameInput of
          Nothing -> loop state
          Just n -> do
            let cleanName = strip n
            versions <- liftIO $ ContentStore.termVersions conn cleanName
            if null versions
              then do
                liftIO $ printError $ "No versions found for term: " ++ cleanName
                loop state
              else do
                liftIO $ do
                  printKeyword "Versions of "
                  printVariable cleanName
                  putStrLn ":"
                  
                  forM_ (zip [1..] versions) $ \(i, (hash, _, ts)) -> do
                    tags <- ContentStore.termToTags conn hash
                    putStr $ show (i :: Int) ++ ". "
                    displayColoredHash hash
                    putStr $ " (" ++ formatTimestamp ts ++ ")"
                    unless (null tags) $ do
                      putStr " ["
                      printKeyword "Tags: "
                      forM_ (zip [0..] tags) $ \(j, tag) -> do
                        printTag (T.unpack tag)
                        when (j < length tags - 1) $ putStr ", "
                      putStr "]"
                    putStrLn ""

                liftIO $ printPrompt "Select version (number or full hash, Enter to cancel): "
                choiceInput <- getInputLine ""
                let choice = strip <$> choiceInput

                selectedHash <- case choice of
                    Just selectedStr | not (null selectedStr) -> do
                        case readMaybe selectedStr :: Maybe Int of
                            Just idx | idx > 0 && idx <= length versions -> do
                                let (h, _, _) = versions !! (idx - 1)
                                return $ Just h
                            _ -> do
                                let potentialHash = T.pack selectedStr
                                let foundByHash = find (\(h, _, _) -> T.isPrefixOf potentialHash h) versions
                                case foundByHash of
                                    Just (h, _, _) -> return $ Just h
                                    Nothing -> do
                                        liftIO $ printError "Invalid selection or hash not found in list."
                                        return Nothing
                    _ -> return Nothing

                case selectedHash of
                    Just hashToSelect -> do
                        let newState = state { replSelectedVersions =
                                              Map.insert cleanName hashToSelect (replSelectedVersions state) }
                        liftIO $ do
                          printSuccess "Selected version "
                          displayColoredHash hashToSelect
                          putStr " for term "
                          printVariable cleanName
                          putStrLn ""
                        loop newState
                    Nothing -> loop state

    handleTag :: REPLState -> InputT IO ()
    handleTag state = case replContentStore state of
      Nothing -> do
        liftIO $ printError "Content store not initialized"
        loop state
      Just conn -> do
        liftIO $ printPrompt "Term hash (full or prefix) or name (most recent version will be used): "
        identInput <- getInputLine ""
        case identInput of
          Nothing -> loop state
          Just ident -> do
            let cleanIdent = strip ident
            
            mFullHash <- liftIO $ resolveIdentifierToHash conn cleanIdent

            case mFullHash of
              Nothing -> do
                liftIO $ printError $ "Could not resolve identifier: " ++ cleanIdent
                loop state
              Just fullHash -> do
                liftIO $ do
                    putStr "Tagging term with hash: "
                    displayColoredHash fullHash
                    putStrLn ""
                tags <- liftIO $ ContentStore.termToTags conn fullHash
                unless (null tags) $ do
                  liftIO $ do
                    printKeyword "Existing tags:"
                    displayTags tags

                liftIO $ printPrompt "Tag to add/set: "
                tagValueInput <- getInputLine ""
                case tagValueInput of
                  Nothing -> loop state
                  Just tv -> do
                    let tagVal = T.pack (strip tv)
                    liftIO $ do
                      ContentStore.setTag conn fullHash tagVal 
                      printSuccess $ "Tag '"
                      printTag (T.unpack tagVal)
                      putStr "' set for term with hash "
                      displayColoredHash fullHash
                      putStrLn ""
                    loop state

    resolveIdentifierToHash :: Connection -> String -> IO (Maybe T.Text)
    resolveIdentifierToHash conn ident
        | T.pack "#" `T.isInfixOf` T.pack ident = do
            let hashPrefix = T.pack ident
            matchingHashes <- liftIO $ query conn "SELECT hash FROM terms WHERE hash LIKE ?" (Only (hashPrefix <> "%")) :: IO [Only T.Text]
            case matchingHashes of
                [Only fullHash] -> return $ Just fullHash
                [] -> do printError $ "No hash found starting with: " ++ T.unpack hashPrefix; return Nothing
                _  -> do printError $ "Ambiguous hash prefix: " ++ T.unpack hashPrefix; return Nothing
        | otherwise = do
            versions <- ContentStore.termVersions conn ident
            if null versions
            then do printError $ "No versions found for term name: " ++ ident; return Nothing
            else return $ Just $ (\(h,_,_) -> h) $ head versions

    interruptHandler :: REPLState -> Interrupt -> InputT IO ()
    interruptHandler state _ = do
      liftIO $ do
        printWarning "Interrupted with CTRL+C"
        printWarning "You can use the !exit command or CTRL+D to exit"
      loop state

    errorHandler :: REPLState -> SomeException -> IO REPLState
    errorHandler state e = do
      printError $ "Error: " ++ displayException e
      return state

    processInput :: REPLState -> String -> IO REPLState
    processInput state input = do
      let asts = parseTricu input
      case asts of
        [] -> return state
        _ -> case replContentStore state of
          Nothing -> do
            printError "Content store not initialized"
            return state
          Just conn -> do
            newState <- foldM (\s astNode -> do
              let varsInAst = Eval.findVarNames astNode
              foldM (\currentSelectionState varName ->
                if Map.member varName (replSelectedVersions currentSelectionState)
                  then return currentSelectionState
                  else do
                    versions <- ContentStore.termVersions conn varName
                    if length versions > 1
                      then do
                        let (latestHash, _, _) = head versions
                        liftIO $ printWarning $ "Multiple versions of '" ++ varName ++ "' found. Using most recent."
                        return currentSelectionState { replSelectedVersions = Map.insert varName latestHash (replSelectedVersions currentSelectionState) }
                      else return currentSelectionState
                ) s varsInAst
              ) state asts

            forM_ asts $ \ast -> do
              case ast of
                SDef name [] body -> do
                  result <- evalAST (Just conn) (replSelectedVersions newState) body
                  hash <- ContentStore.storeTerm conn [name] result
                  
                  liftIO $ do
                    putStr "tricu > "
                    printSuccess "Stored definition: "
                    printVariable name
                    putStr " with hash "
                    displayColoredHash hash
                    putStrLn ""
                    
                    putStr "tricu > "
                    printResult $ formatT (replForm newState) result
                    putStrLn ""

                _ -> do
                  result <- evalAST (Just conn) (replSelectedVersions newState) ast
                  liftIO $ do
                    putStr "tricu > "
                    printResult $ formatT (replForm newState) result
                    putStrLn ""
            return newState

    strip :: String -> String
    strip = dropWhileEnd isSpace . dropWhile isSpace

    watchLoop :: REPLState -> InputT IO ()
    watchLoop state = handle (\Interrupt -> do
                                outputStrLn "\nStopped watching file"
                                when (isJust (replWatcherThread state)) $ do
                                  liftIO $ killThread (fromJust $ replWatcherThread state)
                                loop state { replWatchedFile = Nothing, replWatcherThread = Nothing }) $ do
      liftIO $ threadDelay 1000000
      watchLoop state

    processWatchedFile :: FilePath -> Maybe Connection -> Map.Map String T.Text -> EvaluatedForm -> IO ()
    processWatchedFile filepath mconn selectedVersions outputForm = do
      content <- readFile filepath
      let asts = parseTricu content

      case mconn of
        Nothing -> putStrLn "Content store not initialized for watched file processing."
        Just conn -> do
          forM_ asts $ \ast -> case ast of
            SDef name [] body -> do
              result <- evalAST (Just conn) selectedVersions body
              hash <- ContentStore.storeTerm conn [name] result
              putStrLn $ "tricu > Stored definition: " ++ name ++ " with hash " ++ T.unpack hash
              putStrLn $ "tricu > " ++ name ++ " = " ++ formatT outputForm result
            _ -> do
              result <- evalAST (Just conn) selectedVersions ast
              putStrLn $ "tricu > Result: " ++ formatT outputForm result
          putStrLn $ "tricu > Processed file: " ++ filepath

    formatTimestamp :: Integer -> String
    formatTimestamp ts = formatTime defaultTimeLocale "%Y-%m-%d %H:%M:%S" (posixSecondsToUTCTime (fromIntegral ts))

    displayColoredHash :: T.Text -> IO ()
    displayColoredHash hash = do
      let (prefix, rest) = T.splitAt 16 hash
      setSGR [SetColor Foreground Vivid Cyan]
      putStr $ T.unpack prefix
      setSGR [SetColor Foreground Dull White]
      putStr $ T.unpack rest
      setSGR [Reset]

    coloredHashString :: T.Text -> String
    coloredHashString hash = 
      "\ESC[1;36m" ++ T.unpack (T.take 16 hash) ++ 
      "\ESC[0;37m" ++ T.unpack (T.drop 16 hash) ++ 
      "\ESC[0m"

    withColor :: ColorIntensity -> Color -> IO () -> IO ()
    withColor intensity color action = do
      setSGR [SetColor Foreground intensity color]
      action
      setSGR [Reset]

    printColored :: ColorIntensity -> Color -> String -> IO ()
    printColored intensity color text = withColor intensity color $ putStr text

    printlnColored :: ColorIntensity -> Color -> String -> IO ()
    printlnColored intensity color text = withColor intensity color $ putStrLn text

    printSuccess :: String -> IO ()
    printSuccess = printlnColored Vivid Green

    printError :: String -> IO ()
    printError = printlnColored Vivid Red

    printWarning :: String -> IO ()
    printWarning = printlnColored Vivid Yellow

    printPrompt :: String -> IO ()
    printPrompt = printColored Vivid Blue

    printVariable :: String -> IO ()
    printVariable = printColored Vivid Magenta

    printTag :: String -> IO ()
    printTag = printColored Vivid Yellow

    printKeyword :: String -> IO ()
    printKeyword = printColored Vivid Blue

    printResult :: String -> IO ()
    printResult = printColored Dull White

    displayTags :: [T.Text] -> IO ()
    displayTags [] = return ()
    displayTags tags = do
      putStr "    Tags: "
      forM_ (zip [0..] tags) $ \(i, tag) -> do
        printTag (T.unpack tag)
        when (i < length tags - 1) $ putStr ", "
      putStrLn ""