#!/usr/bin/env python3

import urllib.request
import argparse
import re
import os
import json
import gzip
import io
import html.parser

XML_TAGS_TO_STORE = set(["url", "speaker", "title", "keywords",
                         "content", "date", "description"])
TALK_NOT_AVALABLE = "This performance is unavailable. Watch another TED performance:"

parser = argparse.ArgumentParser()
parser.add_argument('--xml_file')
parser.add_argument('--out_dir')
parser.add_argument('--num_comments', type=int, default=10)
parser.add_argument('--overwrite', action="store_true")
args = parser.parse_args()

def Fetch(url, is_ajax=False):
  print("Fetching URL:", url)
  req = urllib.request.Request(url)
  if is_ajax:
    req.add_header("Accept", "application/json, text/javascript, */*; q=0.01")
    req.add_header("X-Requested-With", "XMLHttpRequest")
  req.add_header("Accept-Encoding", "gzip")
  resp = urllib.request.urlopen(req)
  final_url = resp.geturl()
  if final_url != url:
    print("Redirected from:", url, "to:", resp.geturl())
  data = resp.read()
  if resp.info().get('Content-Encoding') == 'gzip':
    buf = io.BytesIO(data)
    f = gzip.GzipFile(fileobj=buf)
    data = f.read()
  return data.decode('utf-8')

def FetchComments(thread_id, num_comments=10):
  return Fetch(
      "http://www.ted.com/conversation_forums/{}?page=1&per_page={}&sort=newest"
          .format(thread_id, num_comments),
      is_ajax=True)

def GetTalkCommentThreadId(talk_url):
  html = Fetch(talk_url)
  if TALK_NOT_AVALABLE in html:
    return None
  matches = re.findall(R'"threadId":(\d+)', html)
  assert len(matches) == 1, (talk_url, matches)
  return matches[0]

def SanitizeComments(thread_obj, level=0):
  result = []
  for val in thread_obj.values():
    if val['children']:
      val['z_children'] = SanitizeComments(val['children'])
    del val['children']
    result.append(val)
  return result
  
def GetCommentsObj(talk_url):
  thread_id = GetTalkCommentThreadId(talk_url)
  if thread_id is None:
    return {}
  comments_json = FetchComments(thread_id, num_comments=args.num_comments)
  comments_obj = json.loads(comments_json)
  thread_obj = comments_obj["discussion_thread"]["thread"][0]
  return SanitizeComments(thread_obj)

class MyHTMLParser(html.parser.HTMLParser):
  def __init__(self, out_dir):
    super().__init__()
    self.out_dir = out_dir
    self.file_id = None
    self.out_obj = {}
    self.current_tag = None
    self.abort = False

  def handle_starttag(self, tag, attrs):
    self.current_tag = tag
    if tag == "file":
      # attrs == [("id", "1")]
      assert attrs[0][0] == "id"
      self.file_id = attrs[0][1]
      print("Processing file:", self.file_id)
 
  def handle_endtag(self, tag):
    if tag == "file":
      self.StoreJson()
      self.file_id = None
      self.out_obj = {}
      self.current_tag = None

  def handle_data(self, data):
    if self.current_tag in XML_TAGS_TO_STORE:
      data = data.strip()
      if data:
        if self.current_tag in self.out_obj:
          self.out_obj[self.current_tag] += "\n" + data
        else:
          self.out_obj[self.current_tag] = data

  def StoreJson(self):
    assert "url" in self.out_obj
    out_filename = os.path.join(self.out_dir, self.file_id)
    if os.path.exists(out_filename) and not args.overwrite:
      return
    self.out_obj["z_comments"] = GetCommentsObj(self.out_obj["url"])
    file_obj = open(out_filename, "w")
    json.dump(obj=self.out_obj, fp=file_obj, sort_keys=True, indent=2)
    #self.abort = True

def ParseXmlFile(xml_file, json_base_dir):
  parser = MyHTMLParser(json_base_dir)
  for line in open(xml_file):
    parser.feed(line)
    if parser.abort:
      break

def main():
  #print(GetTalkCommentThreadId("http://www.ted.com/talks/amy_cuddy_your_body_language_shapes_who_you_are"))
  ParseXmlFile(args.xml_file, args.out_dir)


if __name__ == '__main__':
  main()
