module TreeGrepper.Match
  ( Match (..),
    Position (..),
    merge,
  )
where

import Data.Aeson (FromJSON)
import Data.Binary (Binary)
import Data.Function (on)
import Data.List (sortBy)
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.List.NonEmpty qualified as N
import Data.Ord (comparing)
import Data.Text (Text)
import Data.Text qualified as T
import GHC.Generics (Generic)

data Match = Match
  { kind :: String,
    name :: String,
    text :: Text,
    start :: Position,
    end :: Position
  }
  deriving (Show, Generic)

instance FromJSON Match

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

instance Ord Position where
  compare = compare `on` (\p -> (p.row, p.column))

instance FromJSON Position

{- tree-grepper (tree-sitter) is unable to match sibling comments blocks as a whole. We thus merge matches if they are line-adjacent. -}
merge :: [Match] -> [Match]
merge matches =
  map mergeGroup
    . chainsBy (\a b -> a.end.row + 1 == b.start.row)
    $ sortBy (comparing (.start)) matches

mergeGroup :: NonEmpty Match -> Match
mergeGroup (m :| ms) =
  m
    { text = text,
      start = start,
      end = end
    }
  where
    mss = m : ms
    text = T.unlines $ map (.text) mss
    start = minimum $ map (.start) mss
    end = maximum $ map (.end) mss

{- 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] -> [NonEmpty a]
chainsBy p = reverse . map N.reverse . go []
  where
    go rs [] = rs
    go [] (x : xs) = go [N.singleton x] xs
    go (ass@((a :| as) : rs)) (x : xs)
      | p a x = go ((x :| a : as) : rs) xs
      | otherwise = go (N.singleton x : ass) xs