{-# LANGUAGE OverloadedRecordDot #-}

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

import Data.Aeson (FromJSON)
import Data.Function (on)
import Data.List (sortBy)
import Data.Ord (comparing)
import GHC.Generics (Generic)

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

instance FromJSON Match

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

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
    . groupBy (\a b -> a.end.row + 1 == b.start.row)
    $ sortBy (comparing (.start)) matches

mergeGroup :: [Match] -> Match
mergeGroup (m : ms) =
  m
    { text = text,
      start = start,
      end = end
    }
  where
    mss = m : ms
    text = 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. -}
groupBy :: (a -> a -> Bool) -> [a] -> [[a]]
groupBy p xs = reverse . map reverse $ go [] p xs
  where
    go rs _ [] = rs
    go [] p (x : xs) = go [[x]] p xs
    go (ass@((a : as) : rs)) p (x : xs)
      | p a x = go ((x : a : as) : rs) p xs
      | otherwise = go ([x] : ass) p xs