1 files changed, 48 insertions, 12 deletions

diff --git a/src/Much/TreeView.hs b/src/Much/TreeView.hs
index 3e94c8a..b3d36e8 100644
--- a/src/Much/TreeView.hs
+++ b/src/Much/TreeView.hs
@@ -16,6 +16,7 @@ module Much.TreeView
     ) where
 
 import Data.ByteString.Lazy.Char8 qualified as LBS8
+import Data.CaseInsensitive qualified as CI
 import Data.Text qualified as T
 import Data.Tree
 import Much.TreeView.Types as Export
@@ -71,7 +72,7 @@ fromMessageForest = map fromMessageTree
 fromMessageTree :: Tree Message -> Tree TreeView
 fromMessageTree (Node m ms) =
     Node (TVMessage m)
-         (xconvHead m <> xconvBody m <> map fromMessageTree ms)
+         (xconvHead m <> xconvBody m <> fromMessageForest ms)
 
 
 xconvHead :: Message -> Forest TreeView
@@ -84,7 +85,8 @@ xconvHead m =
 
 
 xconvBody :: Message -> Forest TreeView
-xconvBody m = map (xconvPart m) (messageBody m)
+xconvBody m =
+    [xconvPart m (messageBody m)]
 
 
 xconvPart :: Message -> MessagePart -> Tree TreeView
@@ -131,16 +133,16 @@ isQuoteLine s0 = do
 -- Loading / Unloading
 --
 
-
 loadSubForest :: TreeView -> IO (Either String (Forest TreeView))
 loadSubForest = \case
     TVMessage m ->
-        Right
+        pure . Right =<<
+            loadPrimaryMessagePart m =<< id
             . unloadPartsWithFilename
             . map unloadReadSubForests
             . concatMap subForest
             . fromMessageForest
-            . findFirsts messageMatch
+            . findTopmost messageMatch
             <$> notmuchShow (termFromMessage m)
 
     TVMessagePart m mp@MessagePart{partContentType="text/html"} ->
@@ -148,10 +150,8 @@ loadSubForest = \case
           Left e -> return $ Left $ show e
           Right html ->
             readProcessWithExitCode
-                "lynx"
-                [ "-dump"
-                , "-nomargins"
-                , "-stdin"
+                "show"
+                [ T.unpack (CI.foldedCase (partContentType mp))
                 ]
                 html
               >>= \case
@@ -172,7 +172,13 @@ loadSubForest = \case
                     $ xconvPart m mp'
 
     TVSearchResult sr ->
-        Right
+        pure . Right =<<
+            mapM
+              (\case
+                Node tv@(TVMessage m) parts -> Node tv <$> loadPrimaryMessagePart m parts
+                node                        -> pure node
+              )
+            =<< id
             . unloadPartsWithFilename
             . map unloadReadSubForests
             . fromMessageForest
@@ -193,6 +199,36 @@ loadSubForest = \case
     termFromSearchResult = unThreadID . searchThread
 
 
+loadPrimaryMessagePart :: Message -> Forest TreeView -> IO (Forest TreeView)
+loadPrimaryMessagePart m parts =
+    mapM (rewriteTreeIO isPrimaryPart loadNode) parts
+  where
+    pid = partID (primaryMessagePart (messageBody m))
+    mid = messageId m
+
+    isPrimaryPart :: TreeView -> Bool
+    isPrimaryPart = \case
+        TVMessagePart m' mp ->
+            messageId m' == mid && partID mp == pid
+        _ -> False
+
+    loadNode :: Tree TreeView -> IO (Tree TreeView)
+    loadNode node@(Node tv _) =
+        loadSubForest tv >>= \case
+            Right subForest' -> pure (Node tv subForest')
+            Left _           -> pure node
+
+
+rewriteTreeIO :: (a -> Bool) -> (Tree a -> IO (Tree a)) -> Tree a -> IO (Tree a)
+rewriteTreeIO shouldRewrite f t@(Node v children)
+  | shouldRewrite v = do
+      Node v' children' <- f t
+      Node v' <$> mapM (rewriteTreeIO shouldRewrite f) children'
+
+  | otherwise = do
+      Node v <$> mapM (rewriteTreeIO shouldRewrite f) children
+
+
 unloadSubForest :: Tree TreeView -> Forest TreeView
 unloadSubForest t = case rootLabel t of
     TVMessage _ ->
@@ -245,8 +281,8 @@ unloadPartsWithFilename =
         in x' { subForest = map rewriteTree $ subForest x' }
 
 
-findFirsts :: (a -> Bool) -> Forest a -> Forest a
-findFirsts p =
+findTopmost :: (a -> Bool) -> Forest a -> Forest a
+findTopmost p =
     concatMap rec
   where
     rec t@Node{..} =