1 files changed, 86 insertions, 28 deletions

diff --git a/src/helpers.cr b/src/helpers.cr
index 1a6037a5..a9866cee 100644
--- a/src/helpers.cr
+++ b/src/helpers.cr
@@ -1,7 +1,57 @@
+class Video
+  module HTTPParamConverter
+    def self.from_rs(rs)
+      HTTP::Params.parse(rs.read(String))
+    end
+  end
+
+  module XMLConverter
+    def self.from_rs(rs)
+      XML.parse_html(rs.read(String))
+    end
+  end
+
+  def initialize(id, info, html, updated, title, views, likes, dislikes, wilson_score)
+    @id = id
+    @info = info
+    @html = html
+    @updated = updated
+    @title = title
+    @views = views
+    @likes = likes
+    @dislikes = dislikes
+    @wilson_score = wilson_score
+  end
+
+  def to_a
+    return [@id, @info, @html, @updated, @title, @views, @likes, @dislikes, @wilson_score]
+  end
+
+  DB.mapping({
+    id:   String,
+    info: {
+      type:      HTTP::Params,
+      default:   HTTP::Params.parse(""),
+      converter: Video::HTTPParamConverter,
+    },
+    html: {
+      type:      XML::Node,
+      default:   XML.parse_html(""),
+      converter: Video::XMLConverter,
+    },
+    updated:      Time,
+    title:        String,
+    views:        Int64,
+    likes:        Int32,
+    dislikes:     Int32,
+    wilson_score: Float64,
+  })
+end
+
 # See http://www.evanmiller.org/how-not-to-sort-by-average-rating.html
 def ci_lower_bound(pos, n)
   if n == 0
-    return 0
+    return 0.0
   end
 
   # z value here represents a confidence level of 0.95
@@ -18,57 +68,67 @@ def elapsed_text(elapsed)
   "#{(millis * 1000).round(2)}µs"
 end
 
-def get_client
-  while POOL.empty?
+def get_client(pool)
+  while pool.empty?
     sleep rand(0..10).milliseconds
   end
 
-  return POOL.shift
+  return pool.shift
 end
 
-def fetch_video(id)
-  # Grab connection from pool
-  client = get_client
-
-  info = client.get("/get_video_info?video_id=#{id}&el=detailpage&ps=default&eurl=&gl=US&hl=en").body
-  info = HTTP::Params.parse(info)
+def fetch_video(id, client)
+  begin
+    info = client.get("/get_video_info?video_id=#{id}&el=detailpage&ps=default&eurl=&gl=US&hl=en").body
+    html = client.get("/watch?v=#{id}").body
+  end
 
-  html = client.get("/watch?v=#{id}").body
   html = XML.parse_html(html)
+  info = HTTP::Params.parse(info)
 
   if info["reason"]?
     raise info["reason"]
   end
 
-  # Return connection to pool
-  POOL << client
+  title = info["title"]
+
+  views = info["view_count"].to_i64
 
-  video = Video.new(id, info, html, Time.now)
+  likes = html.xpath_node(%q(//button[@title="I like this"]/span))
+  likes = likes ? likes.content.delete(",").to_i : 1
+
+  dislikes = html.xpath_node(%q(//button[@title="I dislike this"]/span))
+  dislikes = dislikes ? dislikes.content.delete(",").to_i : 0
+
+  wilson_score = ci_lower_bound(likes, likes + dislikes)
+
+  video = Video.new(id, info, html, Time.now, title, views, likes, dislikes, wilson_score)
 
   return video
 end
 
-def get_video(id, refresh = true)
-  if PG_DB.query_one?("SELECT EXISTS (SELECT true FROM videos WHERE id = $1)", id, as: Bool)
-    video = PG_DB.query_one("SELECT * FROM videos WHERE id = $1", id, as: Video)
+def get_video(id, client, db, refresh = true)
+  if db.query_one?("SELECT EXISTS (SELECT true FROM videos WHERE id = $1)", id, as: Bool)
+    video = db.query_one("SELECT * FROM videos WHERE id = $1", id, as: Video)
 
-    # If record was last updated more than 5 hours ago, refresh (expire param in response lasts for 6 hours)
-    if refresh && Time.now - video.updated > Time::Span.new(0, 5, 0, 0)
-      video = fetch_video(id)
-      PG_DB.exec("UPDATE videos SET info = $2, html = $3, updated = $4 WHERE id = $1", video.to_a)
+    # If record was last updated over an hour ago, refresh (expire param in response lasts for 6 hours)
+    if refresh && Time.now - video.updated > 1.hours
+      video = fetch_video(id, client)
+      db.exec("UPDATE videos SET info = $2, html = $3, updated = $4,\
+       title = $5, views = $6, likes = $7, dislikes = $8, wilson_score = $9 WHERE id = $1", video.to_a)
     end
   else
-    video = fetch_video(id)
-    PG_DB.exec("INSERT INTO videos VALUES ($1, $2, $3, $4)", video.to_a)
+    video = fetch_video(id, client)
+    db.exec("INSERT INTO videos VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", video.to_a)
   end
 
   return video
 end
 
-def search(query)
-  client = get_client
+def search(query, client)
+  begin
+    html = client.get("https://www.youtube.com/results?q=#{query}&sp=EgIQAVAU").body
+  end
 
-  html = client.get("https://www.youtube.com/results?q=#{query}&sp=EgIQAVAU").body
   html = XML.parse_html(html)
 
   html.xpath_nodes(%q(//ol[@class="item-section"]/li)).each do |item|
@@ -80,6 +140,4 @@ def search(query)
       end
     end
   end
-
-  POOL << client
 end