{-# LANGUAGE OverloadedStrings #-}

module TextViewport
  ( Item(..)
  , Buffer(..)
  , WrapStrategy(..)
  , RenderedLine(..)
  , RenderedItem(..)
  , CachedRender(..)
  , RenderCache(..)
  , RenderedBuffer
  , renderBuffer
  , flatten
  , modifyItem
  , updateRenderedItem
  , Viewport(..)
  , defaultViewport
  , scrollUp
  , scrollDown
  , lookupPosition

  , RenderState(..)
  , mkRenderState
  , updateRenderState

  -- Viewport Instance ("simple" façade)
  , ViewportInstance(..)
  , mkViewportInstance
  , visibleLines
  , applyToInstance
  , applyToInstanceRS
  , scrollByI
  , scrollUpI
  , scrollDownI
  , pageUpI
  , pageDownI
  , alignTopI
  , alignBottomI
  , modifyItemI
  , lookupPositionI

  ) where

import Data.DList qualified as DL
import Data.Foldable qualified as F
import Data.HashMap.Strict qualified as HM
import Data.List (minimumBy)
import Data.Ord  (comparing)
import Data.Sequence (Seq)
import Data.Sequence qualified as Seq
import Data.Text (Text)
import Data.Text qualified as T
import Data.Vector (Vector)
import Data.Vector qualified as V
import Text.Hyphenation qualified as H

--------------------------------------------------------------------------------
-- Logical model
--------------------------------------------------------------------------------

data WrapStrategy
  = NoWrap
  | FixedWidthWrap
  | HyphenateWrap
      { hwDict  :: H.Hyphenator
      , hwCache :: HM.HashMap Text [(Text, Text)]
      }

data Item = Item
  { itemText :: Text
  , itemWrap :: WrapStrategy
  }

newtype Buffer = Buffer { unBuffer :: Seq Item }

modifyItem :: Int -> (Item -> Item) -> Buffer -> Buffer
modifyItem ix f (Buffer xs) =
    Buffer (Seq.adjust' f ix xs)

--------------------------------------------------------------------------------
-- Rendering with provenance
--------------------------------------------------------------------------------

data RenderedLine = RenderedLine
  { rlText      :: !Text
  , rlItemIx    :: !Int
  , rlLineIx    :: !Int
  , rlCharStart :: !Int
  } deriving (Show)

type RenderedBuffer = Seq RenderedItem

data RenderedItem = RenderedItem
  { riLines :: !(Vector RenderedLine)
  }

flatten :: RenderedBuffer -> [RenderedLine]
flatten = concatMap (F.toList . riLines) . F.toList

renderBuffer :: Int -> Buffer -> RenderCache -> (RenderCache, RenderedBuffer)
renderBuffer width (Buffer items) (RenderCache cache) =
  let n = Seq.length items
      go i (cAcc, rAcc)
        | i >= n    = (RenderCache cAcc, rAcc)
        | otherwise =
            let item     = Seq.index items i
                mOld     = Seq.index cache i
                newEntry = renderItem width i item mOld
                cAcc'    = Seq.update i (Just newEntry) cAcc
                rAcc'    = rAcc Seq.|> crRendered newEntry
            in go (i + 1) (cAcc', rAcc')
  in go 0 (cache, Seq.empty)

renderItem :: Int -> Int -> Item -> Maybe CachedRender -> CachedRender
renderItem width itemIx (Item txt strategy) mOld =
  case mOld of
    Just old
      | crWidth old == width
      , crText old  == txt
      -> old
    _ ->
      let linesV   = applyStrategy strategy width itemIx txt
          rendered = RenderedItem linesV
      in CachedRender
           { crWidth    = width
           , crStrategy = strategy
           , crText     = txt
           , crRendered = rendered
           }

data CachedRender = CachedRender
  { crWidth    :: !Int
  , crStrategy :: !WrapStrategy
  , crText     :: !Text
  , crRendered :: !RenderedItem
  }

newtype RenderCache = RenderCache { unRenderCache :: Seq (Maybe CachedRender) }

--emptyCache :: RenderCache
--emptyCache = RenderCache Seq.empty

emptyCacheFor :: Buffer -> RenderCache
emptyCacheFor (Buffer items) =
  RenderCache (Seq.replicate (Seq.length items) Nothing)

resizeCache :: Buffer -> RenderCache -> RenderCache
resizeCache (Buffer items) (RenderCache cache) =
  let n = Seq.length items
      m = Seq.length cache
  in RenderCache $
       if m < n then cache <> Seq.replicate (n - m) Nothing
       else Seq.take n cache

--------------------------------------------------------------------------------
-- Wrapping strategies
--------------------------------------------------------------------------------

applyStrategy :: WrapStrategy -> Int -> Int -> Text -> Vector RenderedLine

applyStrategy NoWrap width itemIx txt =
  let rawLines = T.splitOn "\n" txt
      chunks  = map (T.take width) rawLines -- crop
      offsets = scanOffsetsWithNewlines chunks
  in V.fromList
       [ RenderedLine { rlText = chunk, rlItemIx = itemIx, rlLineIx = lineIx, rlCharStart = off }
       | (lineIx, (off, chunk)) <- zip [0..] (zip offsets chunks)
       ]

applyStrategy FixedWidthWrap width itemIx txt =
  let rawLines = T.splitOn "\n" txt
      (dl, _) = foldl step (DL.empty, 0) rawLines
      step (acc, off0) line =
        let chunks   = chunkFixed width line
            offsets  = scanOffsetsFrom off0 chunks
            offNext  = off0 + T.length line + 1
            acc'     = acc `DL.append` DL.fromList (zip offsets chunks)
        in (acc', offNext)
      allChunks = DL.toList dl
  in V.fromList
       [ RenderedLine { rlText = chunk, rlItemIx = itemIx, rlLineIx = lineIx, rlCharStart = off }
       | (lineIx, (off, chunk)) <- zip [0..] allChunks
       ]

applyStrategy (HyphenateWrap dict cache0) width itemIx txt =
  let rawLines = T.splitOn "\n" txt

      -- fold over each physical line, accumulating:
      --   * all rendered (offset, chunk) pairs
      --   * updated hyphenation cache (unused for now)
      --   * running character offset across lines
      (dl, _cache1, _) =
        foldl wrapOneLine (DL.empty, cache0, 0) rawLines
      allChunks = DL.toList dl
  in V.fromList
       [ RenderedLine { rlText = chunk, rlItemIx = itemIx, rlLineIx = lineIx, rlCharStart = off }
       | (lineIx, (off, chunk)) <- zip [0..] allChunks
       ]
  where
    -- Wrap a single physical line using TeX‑lite hyphenation
    --wrapOneLine
    --  :: ([(Int, Text)], HM.HashMap Text [(Text, Text)], Int)
    --  -> Text
    --  -> ([(Int, Text)], HM.HashMap Text [(Text, Text)], Int)
    wrapOneLine (acc, cache, off0) line =
      let (chunks, cache1) = wrapWithHyphenationTeXLite dict width line cache
          offsets          = scanOffsetsFrom off0 chunks
          offNext          = off0 + T.length line + 1
          acc'             = acc `DL.append` DL.fromList (zip offsets chunks)
      in (acc', cache1, offNext)
    --wrapOneLine (acc, cache, off0) line =
    --  let chunks  = wrapWithHyphenationTeXLite dict width line
    --      offsets = scanOffsetsFrom off0 chunks
    --      offNext = off0 + T.length line + 1   -- +1 for newline
    --      acc'    = acc ++ zip offsets chunks
    --  in (acc', cache, offNext)

-- | Compute running character offsets for a list of chunks.
scanOffsetsWithNewlines :: [Text] -> [Int]
scanOffsetsWithNewlines = go 0
  where
    go !_ []     = []
    go !o (l:ls) =
      let off = o
          o'  = o + T.length l + 1   -- +1 for newline
      in off : go o' ls

-- | Chunk a single line into fixed-width pieces.
chunkFixed :: Int -> Text -> [Text]
chunkFixed w t
  | w <= 0    = []
  | T.null t  = [""]
  | otherwise =
      let (h, rest) = T.splitAt w t
      in h : if T.null rest then [] else chunkFixed w rest

--------------------------------------------------------------------------------
-- Hyphenation-aware wrapping (TeX-lite)
--------------------------------------------------------------------------------

-- Compute offsets starting from a base offset
scanOffsetsFrom :: Int -> [Text] -> [Int]
scanOffsetsFrom start = go start
  where
    go !_ []     = []
    go !o (t:ts) = o : go (o + T.length t) ts

wrapWithHyphenationTeXLite
  :: H.Hyphenator
  -> Int
  -> Text
  -> HM.HashMap Text [(Text, Text)]
  -> ([Text], HM.HashMap Text [(Text, Text)])
wrapWithHyphenationTeXLite dict width txt cache0 =
    go cache0 (T.words txt)
  where
    go cache [] = ([], cache)
    go cache ws =
      case lineCandidates dict width cache ws of
        ([], cache1) ->
          let chunks = breakWordSafe width ws
          in (chunks, cache1)
        (cs, cache1) ->
          let (line, rest, _) =
                minimumBy (comparing (scoreCandidate width)) cs
              (more, cache2) = go cache1 rest
          in (line : more, cache2)

-- | Lossless fallback: treat remaining words as one long text and
--   chunk it into width-sized pieces. Never truncates, never drops text.
breakWordSafe :: Int -> [Text] -> [Text]
breakWordSafe width ws =
  chunk (T.unwords ws)
  where
    chunk t
      | T.null t            = []
      | T.length t <= width = [t]
      | otherwise =
          let (c, r) = T.splitAt width t
          in c : chunk r

type Candidate = (Text, [Text], Bool)

lineCandidates
  :: H.Hyphenator
  -> Int
  -> HM.HashMap Text [(Text, Text)]
  -> [Text]
  -> ([(Text, [Text], Bool)], HM.HashMap Text [(Text, Text)])
lineCandidates dict width cache0 ws0 =
  go [] [] cache0 ws0
  where
    go _    acc cache [] = (acc, cache)
    go line acc cache (w:ws) =
      let space   = if null line then "" else " "
          baseTxt = T.unwords line
          wholeTxt = baseTxt <> space <> w
          wholeLen = T.length wholeTxt

          acc1 =
            if wholeLen <= width && not (T.null wholeTxt)
              then (wholeTxt, ws, False) : acc
              else acc

          (hyphs, cache1) =
            case HM.lookup w cache of
              Just hs -> (hs, cache)
              Nothing ->
                let hs = hyphenateWord dict w
                in (hs, HM.insert w hs cache)

          hyphCands =
            [ (preTxt, suf : ws, True)
            | (pre, suf) <- hyphs
            , not (T.null pre)
            , let preTxt = baseTxt <> space <> pre <> "-"
            , T.length preTxt <= width
            ]

          acc2 = hyphCands ++ acc1
      in if wholeLen <= width
           then go (line ++ [w]) acc2 cache1 ws
           else (acc2, cache1)

hyphenateWord :: H.Hyphenator -> Text -> [(Text, Text)]
hyphenateWord dict word =
    let parts = H.hyphenate dict (T.unpack word)
    in [ ( T.pack (concat (take i parts))
         , T.pack (concat (drop i parts))
         )
       | i <- [1 .. length parts - 1]
       ]

scoreCandidate :: Int -> Candidate -> Int
scoreCandidate width (line, _, endsWithHyphen) =
    let len      = T.length line
        remSpace = max 0 (width - len)
        badness  = remSpace * remSpace * remSpace
        hyphenPenalty =
          if endsWithHyphen then 50 else 0
        shortPenalty =
          if len < width `div` 2 then 200 else 0
    in badness + hyphenPenalty + shortPenalty

--------------------------------------------------------------------------------
-- Incremental re-rendering
--------------------------------------------------------------------------------

updateRenderedItem :: Int -> Int -> Buffer -> RenderCache -> Seq RenderedItem -> (RenderCache, Seq RenderedItem)
updateRenderedItem width itemIx (Buffer items) (RenderCache cache) rb =
    let item    = Seq.index items itemIx
        mOld     = Seq.index cache itemIx
        newEntry = renderItem width itemIx item mOld
        newCache = Seq.update itemIx (Just newEntry) cache
        newRB    = Seq.update itemIx (crRendered newEntry) rb
    in (RenderCache newCache, newRB)

--------------------------------------------------------------------------------
-- Render state
--------------------------------------------------------------------------------

data RenderState = RenderState
  { rsBuffer    :: Buffer             -- original items
  , rsCache     :: RenderCache        -- per-item cached renders
  , rsRendered  :: RenderedBuffer     -- fully wrapped + hyphenated lines
  , rsWidth     :: Int                -- wrapping width
  , rsLineCount :: Int                
  }

mkRenderState :: Int -> Buffer -> RenderState
mkRenderState width buf =
  let (cache1, rendered) = renderBuffer width buf (emptyCacheFor buf)
  in RenderState
       { rsBuffer    = buf
       , rsCache     = cache1
       , rsRendered  = rendered
       , rsWidth     = width
       , rsLineCount = length (flatten rendered)
       }

-- RenderState has to be rebuilt whenever the buffer or the width changes.
updateRenderState :: Int -> Buffer -> RenderState -> RenderState
updateRenderState width buf rs =
  let (cache1, rendered) = renderBuffer width buf (rsCache rs)
  in rs
      { rsBuffer    = buf
      , rsCache     = cache1
      , rsRendered  = rendered
      , rsWidth     = width
      , rsLineCount = length (flatten rendered)
      }

modifyItemRS :: Int -> (Item -> Item) -> RenderState -> RenderState
modifyItemRS ix f st =
  let buf'    = modifyItem ix f (rsBuffer st)
      cache'  = resizeCache buf' (rsCache st)
      (cache1, rendered1) = renderBuffer (rsWidth st) buf' cache'
  in st { rsBuffer    = buf'
        , rsCache     = cache1
        , rsRendered  = rendered1
        , rsLineCount = length (flatten rendered1)
        }

insertItem :: Int -> Item -> RenderState -> RenderState
insertItem i newItem st =
  let Buffer items = rsBuffer st
      items'       = Seq.insertAt i newItem items
      buf'         = Buffer items'
      cache'       = resizeCache buf' (rsCache st)
      (cache1, rendered1) = renderBuffer (rsWidth st) buf' cache'
  in st { rsBuffer = buf'
        , rsCache = cache1
        , rsRendered = rendered1
        , rsLineCount = length (flatten rendered1)
        }

deleteItem :: Int -> RenderState -> RenderState
deleteItem i st =
  let Buffer items = rsBuffer st
      items'       = Seq.deleteAt i items
      buf'         = Buffer items'
      cache'       = resizeCache buf' (rsCache st)
      (cache1, rendered1) = renderBuffer (rsWidth st) buf' cache'
  in st { rsBuffer = buf'
        , rsCache = cache1
        , rsRendered = rendered1
        , rsLineCount = length (flatten rendered1)
        }

replaceItem :: Int -> Item -> RenderState -> RenderState
replaceItem i newItem st =
  let Buffer items = rsBuffer st
      items'       = Seq.update i newItem items
      buf'         = Buffer items'
      cache'       = resizeCache buf' (rsCache st)
      (cache1, rendered1) = renderBuffer (rsWidth st) buf' cache'
  in st { rsBuffer = buf'
        , rsCache = cache1
        , rsRendered = rendered1
        , rsLineCount = length (flatten rendered1)
        }

appendItem :: Item -> RenderState -> RenderState
appendItem newItem st =
  insertItem (Seq.length (let Buffer xs = rsBuffer st in xs)) newItem st

clearBuffer :: RenderState -> RenderState
clearBuffer st =
  let buf' = Buffer Seq.empty
      cache' = RenderCache Seq.empty
  in st { rsBuffer = buf'
        , rsCache = cache'
        , rsRendered = Seq.empty
        , rsLineCount = 0
        }

fromList :: Int -> [Item] -> RenderState
fromList width xs =
  let buf     = Buffer (Seq.fromList xs)
      cache0  = RenderCache (Seq.replicate (length xs) Nothing)
      (cache1, rendered) = renderBuffer width buf cache0
  in RenderState
       { rsBuffer    = buf
       , rsCache     = cache1
       , rsRendered  = rendered
       , rsWidth     = width
       , rsLineCount = length (flatten rendered)
       }

fromSeq :: Int -> Seq Item -> RenderState
fromSeq width items =
  let buf     = Buffer items
      cache0  = RenderCache (Seq.replicate (Seq.length items) Nothing)
      (cache1, rendered) = renderBuffer width buf cache0
  in RenderState
       { rsBuffer    = buf
       , rsCache     = cache1
       , rsRendered  = rendered
       , rsWidth     = width
       , rsLineCount = length (flatten rendered)
       }



--------------------------------------------------------------------------------
-- Viewport
--------------------------------------------------------------------------------

data Viewport = Viewport
  { vpWidth  :: !Int
  , vpHeight :: !Int
  , vpOffset :: !Int
  } deriving (Show)

mkViewport :: Int -> Int -> Viewport
mkViewport width height =
  Viewport
    { vpWidth  = width
    , vpHeight = height
    , vpOffset = 0
    }

defaultViewport :: Int -> Int -> RenderState -> Viewport
defaultViewport width height rs =
    alignBottom rs (mkViewport width height)

-- any function that sets vpOffset and can overshoot should use clampViewport
clampViewport :: RenderState -> Viewport -> Viewport
clampViewport rs vp =
  let total  = rsLineCount rs
      maxOff = max 0 (total - vpHeight vp)
      off    = vpOffset vp
  in vp { vpOffset = max 0 (min maxOff off) }

scrollBy :: Int -> RenderState -> Viewport -> Viewport
scrollBy delta rs vp =
  clampViewport rs vp { vpOffset = vpOffset vp + delta }

scrollUp :: Int -> RenderState -> Viewport -> Viewport
scrollUp n = scrollBy (-n)

scrollDown :: Int -> RenderState -> Viewport -> Viewport
scrollDown n = scrollBy n

pageUp :: RenderState -> Viewport -> Viewport
pageUp rs vp =
  scrollBy (-(vpHeight vp)) rs vp

pageDown :: RenderState -> Viewport -> Viewport
pageDown rs vp =
  scrollBy (vpHeight vp) rs vp

alignTop :: Viewport -> Viewport
alignTop vp =
  vp { vpOffset = 0 }

alignBottom :: RenderState -> Viewport -> Viewport
alignBottom rs vp =
  let total = rsLineCount rs
      off   = max 0 (total - vpHeight vp)
  in vp { vpOffset = off }


--------------------------------------------------------------------------------
-- Coordinate lookup
--------------------------------------------------------------------------------

lookupPosition
  :: Int
  -> Int
  -> Viewport
  -> RenderedBuffer
  -> Maybe (Int, Int)
lookupPosition x y vp rb =
  let allLines = flatten rb
      idx      = vpOffset vp + y
  in case drop idx allLines of
       (rl:_) ->
         let charIx = rlCharStart rl + x
         in Just (rlItemIx rl, charIx)
       [] -> Nothing

--------------------------------------------------------------------------------
-- Viewport Instance ("simple" façade)
--------------------------------------------------------------------------------

data ViewportInstance = ViewportInstance
  { viRender :: RenderState
  , viView   :: Viewport
  }

mkViewportInstance :: Int -> Int -> Buffer -> ViewportInstance
mkViewportInstance width height buf =
  let rs = mkRenderState width buf
      vp = clampViewport rs (mkViewport width height)
  in ViewportInstance rs vp


visibleLines :: ViewportInstance -> [RenderedLine]
visibleLines (ViewportInstance rs vp) =
  take (vpHeight vp) . drop (vpOffset vp) . flatten $ rsRendered rs

applyToInstance :: (Viewport -> Viewport) -> ViewportInstance -> ViewportInstance
applyToInstance f (ViewportInstance rs vp) =
    ViewportInstance rs (f vp)

applyToInstanceRS :: (RenderState -> Viewport -> Viewport) -> ViewportInstance -> ViewportInstance
applyToInstanceRS f (ViewportInstance rs vp) =
    ViewportInstance rs (f rs vp)

scrollByI :: Int -> ViewportInstance -> ViewportInstance
scrollByI delta = applyToInstanceRS (scrollBy delta)

scrollUpI :: Int -> ViewportInstance -> ViewportInstance
scrollUpI delta = applyToInstanceRS (scrollUp delta)

scrollDownI :: Int -> ViewportInstance -> ViewportInstance
scrollDownI delta = applyToInstanceRS (scrollDown delta)

pageUpI :: ViewportInstance -> ViewportInstance
pageUpI = applyToInstanceRS pageUp

pageDownI :: ViewportInstance -> ViewportInstance
pageDownI = applyToInstanceRS pageDown

alignTopI :: ViewportInstance -> ViewportInstance
alignTopI = applyToInstance alignTop

alignBottomI :: ViewportInstance -> ViewportInstance
alignBottomI = applyToInstanceRS alignBottom

modifyItemI :: Int -> (Item -> Item) -> ViewportInstance -> ViewportInstance
modifyItemI ix f (ViewportInstance rs vp) =
  let buf' = modifyItem ix f (rsBuffer rs)
      rs'  = mkRenderState (rsWidth rs) buf'
      vp'  = clampViewport rs' vp
  in ViewportInstance rs' vp'

lookupPositionI :: Int -> Int -> ViewportInstance -> Maybe (Int, Int)
lookupPositionI x y (ViewportInstance rs vp) =
    lookupPosition x y vp (rsRendered rs)