#!/usr/bin/env python

import argparse
import json
import os
import re
import warnings
from concurrent import futures
from typing import Any, Dict, Iterable, List, Optional, Union

import requests


class WenKuClient:
    session = requests.Session()
    concurrency = 8

    @staticmethod
    def load_jsonp(content: Union[str, bytes]) -> Any:
        """parse jsonp"""
        if isinstance(content, bytes):
            content = content.decode("u8")

        # find the first '(' and the last ')' and extract the substr between them
        lindex = content.find('(')
        rindex = content.rindex(')')
        return json.loads(content[lindex + 1: rindex])

    @staticmethod
    def parse_doc_id(url: str) -> Optional[str]:
        """parse doc_id from url with regex"""

        pat = re.compile(r'([0-9a-f]{8,})')
        result = pat.findall(url)
        if result:
            return result[0]

    @classmethod
    def get_doc_info(cls, doc_id: str) -> Dict[str, Any]:
        """get doc meta info describing the doc from api

        example url:
        https://wenku.baidu.com/api/doc/getdocinfo?doc_id=1a0992fe5e0e7cd184254b35eefdc8d377ee1453&callback=cb

        the response json like:
        ```json
            {
                "doc_id": "080f08af998fcc22bcd10d77",
                "docInfo": {
                    "docTitle": "doc title",
                },
                "mydoc": 0,
                "xz_top1": 0,
                "htmlBcs": false,
                "isHtml": true,
                "md5sum": "...",
                "bcsParam": [
                    {
                        "merge": "0-5863",
                        "zoom": "&png=0-0&jpg=0-0",
                        "page": 1
                    },

                ],
                "downloadToken": "...",
                "matchStatus": 1,
            }
        ```
        """
        url = 'https://wenku.baidu.com/api/doc/getdocinfo'
        params = {'callback': 'cb', 'doc_id': doc_id}
        content = requests.get(url, params=params).content
        return cls.load_jsonp(content)

    @classmethod
    def parse_image_urls(cls, doc_info: dict) -> List[str]:
        """parse image url from doc meta info

        `bcsParam` field describes image params for bcs(baidu cloud service)
        `md5sum` field describes sign params to fetch image

        the image url and params are described in `bcsParam` field
        """

        md5sum = doc_info['md5sum']
        doc_id = doc_info['doc_id']
        pat = re.compile(r'png=(\d+-\d+)&jpg=(\d+-\d+)')

        urls = []
        for param in doc_info['bcsParam']:
            img_fmt = 'jpg'

            zoom: str = param['zoom']
            res = pat.findall(zoom)

            if not res:
                continue

            png_range, jpg_range = res[0]

            # not image, when doc is in format of word or txt
            if jpg_range == '0-0' and png_range == '0-0':
                continue

            # png image
            if jpg_range == '0-0':
                img_fmt = 'png'
            url = f'https://wkretype.bdimg.com/retype/zoom/{doc_id}?o={ img_fmt }{ md5sum }{ zoom }'
            urls.append(url)
        return urls

    @classmethod
    def parse_doc_content(cls, raw: bytes) -> str:
        """parse text from doc info

        wenku describe every small fragments of word document with json
        """
        items = cls.load_jsonp(raw)['body']
        text = "".join([item['c'] for item in items if item['t'] == 'word'])
        return text

    @classmethod
    def get_text(cls, doc_id: dict) -> Iterable[str]:
        """get text of doc_id"""
        url = f'https://wenku.baidu.com/view/{ doc_id }'
        html = cls.session.get(url).content.decode('gb2312')

        pat = re.compile(r'WkInfo\.htmlUrls\s=\s\'(.*\})', re.IGNORECASE)
        res = pat.findall(html)
        if not res:
            raise ValueError
        json_text = res[0]
        json_text = json_text.replace(r"\x22", '"').replace(r'\\', '')
        doc: dict = json.loads(json_text)
        content_info: list = doc['json']
        urls = [item['pageLoadUrl'] for item in content_info]
        frags = []
        for data in cls.batch_fetch(urls):
            text = cls.parse_doc_content(data)
            frags.append(text)
        return "".join(frags)

    @classmethod
    def batch_fetch(cls, urls: List[str]) -> Iterable[bytes]:
        """fetch content from batch urls concurrency and in order by ThreadPoolExecutor"""
        def get_content(url: str) -> bytes:
            return requests.get(url).content

        tasks = []
        with futures.ThreadPoolExecutor(max_workers=cls.concurrency) as executor:
            for url in urls:
                task = executor.submit(get_content, url)
                tasks.append(task)

            for task in tasks:
                yield task.result()

    @staticmethod
    def mkdir(dirname: str):
        try:
            os.mkdir(dirname)
        except Exception as e:  # noqa
            pass

    @classmethod
    def fetch(cls, url: str):

        doc_id = cls.parse_doc_id(url)
        if not doc_id:
            warnings.warn(f"Can not get doc_id from { url }")
            return

        doc_info = cls.get_doc_info(doc_id)

        title = doc_info['docInfo']['docTitle']
        cls.mkdir(title)

        image_urls = cls.parse_image_urls(doc_info)
        results = cls.batch_fetch(image_urls)
        for idx, data in enumerate(results, 1):
            ext = 'jpg' if 'o=jpg' in image_urls[idx-1] else 'png'
            with open(f"{title}/page-{ idx }.{ ext }", 'wb') as f:
                f.write(data)

        with open(f'{title}/{ title }.txt', 'w') as f:
            f.write(cls.get_text(doc_id))


if __name__ == "__main__":

    description = '''Baidu Wenku Downloader. Usage: python wenku.py <url>'''
    cmd = argparse.ArgumentParser(description=description)
    cmd.add_argument("url", type=str)
    args = cmd.parse_args()
    WenKuClient.fetch(args.url)
