{-# LANGUAGE RecordPuns, ParallelListComp, PatternGuards #-}
{-# OPTIONS -fno-warn-name-shadowing -funbox-strict-fields #-}

module Text.PageIO.Transform where
import Data.Maybe
import Data.Monoid
import Data.Function (on)
import Data.List (find, inits, tails, sortBy, groupBy, intersperse)
import Text.Printf
import qualified Text.PageIO.LabelMap as LM
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as L

import Text.PageIO.Types
import Text.PageIO.Extract (extractPage, SheetResult(..), BlockResult(..), Area)

data Doc = MkDoc
    { docMeta    :: !SheetResult
    , docContent :: !L.ByteString
    }
    deriving (Show, Eq, Ord)

type ValueMap = LabelMap [Value]
data AppliedVariable = MkAppliedVariable
    { avRow     :: !Row
    , avCol     :: !Col
    , avValue   :: !Value
    }
    deriving (Show, Eq, Ord)

data Slot = MkSlot
    { slotSize      :: !Row
    , slotBlocks    :: ![Label]
    }
    deriving (Eq, Ord, Show)

type Ordered a = ([OrderBy (Maybe Value)], a)

data BlockData = MkBlockData
    { dataSize      :: !Row
    , dataAreas     :: ![Ordered Area]
    }
    deriving (Eq, Ord, Show)

instance Monoid BlockData where
    mempty = MkBlockData (error "impossible") []
    mappend x y = x{ dataAreas = dataAreas x ++ dataAreas y }
    mconcat []  = mempty
    mconcat xs  = MkBlockData
        { dataSize  = dataSize (head xs)
        , dataAreas = concatMap dataAreas xs
        }

type PageCapacity = [Slot]
type FitAttempt = LabelMap BlockData

-- The actual data bound to that area.
type PageBinding = LabelMap [Area]

parsePages :: Sheet -> [Page] -> [Doc]
parsePages sheetIn pagesIn = map emitDoc docGroupsOut
    where
    docGroupsOut = case sheetGroupBy sheetIn of
        []      -> [resultsIn]
        lbls    -> groupBy ((==) `on` docKeys) resultsIn
            where
            docKeys (res, _) = map (`LM.lookup` resultFields res) lbls
    resultsIn = [ (res, page) | (Just res, page) <- map (\p -> (extractPage sheetIn p, p)) pagesIn ]

emitDoc :: [(SheetResult, Page)] -> Doc
emitDoc xs = MkDoc meta (packPages pages)
    where
    meta = mconcat results
    results = map fst xs
    pages   = map snd xs


-- Turn a list of page into pages grouped by results
transformPages :: Sheet -> [Page] -> Sheet -> [Page] -> [Doc]
transformPages sheetIn pagesIn sheetOut pagesOut = map (makeDoc sheetOut) docGroupsOut
    where
    bindingsOut = concatMap (bindDoc sheetOut pagesOut) docGroupsIn
    docGroupsOut = case sheetGroupBy sheetOut of
        []      -> [bindingsOut]
        lbls    -> groupBy ((==) `on` docKeys) bindingsOut
            where
            docKeys MkDocBinding{ docResult } = map (`LM.lookup` resultFields docResult) lbls
    docGroupsIn = case sheetGroupBy sheetIn of
        []      -> [resultsIn]
        lbls    -> groupBy ((==) `on` docKeys) resultsIn
            where
            docKeys res = map (`LM.lookup` resultFields res) lbls
    resultsIn = [ res | Just res <- map (extractPage sheetIn) pagesIn ]

makeDoc :: Sheet -> [DocBinding] -> Doc
makeDoc sheetOut xs = MkDoc meta . packPages $ fillVariables sheetOut results
    [ makePage sheetOut b p
    | b <- map docBinding xs
    | p <- map docPage xs
    ]
    where
    meta = mconcat results
    results = map docResult xs

packPages :: [Page] -> L.ByteString
packPages ps = unpages [ L.unlines [ L.fromChunks [l] | l <- pageLines p ] | p <- ps ]
    where
    unpages :: [L.ByteString] -> L.ByteString
    unpages [] = L.empty
    unpages ss = (L.concat $ intersperse nl ss) `L.append` nl -- half as much space
        where nl = L.pack "\x0C\n"

-- Using Sheet+Page as template to rewrite a page
-- First, rewrite SheetResult to eliminate all Variable fields
-- Next, populate all fields
-- Finally, populate frames with blocks and space-ify the rest of the frame lines
--
-- First try fitting everything on the last page.
-- All frames must be consumed by blocks mentioned on that page.
-- If it won't fit, then retry with the last two pages, etc.
--
-- Oh, and we re-extract the result pages to fill in vars! What joy!
--
-- XXX - If data would expand, then maybe repeat the first page indefinitely?
--
data DocBinding = MkDocBinding
    { docResult     :: !SheetResult
    , docPage       :: !Page
    , docBinding    :: !PageBinding
    }
    deriving (Show, Eq, Ord)

bindDoc :: Sheet -> [Page] -> [SheetResult] -> [DocBinding]
bindDoc sheetOut pagesOut resultsIn = case find (\(x, _, _) -> isJust x) pageBindings of
    Just (Just bindings, boundPages, boundResults) -> 
        [ MkDocBinding r p b
        | r <- boundResults
        | b <- bindings
        | p <- boundPages
        ]
    _  -> []
    where
    resultsOut          = map (maybe (error "Roundtrip failed!") id . extractPage sheetOut) pagesOut
    orders              = sheetBlockOrderBys sheetOut
    capacity            = foldr (doCapacity sheetOut) [] resultsOut 
    attempt             = foldr (doAttempt orders) mempty resultsIn
    capacityTails       = repeatTails capacity
    pageTails           = repeatTails pagesOut
    resultTails         = repeatTails resultsIn
    pageBindings        = 
        [ (tryFit pc sortedAttempt, pages', results')
        | pc        <- capacityTails
        | pages'    <- pageTails
        | results'  <- resultTails
        ]
    sortedAttempt       = (`fmap` filteredAttempt) $ \dat ->
        dat{ dataAreas = sortBy (compare `on` fst) $ dataAreas dat }
    filteredAttempt     = attempt `LM.intersection` capacityLabels
    capacityLabels      = LM.fromList [ (k, ()) | k <- concatMap (concatMap slotBlocks) capacity ]

-- Given "abc", generate ["", "c", "bc", "abc", "aabc", "aaabc", "aaaabc"...]
repeatTails :: [a] -> [[a]]
repeatTails [] = []
repeatTails orig@(x:xs) = reverse (tails xs) ++ map (++ orig) infinitePrefixes
    where
    infinitePrefixes = inits (repeat x)

-- We now calculate the applied vars for each page
fillVariables :: Sheet -> [SheetResult] -> [Page] -> [Page]
fillVariables sheet results pages =
    [ fillPageVariables appliedVars p
    | p             <- pages
    | appliedVars   <- map applyVariable ([1..] `zip` valueMaps)
    ]
    where
    vars        = sheetVariableFields sheet
    valueMaps   = map makeValueMap
        [ MkSheetResult rf nr
        | MkSheetResult rf _  <- results
        | MkSheetResult _ nr <- newResults
        ]
    docVals     = LM.unionsWith (++) valueMaps
    newResults  = map (maybe (error "1Roundtrip failed!") id . extractPage sheet) pages
    applyVariable :: (Int, ValueMap) -> LabelMap AppliedVariable
    applyVariable (pageNum, pageVals) = LM.mapMaybeWithKey applyOneVariable vars
        where
        applyOneVariable :: Label -> Field -> Maybe AppliedVariable
        applyOneVariable lbl MkField{ fieldBox, fieldVariable }
            | Just var <- fieldVariable
            , LM.member lbl pageVals  = Just MkAppliedVariable
                { avRow     = boxTop fieldBox
                , avCol     = boxLeft fieldBox
                , avValue   = getVar var
                }
            | otherwise                 = Nothing
            where
            len = boxRight fieldBox - boxLeft fieldBox + 1
            valOf :: Scope -> Label -> [Value]
            valOf scope lbl = maybe [] id $ LM.lookup lbl vals
                where
                vals = case scope of
                    SDoc    -> docVals
                    SPage   -> pageVals
            getVar VPage                = formatNumber pageNum
            getVar (VSum scope label)   = formatDotted $ sum (map valToInt $ valOf scope label)
            getVar (VCount scope label) = formatNumber $ length (valOf scope label)
            getVar (VLiteral lit)       = formatLiteral lit
            getVar (VLabel label)       = formatLiteral $ case valOf SPage label of
                []      -> S.empty
                (x:_)   -> x
            formatLiteral lit = lit `S.append` S.replicate (len - S.length lit) ' '
            formatNumber = S.pack . printf ("%" ++ show len ++ "d")
            formatDotted n = S.pack $ (replicate pad ' ') ++ dottedStr ++ ".00"
                where
                pad         = len - length dottedStr - 3
                str         = show n
                dottedStr   = reverse (addDot $ reverse str)
                addDot (x:y:z:rest@(_:_))   = (x:y:z:',':addDot rest)
                addDot xs                   = xs

makeValueMap :: SheetResult -> ValueMap
makeValueMap MkSheetResult{ resultFields, resultBlocks } = LM.unionsWith (++) (fieldMap:blockMaps)
    where
    fieldMap  = fmap (:[]) resultFields
    blockMaps :: [ValueMap]
    blockMaps = map makeBlockValueMap $ LM.elems resultBlocks
    makeBlockValueMap :: BlockResult -> ValueMap
    makeBlockValueMap (MkBlockResult avs) = LM.unionsWith (++) maps
        where
        maps :: [LabelMap [Value]]
        maps = map (fmap (:[]) . snd) avs

-- Now comes the fun part.
-- For each frame, see if any of its blocks are there.
-- If yes, concat them and pain the rest white.
makePage :: Sheet -> PageBinding -> Page -> Page
makePage MkSheet{ sheetFrames } binding page = foldl makeFrame page sheetFrames
    where
    makeFrame page MkFrame{ frameBox, frameBlocks }
        | null frameBindings = page
        | otherwise          = pageReplaced
        where
        pageCleared   = clearArea frameBox page
        pageReplaced  = replaceArea frameBox (concat frameBindings) pageCleared
        frameBindings = 
            [ fromJust result
            | frameLabel <- LM.keys frameBlocks
            , let result = LM.lookup frameLabel binding
            , isJust result
            ]

-- Start with 
replaceArea :: Box -> [Area] -> Page -> Page
replaceArea _ [] page = page
replaceArea box@MkBox{ boxTop, boxLeft } (area:rest) page = replaceArea box' rest page'
    where
    box'  = box{ boxTop = boxTop + areaRows area}
    page' = fillArea boxTop boxLeft area page

clearArea :: Box -> Page -> Page
clearArea MkBox{ boxTop, boxLeft, boxBottom, boxRight } (MkPage lns) = MkPage $ before ++ cleared ++ after
    where
    (before, rest)  = splitAt (boxTop-1) lns
    (middle, after) = splitAt (boxBottom-boxTop+1) rest
    cleared         = map (fillLine boxLeft . ((,) whiteSpace)) middle
    whiteSpace      = S.replicate (boxRight-boxLeft+1) ' '

{-# INLINE fillArea #-}
fillArea :: Row -> Col -> Area -> Page -> Page
fillArea row col (MkPage areaLines) (MkPage lns) = MkPage $ before ++ replaced ++ after
    where
    (before, rest)  = splitAt (row-1) lns
    (middle, after) = splitAt (length areaLines) rest
    replaced        = map (fillLine col) (areaLines `zip` middle)

{-# INLINE fillLine #-}
fillLine :: Col -> (Value, Value) -> Value
fillLine col (areaLine, origLine) = S.concat [before, areaLine, after]
    where
    origLinePadded
        | padLength <= 0    = origLine
        | otherwise         = origLine `S.append` S.replicate padLength '-' 
    padLength       = (col + areaLength - 1) - S.length origLine
    (before, rest)  = S.splitAt (col-1) origLinePadded
    after           = S.drop areaLength rest
    areaLength      = S.length areaLine

doCapacity :: Sheet -> SheetResult -> [PageCapacity] -> [PageCapacity]
doCapacity MkSheet{ sheetFrames } MkSheetResult{ resultBlocks } pcs = pc:pcs
    where
    pc              = map frameToSlot framesOccured
    framesOccured   = filter (labelOccured . frameBlocks) sheetFrames
    labelOccured lm = any (`LM.member` lm) resultLabels
    resultLabels    = LM.keys resultBlocks
    frameToSlot MkFrame{ frameBox, frameBlocks } = MkSlot
        { slotSize   = boxBottom frameBox - boxTop frameBox + 1
        , slotBlocks = LM.keys frameBlocks
        }

doAttempt :: LabelMap [OrderBy Label] -> SheetResult -> FitAttempt -> FitAttempt
doAttempt orders MkSheetResult{ resultBlocks } att = att'
    where
    att'            = LM.unionsWith mappend [att, LM.mapMaybeWithKey blockToData resultBlocks]
    blockToData lbl (MkBlockResult lms) = case areas of 
        []              -> Nothing
        ((_, area):_)   -> Just MkBlockData
            { dataSize  = areaRows area
            , dataAreas = areas
            }
        where
        orderFrom = case LM.lookup lbl orders of
            Just orderBys   -> \vals -> map (fmap (`LM.lookup` vals)) orderBys
            Nothing         -> const []
        areas = [ (orderFrom vals, area) | (area, vals) <- lms ]

areaRows :: Area -> Row
areaRows (MkPage lns) = length lns

fillPageVariables :: LabelMap AppliedVariable -> Page -> Page
fillPageVariables avs page = foldl fillOneVar page (LM.elems avs)
    where
    fillOneVar page MkAppliedVariable{ avRow, avCol, avValue }
        = fillArea avRow avCol (valueToArea avValue) page

valueToArea :: Value -> Area
valueToArea val = MkPage [val]

sheetVariableFields :: Sheet -> LabelMap Field
sheetVariableFields MkSheet{ sheetFrames, sheetFields } = LM.unions filteredMaps
    where
    filteredMaps = map (LM.filter (isJust . fieldVariable)) fieldMaps
    fieldMaps = LM.mapWithKey addVariable sheetFields : concatMap frameVariableFields sheetFrames
    addVariable lbl fld@MkField{ fieldVariable } = case fieldVariable of
        Just{}  -> fld
        _       -> fld{ fieldVariable = Just (VLabel lbl) }
    frameVariableFields frame@MkFrame{ frameBox } = map blockVariableFields blocks
        where
        blocks    = LM.elems (frameBlocks frame)
        rowOffset = boxTop frameBox - 1
        colOffset = boxLeft frameBox - 1
        blockVariableFields = fmap adjustField . blockFields
        adjustField field@MkField{ fieldBox = MkBox{ boxTop, boxLeft, boxBottom, boxRight } } = field
            { fieldBox = MkBox
                { boxTop    = boxTop      + rowOffset
                , boxLeft   = boxLeft     + colOffset
                , boxBottom = boxBottom   + rowOffset
                , boxRight  = boxRight    + colOffset
                }
            }

sheetBlockOrderBys :: Sheet -> LabelMap [OrderBy Label]
sheetBlockOrderBys MkSheet{ sheetFrames } = LM.unions (map frameBlockOrderBys sheetFrames)
    where
    frameBlockOrderBys MkFrame{ frameBlocks } = LM.fromList
        [ (lbl, order) | (lbl, order@(_:_)) <- LM.toList (fmap blockOrderBy frameBlocks) ]

-- Each try is on a list of Frame Capacities
-- Each frame capacity is a number of rows, and the label of blocks it can consume
-- We try it with the total number of blocks

tryFit :: [PageCapacity] -> FitAttempt -> Maybe [PageBinding]
tryFit pcs att = go pcs ([], att)
    where
    go []     (bindings, att)
        | all (null . dataAreas) (LM.elems att) = Just (reverse bindings)
        | otherwise                             = Nothing
    go (p:ps) (bindings, att) = go ps ((bnd:bindings), att')
        where
        (bnd, att') = fitOnePage p (mempty, att)

fitOnePage :: [Slot] -> (PageBinding, FitAttempt) -> (PageBinding, FitAttempt)
fitOnePage [] x     = x
fitOnePage (s:ss) x = fitOnePage ss $ fitOneSlot s x

fitOneSlot :: Slot -> (PageBinding, FitAttempt) -> (PageBinding, FitAttempt)
fitOneSlot slot@MkSlot{ slotSize, slotBlocks } x@(binding, att) = case slotBlocks of
    []      -> x
    (b:bs)  -> case LM.lookup b att of
        Just dat@MkBlockData{ dataSize, dataAreas } ->
            let consumed     = min (length dataAreas) (slotSize `div` dataSize)
                consumedSize = consumed * dataSize
                newData      = dat{ dataAreas = drop consumed dataAreas }
                newAttempt   = LM.insert b newData att
                newBinding   = LM.insertWith (++) b (map snd $ take consumed dataAreas) binding
                newSlot      = slot
                    { slotSize   = slotSize - consumedSize
                    , slotBlocks = bs
                    }
             in fitOneSlot newSlot (newBinding, newAttempt)
        _   -> fitOneSlot slot{ slotBlocks = bs } x