AO3-DL/python/process.py

#!/usr/bin/env python3
# This file is part of AOMO, "Archive Of My Own", a collection of Python and PHP
# scripts designed act as an improved local backup system
# for works published on https://archiveofourown.org
#
# Copyright (c) 2025  Cyberpro123, except where otherwise noted.
#
# This project is available under the GNU General Public License (GPL) v2.0,
# available at https://www.gnu.org/licenses/old-licenses/gpl-2.0.en.html
# or in the 'LICENSE.md' file that should be distributed alongside this one.
#
# Report issues to https://codeberg.org/Cyberpro123/AOMO
# Contact author at cyberpro123@posteo.com
import datetime
import logging
import os.path
import platform
import traceback

import bs4

import common
import sql


def makeStr(tag) -> str:
    tagStr = ""
    if isinstance(tag, str) or isinstance(tag, bs4.Tag):
        tagStr = str(tag)
    if isinstance(tag, list):
        for i in tag:
            tagStr += makeStr(tag=i)
    return tagStr


def processSingle(
    workID: int,
    config: dict,
    logger: logging.Logger,
) -> None:
    global countWorksTotal
    global countWorksStarted
    global countWorksComplete
    countWorksStarted += 1
    logger.info(
        f"Started Processing [{workID}] - [{countWorksStarted}] / [{countWorksTotal}]"
    )
    sqlInfo = {}
    rawFilename = os.path.join(config["dirRaws"], f"{workID}.html")
    if os.path.getsize(rawFilename) > 1000000:
        logger.warning(
            f"Work [{workID}]'s Raw is larger than 1 MB and may take several minutes to parse."
        )
    with open(rawFilename) as rawFile:
        rawSoup = bs4.BeautifulSoup(rawFile, features="lxml")

    ################################################################
    # Tag Isolation
    ################################################################
    ################################
    # TagsDiv
    ################################
    headerMeta = rawSoup.find(id="preface").div
    workTags = headerMeta.dl.find_all("dd")
    workTagHeadersRaw = headerMeta.dl.find_all("dt")
    workTagHeaders = []
    for tag in workTagHeadersRaw:
        tagStr = tag.string
        if tagStr == "Categories:":
            tagStr = "Category:"
        if tagStr[-2:] == "s:":
            tagStr = f"{tagStr[:-2]}:"
        workTagHeaders.append(tagStr)
    del workTagHeadersRaw
    sqlInfo["rating"] = workTags[workTagHeaders.index("Rating:")].a.string
    sqlInfo["warnings"] = []
    sqlInfo["categories"] = []
    sqlInfo["tagsFandom"] = []
    sqlInfo["tagsShips"] = []
    sqlInfo["tagsChara"] = []
    sqlInfo["tagsOther"] = []
    sqlInfo["language"] = workTags[workTagHeaders.index("Language:")].string
    for i in (
        ("warnings", "Archive Warning:", "Warnings!?!?!"),
        ("categories", "Category:", "Categories"),
        ("tagsFandom", "Fandom:", "Fandom Tags?!?!"),
        ("tagsShips", "Relationship:", "Ship Tags"),
        ("tagsChara", "Character:", "Character Tags"),
        ("tagsOther", "Additional Tag:", "General Tags"),
    ):
        try:
            for tag in workTags[workTagHeaders.index(i[1])].find_all("a"):
                sqlInfo[i[0]].append(tag.string)
        except ValueError:
            logger.debug(f"Work [{workID}] doesn't seem to have any [{i[2]}]")
    seriesList = []
    try:
        for i in workTags[workTagHeaders.index("Serie:")]:
            seriesList.append(i)
    except ValueError:
        logger.debug(f"Work [{workID}] doesn't seem to be part of any series.")
    sqlInfo["serieses"] = []
    for num, tag in enumerate(seriesList):
        if isinstance(tag, bs4.element.Tag) and tag.name == "a":
            hrefSplit = tag["href"].split("/")
            seriesID = int(hrefSplit[1 + hrefSplit.index("series")])
            for chunk in str(seriesList[num - 1]).split():
                if chunk.isdigit():
                    seriesPos = int(chunk)
            sqlInfo["serieses"].append((seriesID, seriesPos, tag.string))
            del hrefSplit, seriesID, seriesPos
    del seriesList
    statsSplit = workTags[workTagHeaders.index("Stat:")].string.split()
    chapterNOs = statsSplit[statsSplit.index("Chapters:") + 1].split("/")
    sqlInfo["nchapters"] = int(chapterNOs[0])
    try:
        sqlInfo["chaptersExpected"] = int(chapterNOs[1])
    except ValueError:
        sqlInfo["chaptersExpected"] = 0
    del chapterNOs
    sqlInfo["dateUp"] = 0
    for i in ("Completed:", "Updated:", "Published:"):
        if not sqlInfo["dateUp"]:
            try:
                sqlInfo["dateUp"] = int(
                    datetime.datetime.fromisoformat(
                        statsSplit[1 + statsSplit.index(i)]
                    ).timestamp()
                )
            except ValueError:
                pass
    if not sqlInfo["dateUp"]:
        logger.error(f"Date Updated couldn't be found in work [{workID}]")
    sqlInfo["datePb"] = int(
        datetime.datetime.fromisoformat(
            statsSplit[1 + statsSplit.index("Published:")]
        ).timestamp()
    )
    sqlInfo["dateDl"] = 0
    sqlInfo["dateEd"] = 0
    version, sqlInfo["dateEd"], sqlInfo["dateDl"] = (
        common.getMetadataFromRawCommentByID(id=workID, config=config)
    )
    sqlInfo["wordCount"] = statsSplit[1 + statsSplit.index("Words:")]
    del version
    ################################
    # Rest of Header
    ################################
    sqlInfo["text"] = {
        0: {
            "title": "",
            "summary": "",
            "notesStart": "",
            "notesEnd": "",
            "body": "",
        }
    }
    sqlInfo["text"][0]["title"] = headerMeta.h1.string
    headerBlockquotes = headerMeta.find_all("blockquote", class_="userstuff")
    match len(headerBlockquotes):
        case 2:
            sqlInfo["text"][0]["notesStart"] = makeStr(
                tag=headerBlockquotes[1].contents,
            )
            sqlInfo["text"][0]["summary"] = makeStr(
                tag=headerBlockquotes[0].contents,
            )
        case 1:
            match headerMeta.p.string:
                case "Summary":
                    hbqDest = "summary"
                case "Notes":
                    hbqDest = "notesStart"
            sqlInfo["text"][0][hbqDest] = makeStr(
                tag=headerBlockquotes[0].contents,
            )
            del hbqDest
        case 0:
            logger.debug(
                f"Work [{workID}] seems to have neither a summary nor start notes."
            )
    sqlInfo["authors"] = []
    for tag in headerMeta.div.find_all("a"):
        paraOpenLoc = tag.string.rfind("(")
        paraCloseLoc = tag.string.rfind(")")
        if paraOpenLoc < paraCloseLoc and paraOpenLoc != -1 and paraCloseLoc != -1:
            # pseud detected
            sqlInfo["authors"].append(tag.string[paraOpenLoc + 1: paraCloseLoc])
        else:
            # not a pseud
            sqlInfo["authors"].append(tag.string)
        del paraOpenLoc
        del paraCloseLoc
    try:
        sqlInfo["text"][0]["notesEnd"] = makeStr(
            tag=rawSoup.find(id="afterword").find(id="endnotes").blockquote.contents,
        )
    except AttributeError:
        logger.debug(f"Work [{workID}] doesn't seem to have work end notes")
    allPrefaces = rawSoup.select("#chapters > div.meta.group")
    allMains = rawSoup.select("#chapters > div.userstuff")
    allPostfaces = {}
    for i in range(1, 1 + len(allMains)):
        try:
            allPostfaces[i] = rawSoup.find(id=f"endnotes{i}")
        except KeyError:
            logger.debug(f"Work [{workID}] Chapter [{i}] doesn't seem to have endnotes")
    logger.debug(
        f"Pref/Main/Post for [{workID}]: [{len(allPrefaces)}] [{len(allMains)}] [{len(allPostfaces)}]"
    )
    cssFileDir = os.path.join(config["dirRaws"], f"{workID}.css")
    if os.path.exists(cssFileDir):
        with open(cssFileDir) as file:
            for line in file:
                sqlInfo["text"][0]["body"] += line
    for chaptNum in range(1, 1 + len(allMains)):
        sqlInfo["text"][chaptNum] = {
            "title": "",
            "summary": "",
            "notesStart": "",
            "notesEnd": "",
            "body": "",
        }
        ################################
        # Real
        ################################
        if len(allPrefaces) != 0:
            sqlInfo["text"][chaptNum]["title"] = makeStr(
                tag=allPrefaces[chaptNum - 1].h2.string,
            )
            preBQs = allPrefaces[chaptNum - 1].find_all("blockquote")
            if len(preBQs) == 2:
                sqlInfo["text"][chaptNum]["summary"] = makeStr(
                    tag=preBQs[0].contents,
                )
                sqlInfo["text"][chaptNum]["notesStart"] = makeStr(
                    tag=preBQs[1].contents,
                )
            elif len(preBQs) == 1:
                pCandidates = allPrefaces[chaptNum - 1].select("div.meta.group > p")
                match pCandidates[-1].string:
                    case "Chapter Notes":
                        bqDest = "notesStart"
                    case "Chapter Summary":
                        bqDest = "summary"
                sqlInfo["text"][chaptNum][bqDest] = makeStr(
                    tag=preBQs[0].contents,
                )
                del pCandidates
                del bqDest
            del preBQs
        elif len(allMains) == 1 and sqlInfo["nchapters"] == 1:
            logger.debug(f"Work [{workID}] seems to be a oneshot")
        else:
            logger.error(
                f"Work [{workID}], which doesn't seem to be a oneshot, had a len(allPrefaces) of 0"
            )
        sqlInfo["text"][chaptNum]["body"] = makeStr(
            tag=allMains[chaptNum - 1].contents,
        )
        tail = allPostfaces.get(chaptNum)
        if tail:
            sqlInfo["text"][chaptNum]["notesEnd"] = makeStr(
                tag=tail.blockquote.contents,
            )

    sql.addWork(
        id=workID,
        info=sqlInfo,
        config=config,
        logger=logger,
    )
    countWorksComplete += 1
    logger.info(
        f"Completed Processing [{workID}] - [{countWorksComplete}] / [{countWorksTotal}]"
    )


def multithreading(
    workIDs: set,
    config: dict,
    logger: logging.Logger,
) -> None:
    if (len(workIDs) < 10) or (platform.system() == "Windows"):
        for workID in workIDs:
            processSingle(
                workID=workID,
                config=config,
                logger=logger,
            )
    else:
        import concurrent.futures

        futures = {}
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=10, thread_name_prefix=common.threadNameBulk
        ) as pool:
            for workID in workIDs:
                futures[workID] = pool.submit(
                    processSingle, workID=workID, config=config, logger=logger
                )
        logger.info("Mostly complete, getting exceptions from futures...")
        for workID in futures:
            try:
                futures[workID].result()
            except FileNotFoundError as ex:
                split = str(ex).split()
                string = ""
                for num, i in enumerate(split):
                    if num == 7:
                        string += i
                    elif num > 7:
                        string += f" {i}"
                common.logger.warning(f"Work [{workID}] couldn't find file [{string}]")
                del split
                del string
            except Exception as ex:
                common.logger.error(
                    f"Work [{workID}] raised [{type(ex).__name__}]: [{ex.args}]"
                )
                traceback.print_exception(ex, file=common.args.outfile)
    logger.info("Complete, process.multithreding exiting")


if __name__ == "__main__":
    common.init(
        description="Extracts information from 'Raw' HTML files and stores it in the database to be used later.",
        all="Process every workid.html file in the Raws directory",
        auto="Process every workid.html file in the Raws directory that hasn't already been processed",
    )
    workIDs = set(())
    if common.args.auto:
        editTimes = sql.getAllEditTimePerWorkID(
            config=common.config, logger=common.logger
        )
        for filename in os.listdir(common.config["dirRaws"]):
            if filename[-5:] == ".html" and filename[:-5].isdecimal():
                workID = int(filename[:-5])
                timestampDB = editTimes.get(workID, 0)
                timestampRaw = common.getMetadataFromRawCommentByID(
                    workID, common.config
                )[1]
                if timestampRaw == timestampDB != 0:
                    common.logger.debug(
                        f"Work [{workID}]'s Raw has the same last edited timestamp [{timestampRaw}] as what's recorded in database, skipping"
                    )
                else:
                    workIDs.add(workID)
                del timestampRaw, timestampDB
    elif common.args.all:
        for filename in os.listdir(common.config["dirRaws"]):
            if filename[-5:] == ".html" and filename[:-5].isdecimal():
                workIDs.add(int(filename[:-5]))
    else:
        inRaw = common.parseInfile(common.args.infile)
        for line in inRaw:
            try:
                workIDs.add(int(line))
            except ValueError:
                common.logger.error(f"Could not convert input [{line}] to int")

    countWorksTotal = len(workIDs)
    countWorksStarted = 0
    countWorksComplete = 0
    multithreading(
        workIDs=workIDs,
        config=common.config,
        logger=common.logger,
    )