{-# LANGUAGE OverloadedStrings #-}
module Table(applySort, applySortDoc, splitTable) where

import Text.XML
import Data.Text
import qualified Data.Map as M

import Data.Maybe
import qualified Data.List as L
import Text.Read (readMaybe)

applySortDoc :: String -> Document -> Document
applySortDoc anchor doc@Document {documentRoot = el} = doc {documentRoot = applySort anchor el}

applySort :: String -> Element -> Element
applySort ('#':'-':'r':'h':'a':'p':'s':'-':'%':anchor) el
    | (id', ord:col) <- L.break (`elem` ['<', '>']) anchor, Just col' <- readMaybe col =
        applySort' id' (ord == '<') col' el
applySort _ el = el

applySort' :: String -> Bool -> Int -> Element -> Element
applySort' ('.':id') asc col el@Element { elementNodes = childs }
    | (ix, subpath) <- L.break (== '.') id', Just ix' <- readMaybe ix =
        el { elementNodes = setAt ix' (rewriteNode subpath) childs }
    | otherwise = el
  where
    rewriteNode p (NodeElement child) = NodeElement $ applySort' p asc col child
    rewriteNode _ x = x
applySort' "" asc col el = applySort'' asc col el

applySort' id' asc col el@Element { elementAttributes = attrs, elementNodes = childs }
    | Just actual <- "id" `M.lookup` M.mapKeys nameLocalName attrs, pack id' == actual =
        applySort'' asc col el
    | otherwise = el { elementNodes = L.map searchNode childs }
  where
    searchNode (NodeElement child) = NodeElement $ applySort' id' asc col child
    searchNode x = x

applySort'' asc col el
    | Just sortable <- table2sorttable el = el {
        elementNodes = annotateTHead header asc col ++
            (L.concatMap (L.map NodeElement . markup) $ L.sortBy compareRows sortable)
            ++ footer
      }
    | otherwise = el
  where
    compareRows (TableRow a _) (TableRow b _)
        | asc = (a !! col) `compare` (b !! col)
        | otherwise = (b !! col) `compare` (a !! col)
    (header, _, footer) = splitTable $ elementNodes el

data TableRow = TableRow { keys :: [Text], markup :: [Element] }

table2sorttable Element {
        elementName = Name "table" _ _,
        elementAttributes = attrs,
        elementNodes = childs
    } | "-rhaps-unsortable" `elem` attrs, (_, body, _) <- splitTable childs =
        trs2sorttable body
table2sorttable _ = Nothing

splitTable :: [Node] -> ([Node], [Element], [Node])
splitTable (NodeElement el@Element { elementName = Name "caption" _ _}:els) =
    let (header, body, footer) = splitTable els in (NodeElement el:header, body, footer)
splitTable (NodeElement el@Element { elementName = Name "colgroup" _ _}:els) =
    let (header, body, footer) = splitTable els in (NodeElement el:header, body, footer)
splitTable (NodeElement el@Element { elementName = Name "thead" _ _}:els) =
    let (body, footer) = splitTableBody els in ([NodeElement el], body, footer)
splitTable (NodeElement el@Element { elementName = Name "tr" _ _, elementNodes = childs}:els)
    | L.all (== "th") [nameLocalName $ elementName el | NodeElement el <- childs] =
        let (body, footer) = splitTableBody els in ([NodeElement el], body, footer)
splitTable els@(NodeElement _:_) =
    let (body, footer) = splitTableBody els in ([], body, footer)
splitTable (_:els) = splitTable els
splitTable [] = ([], [], [])

splitTableBody :: [Node] -> ([Element], [Node])
splitTableBody (NodeElement el@Element { elementName = Name "tbody" _ _, elementNodes = childs }:els) =
    ([el | NodeElement el@Element { elementName = Name "tr" _ _ } <- childs], els)
splitTableBody (NodeElement el@Element { elementName = Name "tr" _ _ }:els) =
    let (body, footer) = splitTableBody els in (el:body, footer)
splitTableBody els@(NodeElement _:_) = ([], els)
splitTableBody (_:els) = splitTableBody els
splitTableBody [] = ([], [])

annotateTHead (NodeElement el@Element { elementName = Name "thead" _ _, elementNodes = childs }:nodes) a c =
    NodeElement el { elementNodes = annotateTHead childs a c } : nodes
annotateTHead (NodeElement el@Element { elementName = Name "tr" _ _, elementNodes = childs }:nodes) a c =
    NodeElement el { elementNodes = annotateTR childs a c 0 } : nodes
annotateTHead (child:childs) a c = child:annotateTHead childs a c
annotateTHead [] _ _ = []

annotateTR (NodeElement el@Element { elementName = Name n _ _, elementAttributes = attrs }:nodes) asc col count
    | n `elem` ["th", "td"], count >= col =
        NodeElement el { elementAttributes = M.insert "aria-sort" asc' attrs }:nodes
    | n `elem` ["th", "td"] = NodeElement el:annotateTR nodes asc col (count + colspan)
  where
    colspan = fromMaybe 1 (readMaybe =<< unpack <$> M.lookup "colspan" attrs')
    attrs' = M.mapKeys nameLocalName attrs
    asc' | asc = "ascending"
        | otherwise = "descending"
annotateTR (node:nodes) a c n = node:annotateTR nodes a c n
annotateTR [] _ _ _ = []

trs2sorttable els@(el@Element { elementName = Name "tr" _ _, elementNodes = childs }:_)
    | Just keys' <- tds2keys [el | NodeElement el <- childs],
      Just (group, rest) <- groupTrs els 1,
      Just rest' <- trs2sorttable rest = Just (TableRow keys' group : rest')
trs2sorttable [] = Just []
trs2sorttable _ = Nothing

tds2keys :: [Element] -> Maybe [Text]
tds2keys (el@Element {elementName = Name n _ _, elementAttributes = attrs, elementNodes = childs }:els)
    | n `elem` ["td", "th"], Just rest <- tds2keys els =
        Just (Prelude.replicate colspan (nodesText childs) ++ rest)
  where
    colspan | Just n <- "colspan" `M.lookup` M.mapKeys nameLocalName attrs,
            Just m <- readMaybe $ unpack n = m
        | otherwise = 1
tds2keys [] = Just []
tds2keys _ = Nothing

groupTrs (el@Element {elementName = Name "tr" _ _}:els) n
    | rowRowspan n el <= 1 = Just (el:[], els)
    | Just (tail, rest) <- groupTrs els $ pred n = Just (el:tail, rest)
groupTrs (_:els) n = groupTrs els n
groupTrs _ _ = Nothing

rowRowspan n Element {elementName = Name "tr" _ _, elementNodes = childs } =
    Prelude.maximum (n : [n |
            NodeElement (Element (Name name _ _) attrs _) <- childs,
            name `elem` ["td", "th"],
            rowspan <- maybeToList ("rowspan" `M.lookup` M.mapKeys nameLocalName attrs),
            n <- maybeToList $ readMaybe $ unpack rowspan])


--- Utils

(+++) = append
nodesText :: [Node] -> Text
nodesText (NodeElement (Element _ attrs children):nodes) = nodesText children +++ nodesText nodes
nodesText (NodeContent text:nodes) = text +++ nodesText nodes
nodesText (_:nodes) = nodesText nodes
nodesText [] = ""

setAt :: Int -> (a -> a) -> [a] -> [a]
setAt i a ls
  | i < 0 = ls
  | otherwise = go i ls
  where
    go 0 (x:xs) = a x : xs
    go n (x:xs) = x : go (n-1) xs
    go _ []     = []