{-# LANGUAGE OverloadedStrings #-} module Main where import Database.MongoDB as DB import Text.HTML.TagSoup import System.Environment import System.IO import System.Process import Text.Printf import System.Exit import Control.Monad.Trans import Network.URI import Data.List.Split import Data.Maybe import Data.List (isPrefixOf,nub) import Data.Char import qualified Data.Text as T import qualified Data.Map as Map import Control.Monad import Data.Classify.Television import Epguides main :: IO () main = do args <- getArgs let addr = "localhost" case args of [port, "magnet"] -> withMeteor addr port run [port, "shows"] -> withMeteor addr port (getTVShows >>= liftIO . mapM_ putStrLn) [port, "add", name] -> withMeteor addr port (addTVShow name) [port, "remove", name] -> withMeteor addr port (removeTVShow name) [port, "epguides", name] -> withMeteor addr port (mkEpguidesEntries name) [port, "set_epguides", name, url] -> withMeteor addr port (setEpguides name url) [port, "rebuild"] -> withMeteor addr port rebuild _ -> printUsage withMeteor addr port fn = do pipe <- runIOE $ connect (Host addr (PortNumber (fromIntegral (read port :: Integer)))) access pipe master "meteor" fn close pipe printUsage :: IO () printUsage = do prog <- getProgName hPutStrLn stderr $ printf "Usage: %s \n" prog hPutStrLn stderr $ printf "Commands:" hPutStrLn stderr $ printf " magnet Fetch magnet links from stdin and add to DB." hPutStrLn stderr $ printf " shows List the recognized TV shows." hPutStrLn stderr $ printf " add Add a name to the list of known TV shows." hPutStrLn stderr $ printf " remove Remove a name to the list of known TV shows." hPutStrLn stderr $ printf " epguides Fetch titles and air-times from stdin." hPutStrLn stderr $ printf " set_epguides Associate with ." hPutStrLn stderr $ printf " rebuild Recalculate computed properties in the DB." exitWith ExitSuccess isMagnetLink :: URI -> Bool isMagnetLink uri = uriScheme uri == "magnet:" getDownloadName :: URI -> String getDownloadName uri = fromMaybe "" $ listToMaybe $ [ dn | ["dn", dn] <- map (splitOn "=") $ splitOn "&" (dropWhile (== '?') $ uriQuery uri) ] getHash :: URI -> String getHash uri = fromMaybe "" $ listToMaybe $ [ xt | ["xt", xt] <- map (splitOn "=") $ splitOn "&" (dropWhile (== '?') $ uriQuery uri) ] getTVShows = do records <- rest =<< find (select [] "meta") {project = [ "name" =: (1::Integer) ]} return $ nub $ mapMaybe (DB.lookup "name") records addTVShow name = insert "meta" ["name" =: name] setEpguides name url = modify (select ["name" =: name] "meta") ["$set" =: ["epguides" =: url]] getEpguidesURL name = do entry <- findOne (select ["name" =: name] "meta") case DB.lookup "epguides" =<< entry of Nothing -> error $ "No epguides URL associated with " ++ show name Just url -> return url removeTVShow name = delete (select ["name" =: name] "meta") getTags :: String -> [String] getTags = words . lower where lower str = [ if isAlphaNum c then c else ' ' | c <- str ] run = do inp <- liftIO $ getContents let hrefs = [ href | TagOpen "a" attrs <- parseTags inp, ("href", href) <- attrs ] shows <- getTVShows insertMagnetLinks shows hrefs insertMagnetLinks shows uris = do let magnets = [ uri | Just uri <- map parseURI uris, isMagnetLink uri ] elts = [ (uri, elt) | uri <- magnets , let dn = getDownloadName uri , Just elt <- [tryClassify shows dn] ] forM_ elts $ \(uri, (name, version, title)) -> do liftIO $ putStrLn $ "Adding element: " ++ show (name, version, title) let query = [ "name" =: name , case version of Version season episode -> "episode" =: [ "season" =: season, "episode" =: episode ] DateVersion year month day -> "airdate" =: [ "year" =: year, "month" =: month, "day" =: day ] ] inst = [ "link" =: show uri, "hash" =: getHash uri, "tags" =: getTags title ] modify (select query "episodes") ["$addToSet" =: ["instances" =: inst]] return () rebuild = do shows <- getTVShows records <- rest =<< find (select ["instances" =: ["$exists" =: True]] "episodes") forM_ records $ \record -> do let record_id = DB.lookup "_id" record :: Maybe ObjectId forM_ (fromMaybe [] (DB.lookup "instances" record)) $ \inst -> case DB.lookup "link" inst of Nothing -> return () Just uri -> do modify (select ["_id" =: record_id] "episodes") ["$pull" =: ["instances" =: inst]] insertMagnetLinks shows [uri] {- {name: TV show ,searchable: string ,airdate: {year: int, month: int, day: int} ,episode: {season: int, episode: int} ,title: string ,instances: [{link: magnet, hash: hash, tags: [string]}] } -} mkEpguidesEntries name = do shows <- getTVShows unless (name `elem` shows) $ error $ "Unknown TV series: " ++ name epguides_url <- getEpguidesURL name inp <- liftIO $ readProcess "wget" ["-q", "-O", "-", epguides_url] [] forM_ (parseEpguide inp) $ \(EpIdx season episode, (year, month, day), title) -> do let query = ["name" =: name, "episode" =: ["season" =: season, "episode" =: episode ]] airdate = ["year" =: year, "month" =: month, "day" =: day] n <- count (select query "episodes") case n of 0 -> do liftIO $ putStrLn $ "Inserting new: " ++ show (season, episode, year, month, day, title) insert_ "episodes" [ "name" =: name , "episode" =: ["season" =: season, "episode" =: episode] , "airdate" =: airdate , "title" =: title ] _ -> do liftIO $ putStrLn $ "Updating: " ++ show (season, episode, year, month, day, title) modify (select query "episodes") ["$set" =: ["title" =: title, "airdate" =: airdate]]