{-# LANGUAGE OverloadedStrings #-}
module KnownHosts
    ( KnownHosts(..)
    , decode
    , encode
    , assignKey
    , parseLine
    , serializeLine ) where

import qualified Data.ByteString.Lazy.Char8 as L
import Data.Char
import Data.List
import Data.Ord
import Data.Maybe
import qualified SSHKey as SSH

data KnownHostsEntry = KnownHostsEntry
    { eLineno :: Int
    , eSp :: [L.ByteString]
    , eComment :: L.ByteString
    , eKeyBlob :: L.ByteString
    , eKey :: Maybe SSH.Key
    , eKeyHeader :: L.ByteString
    , eHosts :: [L.ByteString]
    }
 deriving (Eq,Show)

serializeLine :: KnownHostsEntry -> L.ByteString
serializeLine (KnownHostsEntry { eSp = sp
                               , eComment = comment
                               , eKeyHeader = hdr
                               , eKeyBlob = blob
                               , eKey = _
                               , eHosts = hosts
                               })
    = L.concat $ zipWith L.append [host,hdr,blob] (take 3 sp) ++ [comment]
    -- = L.concat $ zipWith L.append [host,hdr,keyblob] (take 3 sp) ++ [comment]
 where
    host = L.intercalate "," hosts
    -- keyblob = L.dropWhile isSpace $ L.dropWhile (not . isSpace) $ SSH.keyblob key
    

parseLine :: Int -> L.ByteString -> KnownHostsEntry
parseLine num str = KnownHostsEntry
    { eLineno = num
    , eSp = sp
    , eComment = comment
    , eKeyHeader = hdr
    , eKeyBlob = blob
    , eKey = key
    , eHosts = hosts
    }
 where
    gs = L.groupBy (\a b -> isSpace a==isSpace b) str
    (sp,ns) = partition (\x -> isSpace (L.head x)) gs
    hosts = do
        h <- take 1 ns
        let hs = L.groupBy (\_ b -> b/=',') h
            hs' = map L.tail $ drop 1 hs
        h0 <- take 1 hs
        h0 : hs'
    hdr = L.concat $ take 1 (drop 1 ns)
    blob = L.concat $ take 1 (drop 2 ns)
    key = SSH.blobkey $ L.concat $ hdr:" ":blob:[]
    comment = L.concat $ zipWith L.append (drop 3 ns) (drop 3 sp ++ repeat "")

data KnownHosts = KnownHosts
    { khcount :: Int
    , khentries :: [KnownHostsEntry]
    }

renum kh = kh { khentries = zipWith setnum [1..] (khentries kh) }
        where setnum n e = e { eLineno=n }

decode :: L.ByteString -> KnownHosts
decode input = renum $ foldr grok (KnownHosts 0 []) ls
 where
    ls = L.lines input
    grok x (KnownHosts cnt es) =
        KnownHosts (cnt+1) (parseLine 0 x:es)

encode :: KnownHosts -> L.ByteString
encode kh = L.unlines $ map serializeLine (khentries kh)

assignKey :: SSH.Key -> L.ByteString -> KnownHosts -> KnownHosts
assignKey key name kh = renum KnownHosts { khcount = khcount kh - length ds
                                         , khentries = es }
 where
    (ms,ns) = partition (\e -> eKey e == Just key) $ khentries kh
    ms' = map (\e -> e { eHosts = eHosts e `union` [name]}) ms
    ns' = map (\e -> e { eHosts = eHosts e \\ if isJust (eKey e) then [name]
                                                                 else []
                       })
              ns
    (ds,ns'') = partition (\e -> null (eHosts e)) ns'
    es = sortBy (comparing eLineno) $ ms' ++ ns''