module Comment
  ( Comment (..),
    Point (..),
    getComments,
    extractComments,
    CommentStyle (..),
    uncomment,
    comment,
  )
where

import Comment.Language qualified as L
import Control.Applicative (liftA2)
import Control.Exception (catch)
import Data.Binary (Binary)
import Data.ByteString qualified as B
import Data.ByteString.Lazy qualified as LB
import Data.List (find, sortBy)
import Data.List.NonEmpty qualified as N
import Data.Maybe (fromMaybe)
import Data.Ord (comparing)
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Exception qualified as E
import Foreign.Marshal.Alloc (alloca, free)
import Foreign.Marshal.Array (peekArray)
import Foreign.Storable
import GHC.Generics (Generic)
import Git qualified
import System.FilePath (takeExtension)
import TreeSitter qualified as S

data Comment = Comment
  { text :: T.Text,
    language :: L.Language,
    startByte :: Int,
    endByte :: Int,
    startPoint :: Point,
    endPoint :: Point,
    filePath :: FilePath
  }
  deriving (Eq, Show)

data Point = Point
  { row :: Int,
    column :: Int
  }
  deriving (Eq, Show, Generic, Binary)

getComments :: Git.CommitHash -> FilePath -> IO [Comment]
getComments commitHash filePath =
  fmap mergeLineComments
    . extractComments filePath language
    . LB.toStrict
    =<< catch
      (Git.readTextFileOfBS commitHash filePath)
      (\(_ :: E.CannotReadFile) -> pure "")
  where
    language = L.fromPath (takeExtension filePath)

    mergeLineComments :: [Comment] -> [Comment]
    mergeLineComments =
      map mergeGroup
        . chainsBy (\a b -> a.endPoint.row + 1 == b.startPoint.row)
        . sortBy (comparing (liftA2 (,) (.startByte) (.endByte)))

    mergeGroup :: N.NonEmpty Comment -> Comment
    mergeGroup css@(c N.:| cs) =
      c
        { text = T.unlines (map (.text) (c : cs)),
          startByte = first.startByte,
          endByte = last.endByte,
          startPoint = first.startPoint,
          endPoint = last.endPoint
        }
      where
        first = N.head css
        last = N.last css

    {- A version of `Data.List.groupBy` that uses the last added group-member for comparison with new candidates for the group. `Data.List.groupBy` uses the initial member for all subsequent comparisons. -}
    chainsBy :: (a -> a -> Bool) -> [a] -> [N.NonEmpty a]
    chainsBy p = reverse . map N.reverse . go []
      where
        go rs [] = rs
        go [] (x : xs) = go [N.singleton x] xs
        go (ass@((a N.:| as) : rs)) (x : xs)
          | p a x = go ((x N.:| a : as) : rs) xs
          | otherwise = go (N.singleton x : ass) xs

extractComments :: FilePath -> L.Language -> B.ByteString -> IO [Comment]
extractComments filePath language str' =
  alloca $ \nodesPtrPtr -> do
    alloca $ \numNodesPtr -> do
      B.useAsCString str' $ \str ->
        S.extract_comments
          (L.parser language)
          str
          nodesPtrPtr
          numNodesPtr
      numNodes <- peek numNodesPtr
      nodesPtr <- peek nodesPtrPtr
      nodes <- peekArray (fromIntegral numNodes) nodesPtr
      free nodesPtr
      pure $
        map
          ( \node ->
              let startByte = fromIntegral node.startByte
                  endByte = fromIntegral node.endByte
               in Comment
                    { startPoint =
                        Point
                          { row = fromIntegral node.startPoint.row + 1,
                            column = fromIntegral node.startPoint.column + 1
                          },
                      endPoint =
                        Point
                          { row = fromIntegral node.endPoint.row + 1,
                            column = fromIntegral node.endPoint.column + 1
                          },
                      text =
                        T.decodeUtf8
                          . B.take (endByte - startByte)
                          . B.drop startByte
                          $ str',
                      ..
                    }
          )
          nodes

data CommentStyle
  = LineStyle T.Text
  | BlockStyle T.Text T.Text
  deriving (Eq, Show, Generic, Binary)

comment :: CommentStyle -> T.Text -> T.Text
comment (LineStyle linePrefix) = T.unlines . map ((linePrefix <> " ") <>) . T.lines
comment (BlockStyle blockStart blockEnd) = (blockStart <>) . (<> blockEnd)

uncomment :: L.Language -> T.Text -> (CommentStyle, T.Text)
uncomment language rawText =
  maybe
    ( ( LineStyle (L.lineStart language),
        stripLineComments (L.lineStart language) text
      )
    )
    ( \(blockStart, blockEnd) ->
        ( BlockStyle blockStart blockEnd,
          stripBlockComment blockStart blockEnd text
        )
    )
    $ do
      (blockStarts, blockEnd) <- L.block language
      (,blockEnd) <$> find (`T.isPrefixOf` text) blockStarts
  where
    text = stripLines rawText
    stripLines = T.intercalate "\n" . map T.strip . T.lines

stripLineComments :: T.Text -> T.Text -> T.Text
stripLineComments lineStart text =
  onLines
    ( \line ->
        fromMaybe line . fmap T.stripStart $
          T.stripPrefix lineStart line
    )
    text
  where
    onLines f = T.intercalate "\n" . map f . T.lines

stripBlockComment :: T.Text -> T.Text -> T.Text -> T.Text
stripBlockComment blockStart blockEnd text =
  T.strip
    . (fromMaybe text . T.stripSuffix blockEnd)
    . (fromMaybe text . T.stripPrefix blockStart)
    $ text