content: module AnsiDiff (
      ansidiff,
  ) where
  
  import Control.Monad (forM_)
  import Control.Monad.ST (ST)
  import Control.Monad.Trans.State.Strict qualified as S
  import Data.Algorithm.Diff qualified as Diff
  import Data.Primitive.Array qualified as P
  import Data.Primitive.PrimArray qualified as P
  import Data.Primitive.Types (Prim)
  import System.Console.ANSI qualified as ANSI
  import Text.EditDistance qualified as ED
  
  ansidiff :: String -> String -> String
  ansidiff old new = runStringWriter $ do
      writeStr ansiReset
      let linesDiff = Diff.getDiff (lines old) (lines new)
      forM_ (chunks linesDiff) $ \case
          Same xs _  -> forM_ xs $ \x -> writeStrLn $ ' ' : x
          Diff xs ys -> ansiDiff' xs ys
  
  ansiDiff' :: [String] -> [String] -> StringWriter ()
  ansiDiff' old new = do
      -- removed
      forM_ diffs $ \case
          Del x  -> writeStrLn $ ansiRed ++ '-' : x ++ ansiReset
          Ins _  -> return ()
          Swp ds -> removedLine ds
  
      -- added
      forM_ diffs $ \case
          Del _  -> return ()
          Ins y  -> writeStrLn $ ansiGreen ++ '+' : y ++ ansiReset
          Swp ds -> addedLine ds
  
    where
      diffs :: [D Char Char]
      diffs = alignment old new
  
  -- empasise line diffs only if there are similar enough.
  -- this is different from 'metric'
  --   we emphasise a line when there are over half of shorter line is preserved.
  predicate :: Double -> Bool
  predicate d = d < 0.2
  
  -------------------------------------------------------------------------------
  -- Alignment
  -------------------------------------------------------------------------------
  
  alignment :: forall a b. (a ~ Char, b ~ Char) => [String] -> [String] -> [D Char Char]
  alignment xs' ys' = do
      reverse $ walk xn yn
    where
      xs :: P.Array [a]
      xs = P.arrayFromList xs'
  
      ys :: P.Array [b]
      ys = P.arrayFromList ys'
  
      xn :: Int
      xn = length xs
  
      yn :: Int
      yn = length ys
  
      diffs :: P.PrimArray Double
      diffs = createPrimArray (xn * yn) $ \arr -> do
          forM_ [0..xn-1] $ \i -> forM_ [0..yn-1] $ \j -> do
              let xc = P.indexArray xs i
              let yc = P.indexArray ys j
              let xm = length xc
              let ym = length yc
              let mi = min xm ym
              let ma = max xm ym
  
              let r :: Double
                  !r | ma == 0 = 0
                     | mi == 0 = 1
  
                      -- a lower bound of edit-distance,
                      -- if it's not satisfying a predicate for inter-line diff,
                      -- just use it instead of calculating proper edit-distance
                      -- (see tests)
                     | let lb = fromIntegral (ma - mi) / fromIntegral ma
                     , not (predicate lb) = lb
  
                      -- here we use edit-distance, because all the strings we are comparing are different,
                      -- potentially a lot. Thus we care about worse case performance i.e. O (nm) is the same as O(nd).
                     | otherwise = fromIntegral (ED.levenshteinDistance ED.defaultEditCosts xc yc) / fromIntegral ma
  
              P.writePrimArray arr (i + j * xn) r
  
      indexDiff :: Int -> Int -> Double
      indexDiff i j = P.indexPrimArray diffs (i + j * xn)
      {-# INLINE indexDiff #-}
  
      distances :: P.PrimArray Double
      distances = createPrimArray ((xn + 1) * (yn + 1)) $ \arr -> do
          forM_ [0..xn] $ \i -> do
              P.writePrimArray arr i (fromIntegral i)
  
          forM_ [1..yn] $ \j -> do
              P.writePrimArray arr (j * (xn + 1)) (fromIntegral j)
  
              forM_ [1..xn] $ \i -> do
                  let r = indexDiff (i - 1) (j - 1)
  
                  if r == 0
                  then do
                      !z <- P.readPrimArray arr ((i - 1) + (j - 1) * (xn + 1))
                      P.writePrimArray arr (i + j * (xn + 1)) z
                  else do
                      !x <- P.readPrimArray arr ((i    ) + (j - 1) * (xn + 1))
                      !y <- P.readPrimArray arr ((i - 1) + (j    ) * (xn + 1))
                      !z <- P.readPrimArray arr ((i - 1) + (j - 1) * (xn + 1))
  
                      P.writePrimArray arr (i + j * (xn + 1)) $!
                          min (r + z) $ min (x + 1) (y + 1)
  
      indexDistance :: Int -> Int -> Double
      indexDistance i j = P.indexPrimArray distances (i + j * (xn + 1))
      {-# INLINE indexDistance #-}
  
      -- trace back walk
      walk :: Int -> Int -> [D a b]
      walk i j
          | i == 0, j == 0
          = []
  
          | i == 0
          = Ins (P.indexArray ys (j - 1)) : walk i (j - 1)
  
          | j == 0
          = Del (P.indexArray xs (i - 1)) : walk (i - 1) j
  
          | otherwise
          = do
              let !x = indexDistance (i    ) (j - 1)
              let !y = indexDistance (i - 1) (j    )
              let !z = indexDistance (i - 1) (j - 1)
  
              if z <= x
                  then
                      if z <= y
                      -- z is smallest
                      then do
                          let xc = P.indexArray xs (i - 1)
                          let yc = P.indexArray ys (j - 1)
                          let !r = indexDiff (i - 1) (j - 1)
                          if predicate r
                          then do
                              let !d = Diff.getDiff xc yc
                              Swp (chunks d) : walk (i - 1) (j - 1)
                          else
                              Del xc : Ins yc : walk (i - 1) (j - 1)
  
                      -- y is smallest
                      else Del (P.indexArray xs (i - 1)) : walk (i - 1) j
  
                  else
                      if y <= x
  
                      -- y is smallest
                      then Del (P.indexArray xs (i - 1)) : walk (i - 1) j
  
                      -- x is smallest
                      else Ins (P.indexArray ys (j - 1)) : walk i (j - 1)
  
  data D a b
      = Del [a]
      | Ins [b]
      | Swp [Chunk a b]
    deriving Show
  
  -------------------------------------------------------------------------------
  -- primitive extras
  -------------------------------------------------------------------------------
  
  createPrimArray :: Prim a => Int -> (forall s. P.MutablePrimArray s a -> ST s ()) -> P.PrimArray a
  createPrimArray n action = P.runPrimArray $ do
      arr <- P.newPrimArray n
      action arr
      return arr
  
  -------------------------------------------------------------------------------
  -- Printing
  -------------------------------------------------------------------------------
  
  removedLine :: [Chunk Char Char] -> StringWriter ()
  removedLine ds = do
      writeStr ansiRed
      writeChar '-'
  
      forM_ ds $ \case
          Diff s _ -> do
              writeStr ansiRED
              writeStr s
              writeStr ansiReset
          Same s _ -> do
              writeStr ansiRed
              writeStr s
              writeStr ansiReset
  
      writeChar '\n'
  
  addedLine :: [Chunk Char Char] -> StringWriter ()
  addedLine ds = do
      writeStr ansiGreen
      writeChar '+'
  
      forM_ ds $ \case
          Diff _ s -> do
              writeStr ansiGREEN
              writeStr s
              writeStr ansiReset
          Same _ s -> do
              writeStr ansiGreen
              writeStr s
              writeStr ansiReset
  
      writeChar '\n'
  
  -------------------------------------------------------------------------------
  -- String Writer
  -------------------------------------------------------------------------------
  
  type StringWriter = S.State ShowS
  
  runStringWriter :: StringWriter () -> String
  runStringWriter m = S.execState m id ""
  
  writeChar :: Char -> StringWriter ()
  writeChar c = S.modify' $ \s -> s . (c :)
  
  writeStr :: String -> StringWriter ()
  writeStr str = S.modify' $ \s -> s . (str ++)
  
  writeStrLn :: String -> StringWriter ()
  writeStrLn str = S.modify' $ \s -> s . (str ++) . ('\n' :)
  
  -------------------------------------------------------------------------------
  -- Chunks
  -------------------------------------------------------------------------------
  
  -- | Diff chunks
  data Chunk a b
      = Same [a] [b]
      | Diff [a] [b]
    deriving Show
  
  chunks :: [Diff.PolyDiff a b] -> [Chunk a b]
  chunks [] = []
  chunks (Diff.Both x y : ds) = consS (chunks ds) where
      consS (Same xs ys : ds') = Same (x : xs) (y : ys) : ds'
      consS ds'                = Same [x] [y] : ds'
  chunks (Diff.First x  : ds) = consL (chunks ds) where
      consL (Diff xs ys : ds') = Diff (x : xs) ys : ds'
      consL ds'                = Diff [x] [] : ds'
  chunks (Diff.Second y : ds) = consR (chunks ds) where
      consR (Diff xs ys : ds') = Diff xs (y : ys) : ds'
      consR ds'                = Diff [] [y] : ds'
  
  -------------------------------------------------------------------------------
  -- ANSI helpers
  -------------------------------------------------------------------------------
  
  ansiRed :: String
  ansiRed = ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Vivid ANSI.Red]
  
  ansiRED :: String
  ansiRED = ANSI.setSGRCode [ANSI.SetColor ANSI.Background ANSI.Dull ANSI.Red]
  
  ansiGreen :: String
  ansiGreen = ANSI.setSGRCode [ANSI.SetColor ANSI.Foreground ANSI.Vivid ANSI.Green]
  
  ansiGREEN :: String
  ansiGREEN = ANSI.setSGRCode [ANSI.SetColor ANSI.Background ANSI.Dull ANSI.Green]
  
  ansiReset :: String
  ansiReset = ANSI.setSGRCode [ANSI.Reset]
  