content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE QualifiedDo #-}
  {-# LANGUAGE RecordWildCards #-}
  
  -- | This module allows us to diff two 'Text' values.
  module Ormolu.Diff.Text
    ( TextDiff,
      diffText,
      selectSpans,
      printTextDiff,
    )
  where
  
  import Control.Monad (unless, when)
  import Data.Algorithm.Diff qualified as D
  import Data.Foldable (for_)
  import Data.IntSet (IntSet)
  import Data.IntSet qualified as IntSet
  import Data.Maybe (listToMaybe)
  import Data.Text (Text)
  import Data.Text qualified as T
  import GHC.Types.SrcLoc
  import Ormolu.Terminal
  import Ormolu.Terminal.QualifiedDo qualified as Term
  #if !MIN_VERSION_base(4,20,0)
  import Data.List (foldl')
  #endif
  
  ----------------------------------------------------------------------------
  -- Types
  
  -- | Result of diffing two 'Text's.
  data TextDiff = TextDiff
    { -- | Path to the file that is being diffed
      textDiffPath :: FilePath,
      -- | The list of differences
      textDiffDiffList :: DiffList,
      -- | Selected lines. Only hunks that contain selected lines will be
      -- displayed, unless 'textDiffSelectedLines' is empty, in which case the
      -- whole diff will be displayed.
      textDiffSelectedLines :: IntSet
    }
    deriving (Eq, Show)
  
  -- | List of lines tagged by 'D.Both', 'D.First', or 'D.Second'.
  type DiffList = [D.Diff [Text]]
  
  -- | Similar to 'DiffList', but with line numbers assigned.
  type DiffList' = [D.Diff [(Int, Int, Text)]]
  
  -- | Diff hunk.
  data Hunk = Hunk
    { hunkFirstStartLine :: Int,
      hunkFirstLength :: Int,
      hunkSecondStartLine :: Int,
      hunkSecondLength :: Int,
      hunkDiff :: DiffList
    }
    deriving (Show)
  
  ----------------------------------------------------------------------------
  -- API
  
  -- | Diff two texts and produce a 'TextDiff'.
  diffText ::
    -- | Text before
    Text ->
    -- | Text after
    Text ->
    -- | Path to use
    FilePath ->
    -- | The resulting diff or 'Nothing' if the inputs are identical
    Maybe TextDiff
  diffText a b path =
    if all isBoth xs
      then Nothing
      else
        Just
          TextDiff
            { textDiffPath = path,
              textDiffDiffList = xs,
              textDiffSelectedLines = IntSet.empty
            }
    where
      xs = D.getGroupedDiff (lines' a) (lines' b)
      isBoth = \case
        D.Both _ _ -> True
        D.First _ -> False
        D.Second _ -> False
      -- T.lines ignores trailing blank lines
      lines' = T.splitOn "\n"
  
  -- | Select certain spans in the diff (line numbers are interpreted as
  -- belonging to the “before” state). Only selected spans will be printed.
  selectSpans :: [RealSrcSpan] -> TextDiff -> TextDiff
  selectSpans ss textDiff = textDiff {textDiffSelectedLines = xs}
    where
      xs = foldl' addOneSpan (textDiffSelectedLines textDiff) ss
      addOneSpan linesSoFar s =
        let start = srcSpanStartLine s
            end = srcSpanEndLine s
         in IntSet.union
              linesSoFar
              (IntSet.fromAscList [start .. end])
  
  -- | Print the given 'TextDiff' as a 'Term' action. This function tries to
  -- mimic the style of @git diff@.
  printTextDiff :: TextDiff -> Term
  printTextDiff TextDiff {..} = Term.do
    (bold . put . T.pack) textDiffPath
    newline
    for_ (toHunks (assignLines textDiffDiffList)) $ \hunk@Hunk {..} ->
      when (isSelectedLine textDiffSelectedLines hunk) $ Term.do
        cyan $ Term.do
          put "@@ -"
          putShow hunkFirstStartLine
          put ","
          putShow hunkFirstLength
          put " +"
          putShow hunkSecondStartLine
          put ","
          putShow hunkSecondLength
          put " @@"
        newline
        for_ hunkDiff $ \case
          D.Both ys _ ->
            for_ ys $ \y -> Term.do
              unless (T.null y) $
                put "  "
              put y
              newline
          D.First ys ->
            for_ ys $ \y -> red $ Term.do
              put "-"
              unless (T.null y) $
                put " "
              put y
              newline
          D.Second ys ->
            for_ ys $ \y -> green $ Term.do
              put "+"
              unless (T.null y) $
                put " "
              put y
              newline
  
  ----------------------------------------------------------------------------
  -- Helpers
  
  -- | Assign lines to a 'DiffList'.
  assignLines :: DiffList -> DiffList'
  assignLines = go 1 1 id
    where
      go _ _ acc [] = acc []
      go !firstLine !secondLine acc (x : xs) =
        case x of
          D.Both a b ->
            let firstInc = length a
                secondInc = length b
                a' =
                  zip3
                    (iterate (+ 1) firstLine)
                    (iterate (+ 1) secondLine)
                    a
             in go
                  (firstLine + firstInc)
                  (secondLine + secondInc)
                  (acc . (D.Both a' a' :))
                  xs
          D.First a ->
            let firstInc = length a
                a' =
                  zip3
                    (iterate (+ 1) firstLine)
                    (repeat secondLine)
                    a
             in go
                  (firstLine + firstInc)
                  secondLine
                  (acc . (D.First a' :))
                  xs
          D.Second b ->
            let secondInc = length b
                b' =
                  zip3
                    (repeat firstLine)
                    (iterate (+ 1) secondLine)
                    b
             in go
                  firstLine
                  (secondLine + secondInc)
                  (acc . (D.Second b' :))
                  xs
  
  -- | Form 'Hunk's from a 'DiffList''.
  toHunks :: DiffList' -> [Hunk]
  toHunks = go 0 False id id []
    where
      -- How many lines of context (that is, lines present in both texts) to
      -- show per hunk.
      margin = 3
      go ::
        Int ->
        Bool ->
        ([Hunk] -> [Hunk]) ->
        (DiffList' -> DiffList') ->
        [(Int, Int, Text)] ->
        DiffList' ->
        [Hunk]
      go !n gotChanges hunksAcc currentAcc bothHistory = \case
        [] ->
          if gotChanges
            then
              let currentAcc' = addBothAfter p currentAcc
                  p = take margin (reverse bothHistory)
               in case formHunk (currentAcc' []) of
                    Nothing -> hunksAcc []
                    Just hunk -> hunksAcc [hunk]
            else hunksAcc []
        (x : xs) ->
          case x of
            D.Both a _ ->
              let currentAcc' = addBothAfter p currentAcc
                  p = reverse (drop (n' - margin) bothHistory')
                  hunksAcc' =
                    case formHunk (currentAcc' []) of
                      Nothing -> hunksAcc
                      Just hunk -> hunksAcc . (hunk :)
                  bothHistory' = reverse a ++ bothHistory
                  lena = length a
                  n' = n + lena
               in if gotChanges && n' > margin * 2
                    then go 0 False hunksAcc' id bothHistory' xs
                    else go n' gotChanges hunksAcc currentAcc bothHistory' xs
            piece ->
              if gotChanges
                then
                  let currentAcc' = currentAcc . addBothBefore p (piece :)
                      p = reverse bothHistory
                   in go 0 True hunksAcc currentAcc' [] xs
                else
                  let currentAcc' = addBothBefore p (piece :)
                      p = reverse (take margin bothHistory)
                   in go 0 True hunksAcc currentAcc' [] xs
      addBothBefore [] acc = acc
      addBothBefore p acc = (D.Both p p :) . acc
      addBothAfter [] acc = acc
      addBothAfter p acc = acc . (D.Both p p :)
  
  -- | Form a 'Hunk'.
  formHunk :: DiffList' -> Maybe Hunk
  formHunk xsRaw = do
    let xs = trimEmpty xsRaw
    hunkFirstStartLine <- listToMaybe xs >>= firstStartLine
    let hunkFirstLength = firstLength xs
    hunkSecondStartLine <- listToMaybe xs >>= secondStartLine
    let hunkSecondLength = secondLength xs
        hunkDiff = mapDiff (fmap third) xs
    return Hunk {..}
  
  -- | Trim empty “both” lines from beginning and end of a 'DiffList''.
  trimEmpty :: DiffList' -> DiffList'
  trimEmpty = go True id
    where
      go isFirst acc = \case
        [] -> acc []
        [D.Both x _] ->
          let x' = reverse $ dropWhile (T.null . third) (reverse x)
           in go False (acc . (D.Both x' x' :)) []
        (D.Both x _ : xs) ->
          let x' = dropWhile (T.null . third) x
           in if isFirst
                then go False (acc . (D.Both x' x' :)) xs
                else go False (acc . (D.Both x x :)) xs
        (x : xs) ->
          go False (acc . (x :)) xs
  
  firstStartLine :: D.Diff [(Int, Int, a)] -> Maybe Int
  firstStartLine = \case
    D.Both ((x, _, _) : _) _ -> Just x
    D.First ((x, _, _) : _) -> Just x
    D.Second ((x, _, _) : _) -> Just x
    _ -> Nothing
  
  firstLength :: [D.Diff [(Int, Int, a)]] -> Int
  firstLength = go 0
    where
      go n [] = n
      go !n (x : xs) = case x of
        D.Both as _ -> go (n + length as) xs
        D.First as -> go (n + length as) xs
        D.Second _ -> go n xs
  
  secondStartLine :: D.Diff [(Int, Int, a)] -> Maybe Int
  secondStartLine = \case
    D.Both ((_, x, _) : _) _ -> Just x
    D.First ((_, x, _) : _) -> Just x
    D.Second ((_, x, _) : _) -> Just x
    _ -> Nothing
  
  secondLength :: [D.Diff [(Int, Int, a)]] -> Int
  secondLength = go 0
    where
      go n [] = n
      go !n (x : xs) = case x of
        D.Both as _ -> go (n + length as) xs
        D.First _ -> go n xs
        D.Second as -> go (n + length as) xs
  
  mapDiff :: (a -> b) -> [D.Diff a] -> [D.Diff b]
  mapDiff f = fmap $ \case
    D.Both a b -> D.Both (f a) (f b)
    D.First a -> D.First (f a)
    D.Second b -> D.Second (f b)
  
  third :: (Int, Int, Text) -> Text
  third (_, _, x) = x
  
  isSelectedLine :: IntSet -> Hunk -> Bool
  isSelectedLine selected Hunk {..} =
    -- If the set of selected lines is empty, everything is selected.
    IntSet.null selected
      || not (IntSet.disjoint selected hunkOriginalLines)
    where
      hunkOriginalLines =
        IntSet.fromAscList (take hunkFirstLength [hunkFirstStartLine ..])
  