Algorithm.hs

{-# LANGUAGE OverloadedStrings #-}

-- | ProjectedContext layer: diff/patch algorithm implementations.
-- Contains the Myers diff algorithm, LCS, hunk building, hunk application,
-- and fuzzy patching helpers.
-- All functions are pure; no IO is performed in this module.
module AIAgent.DiffPatch.Unified.ProjectedContext.Algorithm
  ( lcs
  , myers
  , myersSearch
  , buildEdits
  , buildHunks
  , applyHunk
  , matchLines
  , buildReplacement
  , generateSearchIndices
  , findFirstMatch
  , findNearestMismatch
  , applyHunkFuzzy
  , Frontier
  ) where

import qualified Data.Text as T
import Data.Text (Text)
import Data.Array (listArray, array, (!))
import AIAgent.DiffPatch.Unified.CoreModel.Type
import AIAgent.DiffPatch.Unified.CoreModel.Constant

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU


-- | Compute an edit script from two lists of lines using the LCS algorithm.
-- Each element of the result indicates whether a line was kept, deleted, or inserted.
-- Uses dynamic programming to build an LCS table, then backtracks to produce the edit list.
lcs :: [Text] -> [Text] -> [Edit]
lcs old new = backtrack m n []
  where
    m    = length old
    n    = length new
    oldV = listArray (1, max 1 m) old
    newV = listArray (1, max 1 n) new

    -- DP table: tbl ! (i,j) = LCS length of old[1..i] and new[1..j]
    tbl = array ((0, 0), (m, n))
            [ ((i, j), cell i j) | i <- [0..m], j <- [0..n] ]

    cell 0 _ = (0 :: Int)
    cell _ 0 = 0
    cell i j
      | oldV ! i == newV ! j = tbl ! (i-1, j-1) + 1
      | otherwise            = max (tbl ! (i-1, j)) (tbl ! (i, j-1))

    -- Backtrack using an accumulator to build the list in O(n) time.
    -- To ensure the [Delete, Insert] order, we prioritize the Insert branch
    -- when scores are tied, pushing Insert further towards the end of the final list.
    backtrack i j acc
      | i == 0 && j == 0 = acc

      -- Match: No change
      | i > 0 && j > 0 && oldV ! i == newV ! j =
          backtrack (i-1) (j-1) (Keep (oldV ! i) : acc)

      -- Prefer Insert branch when scores are tied to maintain
      -- the conventional "Delete before Insert" ordering in the final result.
      | j > 0 && (i == 0 || tbl ! (i, j-1) >= tbl ! (i-1, j)) =
          backtrack i (j-1) (Insert (newV ! j) : acc)

      -- Delete branch
      | i > 0 =
          backtrack (i-1) j (Delete (oldV ! i) : acc)

      | otherwise = acc


-- ---------------------------------------------------------------------------
-- Myers diff
-- ---------------------------------------------------------------------------

-- | Unboxed vector of Int used to represent the furthest-reaching x position
-- on each diagonal k.  Element at logical diagonal k is stored at physical
-- index k + offset, where offset = m + n.
type Frontier = VU.Vector Int

-- | Compute an edit script from two lists of lines using the Myers diff algorithm.
-- Implements the O(ND) algorithm from Myers (1986):
--   "An O(ND) Difference Algorithm and Its Variations"
--
-- Key concepts:
--   x        : position in old (0-based, lines consumed so far, 0..m)
--   y        : position in new (0-based, lines consumed so far, 0..n)
--   k = x-y  : diagonal index
--   v[k]     : furthest x reached on diagonal k for a given edit count d
--   snake    : free diagonal advances where old[x] == new[y]
--   snapshot : immutable copy of v after each d-round, used for backtracking
--
-- CRITICAL INVARIANT: when computing the new frontier v' for edit count d,
-- all reads of v[k-1] and v[k+1] must come from the *previous* v (before any
-- updates for this d-round).  This is enforced by computing all (k, x) pairs
-- first, then bulk-writing them into v'.
myers :: [Text] -> [Text] -> [Edit]
myers []  []  = []
myers old []  = map Delete old
myers []  new = map Insert new
myers oldL newL = buildEdits m n old new (myersSearch m n old new)
  where
    -- Convert lists to vectors for O(1) indexed access.
    old = V.fromList oldL
    new = V.fromList newL
    m   = V.length old
    n   = V.length new

-- | Forward phase of the Myers algorithm.
-- Iterates d = 0, 1, ..., m+n.  For each d, computes the furthest x reachable
-- on every active diagonal k, reading exclusively from the previous round's
-- frontier (not the partially-updated one).  Appends a snapshot of the updated
-- frontier after each round.  Stops at the first d for which the endpoint
-- (x=m, y=n) is reached.
myersSearch :: Int -> Int -> V.Vector Text -> V.Vector Text -> [Frontier]
myersSearch m n old new = go 0 v0 []
  where
    offset = m + n
    size   = 2 * offset + 1
    v0     = VU.replicate size 0

    -- Advance diagonally from 0-based position (x, y) as long as lines match.
    snake x y
      | x < m && y < n && (old V.! x) == (new V.! y) = snake (x+1) (y+1)
      | otherwise = x

    -- Compute the new x for diagonal k during edit round d,
    -- reading exclusively from the previous frontier v.
    -- Boundary rules:
    --   k == -d : leftmost diagonal, must come from above (Insert)
    --   k ==  d : rightmost diagonal, must come from below (Delete)
    --   tie     : prefer Insert (xIns >= xDel) to match the backtracking condition
    newX v d k =
      let xDel = (v VU.! (k - 1 + offset)) + 1
          xIns = (v VU.! (k + 1 + offset))
          x0 | k == -d      = xIns
             | k ==  d      = xDel
             | xIns >= xDel = xIns
             | otherwise    = xDel
      in  snake x0 (x0 - k)

    go d v acc =
      let -- Compute all new x values from the current (old) frontier.
          updates = [(k + offset, newX v d k) | k <- [-d, -d+2 .. d]]
          -- Bulk-apply updates to produce the new frontier.
          v'      = v VU.// updates
          acc'    = v' : acc
          xEnd    = v' VU.! (m - n + offset)
          yEnd    = xEnd - (m - n)
      in  if xEnd >= m && yEnd >= n
            then acc'
            else go (d+1) v' acc'

-- | Backward (trace) phase of the Myers algorithm.
-- Reconstructs the edit script by walking from (m, n) back to (0, 0) through
-- the frontier snapshots in reverse order.  At each step, recovers the Delete
-- or Insert move that was taken, plus any preceding snake (Keep) moves.
-- Returns the edit script in correct forward order.
buildEdits :: Int -> Int -> V.Vector Text -> V.Vector Text -> [Frontier] -> [Edit]
buildEdits m n old new snaps = go m n snaps []
  where
    offset = m + n

    -- Walk backward along a snake: while old[x-1]==new[y-1] it was a Keep move.
    unsnake x y acc
      | x > 0 && y > 0 && (old V.! (x-1)) == (new V.! (y-1)) =
          unsnake (x-1) (y-1) (Keep (old V.! (x-1)) : acc)
      | otherwise = (x, y, acc)

    -- Base: reached origin.
    go 0 0 _ acc = acc
    -- d=0 case: the entire path was a snake from the origin.
    go x y [_] acc =
      let (_, _, acc') = unsnake x y acc in acc'
    -- General case: recover one edit step, then recurse with prevV.
    go x y (_v:prevV:rest) acc =
      let (xMid, yMid, accWithSnakes) = unsnake x y acc
          k = xMid - yMid
          d = length rest + 1

          xFromDel = (prevV VU.! (k - 1 + offset)) + 1
          xFromIns = (prevV VU.! (k + 1 + offset))

          -- Mirror the tie-breaking rule used in newX (forward phase).
          cameFromInsert
            | k == -d   = True
            | k ==  d   = False
            | otherwise = xFromIns >= xFromDel

          (xPrev, yPrev, edit) =
            if cameFromInsert
              then (xMid, yMid - 1, Insert (new V.! (yMid - 1)))
              else (xMid - 1, yMid, Delete (old V.! (xMid - 1)))
      in go xPrev yPrev (prevV:rest) (edit : accWithSnakes)
    -- Unreachable: myersSearch guarantees snapshots cover all steps.
    go _ _ [] acc = acc


-- ---------------------------------------------------------------------------
-- buildHunks
-- ---------------------------------------------------------------------------

-- | Build a list of hunks from an edit script.
-- The Int argument specifies the number of unchanged context lines surrounding each hunk.
-- Adjacent change groups separated by 2*n or fewer Keep lines are merged into one hunk.
-- Returns an empty list when the edit script contains no Delete or Insert.
buildHunks :: Int -> [Edit] -> [Hunk]
buildHunks ctx edits
  | null edits = []
  | otherwise  = map makeHunk $ mergeGroups ctx $ groupEdits ctx numberedEdits
  where
    -- Annotate each Edit with 1-based old and new line numbers.
    numberedEdits = numberEdits 1 1 edits

    -- Assign (oldLine, newLine) to each Edit.
    -- Keep and Delete advance oldLine; Keep and Insert advance newLine.
    numberEdits :: Int -> Int -> [Edit] -> [(Int, Int, Edit)]
    numberEdits _  _  []                 = []
    numberEdits ol nl (e@(Keep   _):es)  = (ol, nl, e) : numberEdits (ol+1) (nl+1) es
    numberEdits ol nl (e@(Delete _):es)  = (ol, nl, e) : numberEdits (ol+1) nl     es
    numberEdits ol nl (e@(Insert _):es)  = (ol, nl, e) : numberEdits ol     (nl+1) es

    -- A "group" is a window of numbered edits that will become one hunk.
    -- It is represented as a flat list of (oldLine, newLine, Edit).

    -- Split the numbered edit list into raw change groups, each surrounded by
    -- up to ctx context lines.
    groupEdits :: Int -> [(Int, Int, Edit)] -> [[(Int, Int, Edit)]]
    groupEdits n numbered =
      let -- Locate all indices of Delete/Insert edits.
          changeIdxs = [ i | (i, (_, _, e)) <- zip [0..] numbered, isChange e ]
      in  if null changeIdxs
            then []
            else buildGroups n numbered changeIdxs

    isChange :: Edit -> Bool
    isChange (Keep _) = False
    isChange _        = True

    -- Collect groups, merging when inter-group Keep gap <= 2*ctx.
    buildGroups :: Int -> [(Int, Int, Edit)] -> [Int] -> [[(Int, Int, Edit)]]
    buildGroups n numbered idxs = go clusters
      where
        total = length numbered

        -- Cluster change indices into spans [firstChange..lastChange],
        -- merging spans whose gap is <= 2*n.
        clusterIdxs :: [Int] -> [(Int, Int)]
        clusterIdxs []     = []
        clusterIdxs (i:is) = merge i i is
          where
            merge lo hi []     = [(lo, hi)]
            merge lo hi (x:xs)
              | x - hi <= 2 * n + 1 = merge lo x xs
              | otherwise           = (lo, hi) : merge x x xs

        clusters = clusterIdxs idxs

        -- Expand each cluster span to include ctx context lines on each side.
        go [] = []
        go ((lo, hi):cs) =
          let start = max 0 (lo - n)
              end   = min (total - 1) (hi + n)
          in  slice start end numbered : go cs

    -- Extract elements from index start to end (inclusive).
    slice :: Int -> Int -> [a] -> [a]
    slice s e xs = take (e - s + 1) (drop s xs)

    -- Build a Hunk from a group of numbered edits.
    makeHunk :: [(Int, Int, Edit)] -> Hunk
    makeHunk [] = Hunk 0 0 0 0 []
    makeHunk group =
      let (oldStart, newStart, _) = head group
          oldCount = length [ () | (_, _, e) <- group, isOldLine e ]
          newCount = length [ () | (_, _, e) <- group, isNewLine e ]
          hlines   = map toLine group
      in  Hunk oldStart oldCount newStart newCount hlines

    isOldLine :: Edit -> Bool
    isOldLine (Insert _) = False
    isOldLine _          = True

    isNewLine :: Edit -> Bool
    isNewLine (Delete _) = False
    isNewLine _          = True

    toLine :: (Int, Int, Edit) -> Line
    toLine (_, _, Keep   t) = Line Context t
    toLine (_, _, Delete t) = Line Removed t
    toLine (_, _, Insert t) = Line Added   t

    -- Merge adjacent groups when their Keep gap falls within 2*ctx.
    -- groupEdits already handles merging via clusterIdxs; this is a no-op
    -- pass-through kept for symmetry.
    mergeGroups :: Int -> [[(Int, Int, Edit)]] -> [[(Int, Int, Edit)]]
    mergeGroups _ = id


-- ---------------------------------------------------------------------------
-- applyHunk
-- ---------------------------------------------------------------------------

-- | Apply a single hunk to a list of lines.
-- The source lines must match the Context and Removed lines declared in the hunk;
-- otherwise a PatchError is returned.
--
-- Algorithm:
--   1. Validate that _oldStartHunk is within range (1 .. length src + 1).
--   2. Split src into before / target / after blocks using _oldStartHunk and _oldCountHunk.
--   3. Walk _linesHunk, verifying each Context or Removed line against target.
--      Return HunkMismatch (with 1-based line number) on the first discrepancy.
--   4. Build the replacement block from Context and Added lines in _linesHunk.
--   5. Return Right (before ++ replacement ++ after).
applyHunk :: [Text] -> Hunk -> Either PatchError [Text]
applyHunk src hunk
  -- _oldStartHunk must be in the range [1 .. length src + 1].
  -- (+1 allows a hunk that starts right after the last line, e.g. pure insertion at end.)
  | _oldStartHunk hunk < 1 || _oldStartHunk hunk > length src + 1 =
      Left $ InvalidPatch $ T.pack $
        "_oldStartHunk " ++ show (_oldStartHunk hunk) ++
        " is out of range for source of " ++ show (length src) ++ " lines"
  | otherwise =
      let startIdx    = _oldStartHunk hunk - 1          -- convert to 0-based
          before      = take startIdx src
          target      = take (_oldCountHunk hunk) (drop startIdx src)
          after       = drop (startIdx + _oldCountHunk hunk) src
      in  case matchLines (_oldStartHunk hunk) (_linesHunk hunk) target of
            Left err      -> Left err
            Right ()      -> Right (before ++ buildReplacement (_linesHunk hunk) ++ after)

-- | Verify that Context and Removed lines in the hunk match the target lines in order.
-- Returns HunkMismatch with the 1-based source line number on the first discrepancy.
matchLines :: Int -> [Line] -> [Text] -> Either PatchError ()
matchLines lineNo hlines target = go lineNo hlines target
  where
    go _ [] _ = Right ()
    go ln (line : rest) src =
      case _typeLine line of
        -- Added lines do not consume a source line; skip them.
        Added   -> go ln rest src
        -- Context and Removed lines must match the next source line.
        Context -> checkLine ln (_contentLine line) rest src
        Removed -> checkLine ln (_contentLine line) rest src

    checkLine ln expected rest src =
      case src of
        []     -> Left $ HunkMismatch ln expected
        (s:ss) ->
          if s == expected
            then go (ln + 1) rest ss
            else Left $ HunkMismatch ln expected

-- | Extract the replacement lines from a hunk: Context and Added lines only.
-- Removed lines are dropped (they are replaced by Added lines).
buildReplacement :: [Line] -> [Text]
buildReplacement = foldr step []
  where
    step line acc =
      case _typeLine line of
        Removed -> acc
        _       -> _contentLine line : acc


-- ---------------------------------------------------------------------------
-- Fuzzy Patching helpers
-- ---------------------------------------------------------------------------

-- | Generate candidate 0-based indices in priority order: center first,
-- then alternating outward (center+1, center-1, center+2, center-2, ...).
-- Indices outside [0, maxLen] are excluded.
generateSearchIndices :: Int -> Int -> Int -> [Int]
generateSearchIndices center fuzz maxLen =
  filter (\i -> i >= 0 && i <= maxLen) $
    center : concatMap (\d -> [center + d, center - d]) [1 .. fuzz]

-- | Try each candidate 0-based index in order and return the first one for
-- which matchLines succeeds.  Does not apply the hunk; only confirms position.
findFirstMatch :: [Text] -> Hunk -> [Int] -> Maybe Int
findFirstMatch _   _    []         = Nothing
findFirstMatch src hunk (idx:idxs) =
  let target = take (_oldCountHunk hunk) (drop idx src)
      lineNo  = idx + 1   -- matchLines expects a 1-based line number
  in  case matchLines lineNo (_linesHunk hunk) target of
        Right () -> Just idx
        Left  _  -> findFirstMatch src hunk idxs

-- | Find the candidate position that produced the deepest (nearest) mismatch.
-- Returns Just (1-based line number, expected text, actual text) for the best candidate,
-- or Nothing when the candidate list is empty or all candidates have empty targets.
-- Used to build diagnostic messages when fuzzy search fails entirely.
findNearestMismatch :: [Text] -> Hunk -> [Int] -> Maybe (Int, Text, Text)
findNearestMismatch src hunk idxs = go idxs Nothing (0 :: Int)
  where
    -- Walk through candidates, tracking the best (deepest) mismatch seen so far.
    go :: [Int] -> Maybe (Int, Text, Text) -> Int -> Maybe (Int, Text, Text)
    go []         best _        = best
    go (idx:rest) best bestDepth =
      let target = take (_oldCountHunk hunk) (drop idx src)
          lineNo  = idx + 1
      in  case probeDepth lineNo (_linesHunk hunk) target 0 of
            Nothing            -> go rest best bestDepth  -- empty target
            Just (depth, info) ->
              if depth > bestDepth
                then go rest (Just info) depth
                else go rest best bestDepth

    -- Walk hlines against target, counting how many lines matched before failing.
    -- Returns Just (matchDepth, (1-based lineNo, expected, actual)) on mismatch,
    -- Nothing when no mismatch was found (full match or empty hlines).
    probeDepth :: Int -> [Line] -> [Text] -> Int -> Maybe (Int, (Int, Text, Text))
    probeDepth _ []           _     _     = Nothing
    probeDepth ln (line:rest) src'  depth =
      case _typeLine line of
        Added   -> probeDepth ln rest src' depth
        Context -> checkProbe ln (_contentLine line) rest src' depth
        Removed -> checkProbe ln (_contentLine line) rest src' depth

    checkProbe :: Int -> Text -> [Line] -> [Text] -> Int -> Maybe (Int, (Int, Text, Text))
    checkProbe ln expected rest src' depth =
      case src' of
        []     -> Just (depth, (ln, expected, ""))
        (s:ss) ->
          if s == expected
            then probeDepth (ln+1) rest ss (depth+1)
            else Just (depth, (ln, expected, s))

-- | Format a hunk header as "@@ -os,oc +ns,nc @@" for use in diagnostic messages.
formatHunkHeader :: Hunk -> String
formatHunkHeader h =
  "@@ -" ++ show (_oldStartHunk h) ++ "," ++ show (_oldCountHunk h)
  ++ " +"  ++ show (_newStartHunk h) ++ "," ++ show (_newCountHunk h)
  ++ " @@"

-- | Fuzzy-aware variant of applyHunk.
-- Takes a 1-based hunk index for diagnostic messages.
-- Uses the delta-corrected _oldStartHunk as a hint and scans ±_FUZZY_RANGE lines
-- for the first position where context lines match, then applies the hunk there.
-- Falls back to a rich InvalidPatch message when no candidate matches.
applyHunkFuzzy :: Int -> [Text] -> Hunk -> Either PatchError [Text]
applyHunkFuzzy hunkIdx src hunk =
  let preferredIdx = _oldStartHunk hunk - 1   -- convert 1-based hint to 0-based
      candidates   = generateSearchIndices preferredIdx _FUZZY_RANGE (length src)
      loLine       = preferredIdx - _FUZZY_RANGE
      hiLine       = preferredIdx + _FUZZY_RANGE
      -- Count Context+Removed lines declared in the hunk.
      actualOldLines = length [ l | l <- _linesHunk hunk
                                  , _typeLine l `elem` [Context, Removed] ]
      countMismatch  = actualOldLines /= _oldCountHunk hunk
  in  case findFirstMatch src hunk candidates of
        Just bestIdx ->
          applyHunk src (hunk { _oldStartHunk = bestIdx + 1 })  -- restore 1-based
        Nothing ->
          Left $ InvalidPatch $ T.pack $ unlines $
            [ "fuzzy search failed"
            , "  hunk #" ++ show hunkIdx ++ ": " ++ formatHunkHeader hunk
            , "  preferred line: " ++ show (_oldStartHunk hunk)
            , "  searched: " ++ show (max 1 loLine) ++ ".." ++ show (min (length src) hiLine)
            , "  old lines in hunk (Context+Removed): " ++ show actualOldLines
              ++ ", oldCount in header: " ++ show (_oldCountHunk hunk)
              ++ if countMismatch then "  [MISMATCH]" else ""
            ] ++
            case findNearestMismatch src hunk candidates of
              Nothing              -> []
              Just (ln, exp', act) ->
                [ "  nearest mismatch at line " ++ show ln ++ ":"
                , "    expected: " ++ show exp'
                , "    actual:   " ++ show act
                ]