{-# 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.DataTypes
import qualified Data.Classify.Parser as Classify
import qualified Data.Classify.Rank as Classify
import Epguides
main :: IO ()
main = do
args <- getArgs
case args of
[addr, port, "magnet"] -> withMeteor addr port run
[addr, port, "shows"] -> withMeteor addr port (getTVShows >>= liftIO . mapM_ putStrLn)
[addr, port, "add", name] -> withMeteor addr port (addTVShow name)
[addr, port, "remove", name] -> withMeteor addr port (removeTVShow name)
[addr, port, "epguides", name] -> withMeteor addr port (mkEpguidesEntries name)
[addr, port, "set_epguides", name, url] ->
withMeteor addr port (setEpguides name url)
[addr, 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 <mongo_addr> <mongo_port> <command>\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 <name> Add a name to the list of known TV shows."
hPutStrLn stderr $ printf " remove <name> Remove a name to the list of known TV shows."
hPutStrLn stderr $ printf " epguides <name> Fetch titles and air-times from stdin."
hPutStrLn stderr $ printf " set_epguides <name> <url> Associate <url> with <name>."
hPutStrLn stderr $ printf " rebuild Recalculate computed properties in the DB."
exitWith ExitSuccess
-- stripName "The Big Bang Theory" = "bigbangtheory"
-- stripName "CSI: Miami" = "csimiami"
stripName :: String -> String
stripName = worker . filter isAlpha . filter (not . isSpace) . map toLower
where
worker lst | "the" `isPrefixOf` lst = worker (drop 3 lst)
worker lst = lst
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 ]
idx = Map.fromList [ (stripName name, name) | name <- shows ]
elts = [ (uri, Classify.selectBest elts)
| uri <- magnets
, let dn = getDownloadName uri
, let elts = Classify.run (Classify.parseElement idx) dn
, not (null elts) ]
forM_ elts $ \(uri, element) -> do
liftIO $ putStrLn $ "Adding element: " ++ show element
let query = [ "name" =: name element
, case version element 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 element) ]
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
, "searchable" =: map stripName (words 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]]