module Network.DNS.Named.Zone
    ( ZoneFinder(..)
    , fromZones

    , Zone(zoneName)
    , SomeZone(..)
    , ExternalZone(..)
    , DynamicZone(..)
    , StaticZone(..)
    , fromRecords
    , isInZone
    , wrapZone

    , isAuthoritativeZone
    , getRecords
    )
    where

import           Control.Monad
import           Data.List
import qualified Data.Map as M
import           Data.Map (Map)
import           Data.Maybe
import           Data.Typeable
import           Network.DNS.Message
import           Network.DNS.DomainMap


class ZoneFinder a where
    findZone :: a -> DomainName -> IO (Maybe SomeZone)

instance ZoneFinder (DomainMap SomeZone) where
    findZone = (return .) . flip nearest

instance ZoneFinder (IO (DomainMap SomeZone)) where
    findZone = flip (fmap . nearest)

instance ZoneFinder (DomainMap (IO SomeZone)) where
    findZone m n
        = do let getZoneM = nearest n m
             case getZoneM of
               Just getZone -> liftM Just getZone
               Nothing      -> return Nothing

instance ZoneFinder (DomainName -> IO (Maybe SomeZone)) where
    findZone = id


fromZones :: [SomeZone] -> DomainMap SomeZone
fromZones = fromList . map toPair
    where
      toPair z = (zoneName z, z)


class Zone a where
    zoneName          :: a -> DomainName
    getRecordNames    :: a -> IO [DomainName]
    getRecordsForName :: a -> DomainName -> IO [SomeRR]


data SomeZone = forall a. Zone a => SomeZone a

instance Zone SomeZone where
    zoneName          (SomeZone a) = zoneName a
    getRecordNames    (SomeZone a) = getRecordNames a
    getRecordsForName (SomeZone a) = getRecordsForName a


data ExternalZone
    = ExternalZone {
        ezName           :: !DomainName
      , ezRecordNames    :: !(IO [DomainName])
      , ezRecordsForName :: !(DomainName -> IO [SomeRR])
      }
instance Zone ExternalZone where
    zoneName          = ezName
    getRecordNames    = ezRecordNames
    getRecordsForName = ezRecordsForName


data DynamicZone
    = DynamicZone {
        dzName    :: !DomainName
      , dzRecords :: !(IO (Map DomainName [SomeRR]))
      }
instance Zone DynamicZone where
    zoneName          = dzName
    getRecordNames    = liftM M.keys . dzRecords
    getRecordsForName = flip (fmap . (fromMaybe [] .) . M.lookup) . dzRecords


data StaticZone
    = StaticZone {
        szName    :: !DomainName
      , szRecords :: !(Map DomainName [SomeRR])
      }
    deriving (Show, Eq)

instance Zone StaticZone where
    zoneName          = szName
    getRecordNames    = return . M.keys . szRecords
    getRecordsForName = ((return . fromMaybe []) .) . flip M.lookup . szRecords


fromRecords :: [SomeRR] -> Map DomainName [SomeRR]
fromRecords = foldl ins M.empty
    where
      ins m (SomeRR rr)
          = case M.lookup (rrName rr) m of
              Just l
                  -> M.insert (rrName rr) (l ++ [SomeRR rr]) m
              Nothing
                  -> M.insert (rrName rr) [SomeRR rr] m


isInZone :: Zone z => DomainName -> z -> Bool
isInZone name zone = zoneName zone `isZoneOf` name


wrapZone :: Zone z => z -> SomeZone
wrapZone = SomeZone


isAuthoritativeZone :: Zone z => z -> IO Bool
isAuthoritativeZone z
    = do let q = Question {
                   qName  = zoneName z
                 , qType  = SOA
                 , qClass = IN -- Should we consider any classes other than the Internet?
                 }
         rs <- getRecords z q
         case rs of
           [] -> return False
           _  -> return True


getRecords :: (QueryType qt, QueryClass qc, Zone z) =>
              z
           -> Question qt qc
           -> IO [SomeRR]
getRecords z q
    | cast (qType q) == Just AXFR
        = getRecordsForAXFR

    | otherwise
        = do rs <- getRecordsForName z (qName q)
             return $ filterRecords q rs

    where
      getRecordsForAXFR
          | qName q == zoneName z
              = do names      <- getRecordNames z
                   (soaM, rs) <- liftM (spitSOA . concat) $ mapM (getRecordsForName z) names
                   case soaM of
                     Just soa -> return ([soa] ++ rs ++ [soa])
                     Nothing  -> return []
          | otherwise
              = return []

      spitSOA :: [SomeRR] -> (Maybe SomeRR, [SomeRR])
      spitSOA xs = (findSOA xs, collectNonSOA xs)

      findSOA :: [SomeRR] -> Maybe SomeRR
      findSOA []     = Nothing
      findSOA (SomeRR x : xs)
          | typeOf (rrType x) == typeOf SOA = Just (SomeRR x)
          | otherwise                       = findSOA xs

      collectNonSOA :: [SomeRR] -> [SomeRR]
      collectNonSOA []     = []
      collectNonSOA (SomeRR x : xs)
          | typeOf (rrType x) == typeOf SOA = collectNonSOA xs
          | otherwise                       = SomeRR x : collectNonSOA xs


filterRecords :: (QueryType qt, QueryClass qc) =>
                 Question qt qc
              -> [SomeRR]
              -> [SomeRR]
filterRecords q = filter predicate
    where
      predicate rr
          = predForType rr && predForClass rr

      predForType (SomeRR rr)
          | typeOf (qType q) == typeOf ANY
              = True

          | typeOf (qType q) == typeOf MAILB
              = typeOf (rrType rr) == typeOf MR ||
                typeOf (rrType rr) == typeOf MB ||
                typeOf (rrType rr) == typeOf MG ||
                typeOf (rrType rr) == typeOf MINFO

          | otherwise
              = typeOf (rrType rr) == typeOf (qType q) ||
                typeOf (rrType rr) == typeOf CNAME

      predForClass (SomeRR rr)
          | typeOf (qClass q) == typeOf ANY
              = True

          | otherwise
              = typeOf (rrClass rr) == typeOf (qClass q)