import requests
from bs4 import BeautifulSoup as bs
from tqdm import tqdm, trange
import json
import os
import tables
import pandas as pd
import numpy as np
import pdb
import traceback
import os
import multiprocessing as mp

import argparse



BROWSE_FANFIC = {
    'anime': 'https://www.fanfiction.net/anime/',
    'book':  'https://www.fanfiction.net/book/',
    'cartoon': 'https://www.fanfiction.net/cartoon/',
    'comic': 'https://www.fanfiction.net/comic/',
    'game': 'https://www.fanfiction.net/game/',
    'misc': 'https://www.fanfiction.net/misc/',
    'play': 'https://www.fanfiction.net/play/',
    'movie': 'https://www.fanfiction.net/movie/',
    'tv': 'https://www.fanfiction.net/tv/'
}

def list_universes():
    """
    List all story universes available on fanfiction.net and the number of stories they have
    """
    rows = []
    for fanfic_type, browse_page in tqdm(BROWSE_FANFIC.items(), position=0):

        page = bs(requests.get(browse_page).content, 'lxml')
        link_groups = page.find(id="list_output").find_all('div')

        for l in tqdm(link_groups, position=1):
            link = l.find('a')
            story_num = l.find('span').text.lstrip('(').rstrip(')')
            if story_num.endswith('K'):
                story_num = int(float(story_num.rstrip('K'))*1000)
            else:
                story_num = int(story_num)
            rows.append((fanfic_type, link.text, story_num, 'https://www.fanfiction.net'+link.get('href')))

    return pd.DataFrame.from_records(rows, columns=('type', 'universe', 'n_stories', 'url'))

def list_stories(row):
    """
    For a particular story universe, list all stories and return their metadata.

    Args:
        row (tables.Row): row containing metadata from list_universes()
    """
    try:
        i, row = row
        # FIXME: hardcoding base_dir
        base_dir = os.path.join(os.getcwd(), 'stories')
        save_fn = os.path.join(base_dir, row.universe.replace('/','-')+'.pck.gz')
        if os.path.exists(save_fn):
            return

        stories = []
        first_page = bs(requests.get(row.url+"?&srt=1&r=10").content, 'lxml')


        # find last page
        try:
            last_page_url = first_page.find('center').find('a', text='Last').get('href')
            last_page = int(last_page_url.split('&p=')[-1])
        except AttributeError:
            last_page = 1

        # get current process number to place progress bar
        try:
            current = mp.current_process()
            tqdm_position = current._identity[0]+1
        except:
            tqdm_position = 0

        story_pbar = tqdm(position=tqdm_position, total=row.n_stories)

        # iterate through pages listing stories and get their metadata
        for page_num in range(1,last_page+1):
            page_url = row.url + '?&srt=1&r=10&p={}'.format(page_num)
            page = bs(requests.get(page_url).content, 'lxml')

            links = page.find_all(class_='z-list')
            for l in links:
                story_url = l.find(class_='stitle').get('href')
                story_id = int(story_url.split('/')[2])
                story_name = l.find(class_='stitle').text
                story_url = 'https://fanfiction.net' + story_url

                stories.append((row.universe, story_id, story_url, story_name))
                story_pbar.update()

        story_df = pd.DataFrame.from_records(stories, columns=('universe', 'id', 'url', 'name'))
        story_df.to_pickle(save_fn, compression='gzip')
    except Exception as e:
        print(e)
        pdb.set_trace()

def list_hp(url):


    page = bs(requests.get(url).content, 'lxml')

    links = page.find_all(class_='z-list')
    stories = []
    for l in links:
        story_url = l.find(class_='stitle').get('href')
        story_id = int(story_url.split('/')[2])
        story_name = l.find(class_='stitle').text
        story_url = 'https://fanfiction.net' + story_url

        stories.append(('Harry Potter', story_id, story_url, story_name))

    story_df = pd.DataFrame.from_records(stories, columns=('universe', 'id', 'url', 'name'))
    return story_df



def list_harry_potter():
    pool = mp.Pool(12)

    base_url = 'https://www.fanfiction.net/book/Harry-Potter/'

    # FIXME: hardcoding base_dir
    base_dir = os.path.join(os.getcwd(), 'stories')

    save_fn = os.path.join(base_dir, 'Harry Potter'+'.pck.gz')
    if os.path.exists(save_fn):
        return
    first_page = bs(requests.get(base_url+"?&srt=1&r=10").content, 'lxml')


    # find last page
    try:
        last_page_url = first_page.find('center').find('a', text='Last').get('href')
        last_page = int(last_page_url.split('&p=')[-1])
    except AttributeError:
        last_page = 1

    page_urls = [base_url + '?&srt=1&r=10&p={}'.format(n) for n in range(1, last_page+1)]

    pbar = tqdm(total=len(page_urls), position=0)

    pdb.set_trace()
    results = pool.imap_unordered(list_hp, page_urls)

    pages = []
    for r in results:
        result = r.get()
        pages.append(result)
        pbar.update()

    hp_df = pd.concat(pages)
    hp_df.to_pickle(save_fn, compression="gzip")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--universes', help='location of the universes.json file that tracks which stories we need to reload')
    parser.add_argument('--stories', help='location of stories directory - links to all stories')
    parser.add_argument('--list', help='only list stories')
    args = parser.parse_args()

    if args.universes:
        universe_path = args.universes
    else:
        universe_path = 'universes.json'

    if args.stories:
        story_path = args.stories
    else:
        story_path = os.path.join(os.getcwd(),'stories')

    if not os.path.exists(story_path):
        os.mkdir(story_path)

    # list universes, if we have a previously available universes file,
    # find universes that have more stories than we had last time.
    universes = list_universes()
    try:
        past_universes = pd.read_json(universe_path)
        universes = universes.loc[universes.n_stories > past_universes.n_stories,]
    except:
        print('\ncouldnt load old universes directory, loading all\n')

    universes.to_json(universe_path)

    # start a pool of workers to list stories
    pool = mp.Pool(12)
    #
    # list_harry_potter()

    pbar = tqdm(total=universes.shape[0], position=0)

    results = pool.imap_unordered(list_stories, universes.iterrows())

    for r in results:
        try:
            r.get()
            pbar.update()
        except AttributeError:
            pbar.update()

import requests
from bs4 import BeautifulSoup as bs
from tqdm import tqdm, trange
import json
import os
import tables
import pandas as pd
import numpy as np
import pdb
import traceback
import os
import multiprocessing as mp

import argparse


class Metadata(tables.IsDescription):
    """
    Class to describe columns in metadata table
    """
    page_id = tables.UInt32Col()
    page_title = tables.StringCol(512)
    title = tables.StringCol(512)
    description = tables.StringCol(2048)
    rating = tables.StringCol(64)
    language = tables.StringCol(128)
    chapters = tables.UInt16Col()
    chapter = tables.UInt16Col()
    words = tables.UInt64Col()
    reviews = tables.UInt32Col()
    favs = tables.UInt32Col()
    follows = tables.UInt32Col()
    updated = tables.StringCol(64)
    published = tables.StringCol(64)
    text_idx = tables.UInt64Col()
    genre = tables.StringCol(256)
    characters = tables.StringCol(512)
    universe = tables.StringCol(1024)


def scrape_page(page, page_id, chapter=1, row=None):
    """
    Scrape a single page
    """

    # gather page elements with metadata
    # small sub element with hyphen-separated descriptors
    profile = page.find(id="profile_top")
    subheader = list(profile.find(class_='xgray xcontrast_txt').children)


    subhead_1 = [t.strip() for t in subheader[2].split('-')]
    subhead_2 = [t.strip() for t in subheader[4].split('-')]

    #############################
    # gather metadata -- idiosyncratic selection criteria for the page and its format.
    # after scraping several thousand pages, there are no obvious failures.
    metadata = {}

    # these ones are more or less always present, no position correction needed
    metadata['page_id'] = page_id
    metadata["page_title"] = page.find("title").text.encode('ascii', errors="ignore")
    #metadata["title"] = profile.find("b").text.encode('ascii', errors="ignore")
    metadata["description"] = profile.find('div', recursive=False).text.encode('ascii', errors="ignore")
    metadata["rating"] = subheader[1].text.encode('ascii', errors="ignore")
    metadata["language"], metadata["genre"], metadata["characters"] = subhead_1[1].encode('ascii', errors="ignore"), subhead_1[2].encode('ascii', errors="ignore"), subhead_1[3].encode('ascii', errors="ignore")
    metadata['chapter'] = int(chapter)

    # the following might be present, and so they require special checks
    try:
        metadata["chapters"] = [int(c.strip('Chapters: ')) for c in subhead_1 if c.startswith("Chapters:")][0]
    except:
        metadata['chapters'] = 1


    metadata["words"] = [int(c.strip('Words: ').replace(',', '')) for c in subhead_1 if c.startswith("Words: ")][0]

    try:
        reviews_item = [r for r in subheader if 'Reviews' in r][0]
        reviews_idx = subheader.index(reviews_item)
        metadata["reviews"] = subheader[reviews_idx + 1].text

    except:
        pass

    try:
        metadata["favs"] = [int(c.strip("Favs: ").replace(',', '')) for c in subhead_2 if c.startswith('Favs')][0]
    except:
        metadata['favs'] = 0

    try:
        metadata["follows"] = [int(c.strip("Follows: ").replace(',', '')) for c in subhead_2 if c.startswith('Follows')][0]
    except:
        metadata['follows'] = 0

    try:
        updated_item = [r for r in subheader if 'Published' in r][0]
        updated_idx = subheader.index(updated_item)
        metadata["updated"] = subheader[updated_idx+1].text
    except:
        pass

    try:
        published_item = [r for r in subheader if 'Published' in r][0]
        published_idx = subheader.index(published_item)
        metadata["published"] = subheader[published_idx+1].text
    except:
        pass

    #metadata["text_idx"] = texts.nrows

    # add data from the row
    metadata['universe'] = row.universe
    metadata['title'] = row['name']


    text = page.find(id="storytext").text

    return metadata, text

def scrape_story(row):
    """
    Scrape all pages of a story, including chapters.
    """
    try:
        i, row = row

        global do_multi
        if do_multi:
            current = mp.current_process()
            tqdm_position = current._identity[0]+1
        else:
            tqdm_position = 1

        try:
            url_id = str(row.id).zfill(7)
        except:
            url_id = row.id

        page_url = "https://www.fanfiction.net/s/{}/".format(url_id)
        page = bs(requests.get(page_url).content, 'lxml')

        # determine whether a story is had here, if not skip this page.
        try:
            warning = page.find(class_="gui_warning")
            if warning.text.startswith("Story Not Found"):
                with open('fanfic_log.txt', 'a') as f:
                    f.write("Story not found, id: {}".format(row.id))
                return [False], [False]
        except AttributeError:
            # didn't find warning, so Nonetype has no 'text'
            pass

        try:
            # new error - weird blank page
            warning = page.find(class_="gui_normal")
            if warning.text.endswith("Please check to see you are not using an outdated url."):
                with open('fanfic_log.txt', 'a') as f:
                    f.write("Outdated URL, id: {}".format(row.id))
                return [False], [False]
        except AttributeError:
            pass

        # figure out if story has chapters
        chapter_select = page.find(id="chap_select")
        if chapter_select:
            has_chapters = True
        else:
            has_chapters = False

        # scrape the first page and save data

        metadata = []
        text = []

        metadata_1, text_1 = scrape_page(page, page_id=url_id, row=row)
        metadata.append(metadata_1)
        text.append(text_1)

        # if the text has chapters, iterate.
        if has_chapters:
            n_chapters = len(list(chapter_select.children))
            pbar = tqdm(position=tqdm_position, total=n_chapters)
            for chapter_number in range(2, n_chapters + 1):
                chap_url = page_url + str(chapter_number) + "/"
                page = bs(requests.get(chap_url).content, 'lxml')
                metadata_chap, text_chap = scrape_page(page, page_id=url_id, chapter=chapter_number, row=row)
                metadata.append(metadata_chap)
                text.append(text_chap)
                pbar.update()

        return metadata, text

    except Exception as e:


        with open('fanfic_log.txt', 'a') as f:
            f.write(str(row))
            f.write(str(e))
            f.write(traceback.format_exc())

        return [False], [False]



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--list', help='index file')
    parser.add_argument('--output', help='output file path')
    args = parser.parse_args()

    if args.list:
        list_path = args.list
    else:
        list_path = 'story_idx.pck.gz'

    if args.output:
        output_path = args.output
    else:
        output_path = 'fanfiction.h5'

    print('\nreading story index')
    stories = pd.read_pickle(list_path, compression='gzip')
    print('\nstory index read, scraping {} stories'.format(stories.shape[0]))

    # open hdf5 file to write to

    filter = tables.Filters(complib='lzo', complevel=5)
    h5f = tables.open_file(output_path, mode="a", title="fanfiction.net", filters=filter)

    # if metadata table doesn't exist, make it, otherwise get the reference to it
    try:
        tab = h5f.create_table('/', "metadata", description=Metadata, filters=filter)
    except tables.exceptions.NodeError:
        tab = h5f.get_node('/', "metadata")

    # same with texts, though we'll use a variable-length unicode
    # format, which can only be a single column array.
    try:
        texts = h5f.create_vlarray(h5f.root, 'texts',
                                   tables.VLUnicodeAtom(), filters=filter)
    except tables.exceptions.NodeError:
        texts = h5f.get_node('/', 'texts')

    # remove stories that have already been scraped
    print('loading ids of stories that have already been scraped...')
    #scraped_ids = tab.col('page_id')
    all_scraped_ids = []
    scraped_ids = set()
    i = 0
    for row in tqdm(tab, total=tab.nrows):
        if i>0 and i % 1000000 == 0:
            all_scraped_ids.append(scraped_ids)
            scraped_ids = set()
        scraped_ids.add(row['page_id'])
        i += 1

    all_scraped_ids.append(scraped_ids)
    scraped_ids = set().union(*all_scraped_ids)

    print('looking for stories that have already been scraped (this might take a minute)...')

    already_got = stories['id'].isin(scraped_ids)
    print('\nAlready have {} stories, skipping those'.format(np.sum(already_got)))

    stories = stories[~already_got]


    # the row class allows us to write iteratively to the pytable
    tab_row = tab.row

    # figure out where we start & end. each page number is a 7-8 digit int.
    global do_multi
    do_multi = True

    pbar = tqdm(total=stories.shape[0], position=0)

    if not do_multi:

        for arow in stories.iterrows():
            metadata, text = scrape_story(arow)
            for meta, tex in zip(metadata, text):
                if not meta:
                    #print(meta, tex)
                    continue
                for k, v in meta.items():
                    if isinstance(v, str):
                        v = v.encode('ascii', errors='ignore')
                    try:
                        tab_row[k] = v
                    except TypeError:
                        # could be a number with a comma..
                        v = int(v.replace(b',', b''))
                        tab_row[k] = v

                tab_row.append()
                texts.append(str(tex))

            tab.flush()
            texts.flush()
            pbar.update()


    else:



        # create pool of workers
        pool = mp.Pool(24)



        results = pool.imap_unordered(scrape_story, stories.iterrows())

        for r in results:
            try:
                #try:
                #pdb.set_trace()
                metadata, text = r
                for meta, tex in zip(metadata, text):
                    if (not meta) and (not tex):
                        with open('fanfic_log.txt', 'a') as f:
                            f.write(str('skipping'))
                        continue
                    for k, v in meta.items():
                        if isinstance(v, str):
                            v = v.encode('ascii', errors='ignore')
                        try:
                            tab_row[k] = v
                        except TypeError:
                            # could be a number with a comma..
                            v = int(v.replace(b',', b''))
                            tab_row[k] = v
                    tab_row.append()
                    texts.append(str(tex))

                tab.flush()
                texts.flush()
                pbar.update()

            except Exception as e:
                with open('fanfic_log.txt', 'a') as f:
                    f.write(str(e))
                    f.write(traceback.format_exc())


    h5f.flush()
    h5f.close()