#! /usr/bin/env python2.7

import urlparse
import urllib2
import re
import os
import youtube_dl

from bs4 import BeautifulSoup
from ast import literal_eval

ARGS = None

def postprocess(data):
    if data['status'] == 'finished':
        os.system(ARGS.postprocess)

def get_lecture(name, url):
    response = urllib2.urlopen(url)
    soup = BeautifulSoup(response.read(), 'html.parser')
    ocw_embed_chapter_media = soup.find('script', string=re.compile('^ocw_embed_chapter_media')).text
    media_info = literal_eval(ocw_embed_chapter_media[len('ocw_embed_chapter_media'):])

    youtube_link = media_info[1]
    srt_link = media_info[-1]

    if ARGS.archive:
        raise Exception('not implemented')
    else:
        ydl_opts = {
                'format': ARGS.format,
                'restrictfilenames': True,
                'writesubtitles': True,
                'outtmpl': name + '.%(ext)s',
                'progress_hooks': [postprocess],
                }
        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
            while True:
                try:
                    ydl.download([youtube_link])
                    break
                except Exception as e:
                    print e.message
                    pass

def crawl_links(hostname, path):
    def is_good_link(href):
        if not href.startswith('/'):
            parsed_url = urlparse.urlparse(href)
            if parsed_url.hostname != hostname: return False
            href = parsed_url.path

        if not href.lstrip('/').startswith(path): return False

        tokens = href.strip('/').split('/')
        if tokens[0] not in ('courses', 'resources'): return False
        return True

    video_links = []
    material_links = []

    link_queue = [path]
    visited_links = set()

    while len(link_queue) > 0:
        link = link_queue.pop(0)
        if link in visited_links: continue
        visited_links.add(link)
        print "visiting: ", link

        if re.search('\..+$', link) and not re.search('\.(htm|xml)$', link):
            material_links.append(link)
            continue

        url = '/'.join((hostname, link))
        soup = BeautifulSoup(urllib2.urlopen(url).read(), 'html.parser')

        if soup.find('script', string=re.compile('^ocw_embed_chapter_media')):
            title = soup.find('h1', class_='title').text.strip()
            video_links.append((title, link))
            continue

        for link in soup.find_all('a', href=is_good_link):
            link_queue.append(link['href'])
    return (video_links, material_links)

def get_videos(url):
    parsed_url = urlparse.urlparse(url)
    hostname = '%s://%s' % (parsed_url.scheme, parsed_url.netloc)

    tokens = parsed_url.path.strip('/').split('/')
    if not len(tokens) >= 3 and not tokens[0] == 'resources':
        raise Exception('Invalid course link!')
    basepath = '/'.join(tokens[:3])

    response = urllib2.urlopen(url)
    soup = BeautifulSoup(response.read(), 'html.parser')
    course_name = soup.find('h1', attrs={'class': 'title'}).text

    if not os.path.isdir(course_name):
        os.mkdir(course_name)
    os.chdir(course_name)

    video_links, material_links = crawl_links(hostname, basepath)

    print "Videos found: ", len(video_links), "videos."
    print "Material found: ", len(material_links), "items."

    for link in material_links:
        target = '/'.join(link.lstrip('/').split('/')[3:])
        print "Downloading material:", link, '-->', target
        if not os.path.exists(os.path.dirname(target)):
            os.makedirs(os.path.dirname(target))
        with open(target, 'w') as fh:
            fh.write(urllib2.urlopen('/'.join((hostname, link))).read())

    for idx, lecture in enumerate(video_links):
        name, link = lecture
        name = name.replace('/', ' ')
        get_lecture('%d - %s' % (idx, name), '/'.join((hostname, link)))

#--------------------------------------------------------------------------------
# top-level function

def download():
    for url in ARGS.url:
        get_videos(url)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()

    parser.add_argument('url',                  nargs='+', help='Course URL, such as "http://ocw.mit.edu/courses/mathematics/18-085-computational-science-and-engineering-i-fall-2008/"')
    parser.add_argument('--format', '-f',       action='store', default='best', help='Video / audio format, as specified by youtube-dl. Default is "best".')
    parser.add_argument('--postprocess', '-p',  action='store', default='', help='Post-processing command called after each file is downloaded. Default nothing.')
    parser.add_argument('--archive', '-a',      action='store_true', default=False, help='Download from archive, default off')

    ARGS = parser.parse_args()
    download()

