#!/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 concurrent.futures
import traceback

import bs4
import requests
import requests.exceptions

import common
import download


def getWorksFromSeriesPage(
    soup: bs4.BeautifulSoup,
    logger: common.logging.Logger,
) -> list:
    return getWorksFromPage(mode="series", soup=soup, logger=logger)


def getWorksFromAuthorPage(
    soup: bs4.BeautifulSoup,
    logger: common.logging.Logger,
) -> list:
    return getWorksFromPage(mode="author", soup=soup, logger=logger)


def getWorksFromPage(
    mode: str,
    soup: bs4.BeautifulSoup,
    logger: common.logging.Logger,
) -> list:
    # Significant portions of this function's code are derrived from 'AO3 API',
    # https://github.com/wendytg/ao3_api/
    # as of commit 02e349985d927bd8693f905f440e1ef0539f1984
    # available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
    match mode:
        case "series":
            target = soup.find("ul", {"class": "series work index group"})
        case "author":
            target = soup.find("ol", {"class": "work index group"})
    workIDs = []
    for li in target.find_all("li", {"role": "article"}):
        for i in li.h4.a["href"].split("/"):
            try:
                int(i)
            except ValueError:
                pass
            else:
                workIDs.append(int(i))
    return workIDs


def getAllWorksFromSeries(
    seriesID: int,
    session: requests.Session,
    logger: common.logging.Logger,
    retries: int = common.loopRetries,
    returnSoup: bool = False,
) -> list | tuple:
    # Significant portions of this function's code are derrived from 'AO3 API',
    # https://github.com/wendytg/ao3_api/
    # as of commit 02e349985d927bd8693f905f440e1ef0539f1984
    # available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
    soup1 = bs4.BeautifulSoup(
        download.getPage(
            url=f"https://archiveofourown.org/series/{seriesID}",
            logger=logger,
            retries=retries,
            session=session,
            tryAnon=False,
        ),
        "lxml",
    )
    work_count = 0
    stats = (
        soup1.find("dl", {"class": "series meta group"})
        .find("dl", {"class": "stats"})
        .find_all()
    )
    tail = None
    for curr in stats:
        if curr.name == "dt":
            tail = curr.string
        elif tail == "Works:":
            work_count = int(curr.string.replace(",", ""))
    del stats
    del tail
    page_count = int(work_count / 20) + int(bool(work_count % 20))
    workIDs = getWorksFromSeriesPage(soup=soup1, logger=logger)
    if not returnSoup:
        del soup1
    for pageNo in range(2, 1 + page_count):
        workIDs.extend(
            getWorksFromSeriesPage(
                soup=bs4.BeautifulSoup(
                    download.getPage(
                        url=f"https://archiveofourown.org/series/{seriesID}?page={pageNo}",
                        logger=logger,
                        retries=retries,
                        session=session,
                        tryAnon=False,
                    ),
                    "lxml",
                ),
                logger=logger,
            )
        )
    if len(workIDs) != work_count:
        logger.error(
            f"Series [{seriesID}] claims to have [{work_count}] works, but [{len(workIDs)}] were found."
        )
    if returnSoup:
        return workIDs, soup1
    else:
        return workIDs


def getAllWorksFromAuthor(
    username: str,
    session: requests.Session,
    logger: common.logging.Logger,
    retries: int = common.loopRetries,
) -> list:
    # Significant portions of this function's code are derrived from 'AO3 API',
    # https://github.com/wendytg/ao3_api/
    # as of commit 02e349985d927bd8693f905f440e1ef0539f1984
    # available under the MIT License Copyright (c) 2019 Francisco Patrício Rodrigues
    soup1 = bs4.BeautifulSoup(
        download.getPage(
            url=f"https://archiveofourown.org/users/{username}/works",
            logger=logger,
            retries=retries,
            session=session,
            tryAnon=False,
        ),
        "lxml",
    )
    work_count = 0
    h2 = soup1.find(
        "div", {"class": "works-index dashboard filtered region"}
    ).h2.text.split()
    if "of" in h2:
        target = h2[1 + h2.index("of")]
    else:
        target = h2[0]
    work_count = int(target.replace(",", ""))
    del target
    del h2
    page_count = int(work_count / 20) + int(bool(work_count % 20))
    workIDs = getWorksFromAuthorPage(soup=soup1, logger=logger)
    del soup1
    for pageNo in range(2, 1 + page_count):
        workIDs.extend(
            getWorksFromAuthorPage(
                soup=bs4.BeautifulSoup(
                    download.getPage(
                        url=f"https://archiveofourown.org/users/{username}/works?page={pageNo}",
                        logger=logger,
                        retries=retries,
                        session=session,
                        tryAnon=False,
                    ),
                    "lxml",
                ),
                logger=logger,
            )
        )
    if len(workIDs) != work_count:
        logger.error(
            f"Author [{username}] claims to have [{work_count}] works, but [{len(workIDs)}] were found."
        )
    return workIDs


if __name__ == "__main__":
    common.init(
        description="Takes links to AO3 pages and returns the corresponding Work IDs.",
        secondOut="WorkID Output",
    )
    config = common.config

    workIDs = set(())
    seriesIDs = set(())
    authors = set(())

    inRaw = common.parseInfile(common.args.infile)
    for line in inRaw:
        if line.isdecimal():
            common.logger.warning(
                f"Plain int [{line}] passed to parser.py, assuming it is a workID."
            )
            workIDs.add(int(line))
        elif line[1:].isdecimal():
            match line[:1].lower():
                case "w":
                    workID = int(line[1:])
                    common.logger.debug(f"Parsed [{line}] into workID [{workID}]")
                    workIDs.add(workID)
                case "s":
                    seriesID = int(line[1:])
                    common.logger.debug(f"Parsed [{line}] into series ID [{seriesIDs}]")
                    seriesIDs.add(seriesID)
        elif line[:4] == "http" or -1 < line.find("archiveofourown.org") <= 12:
            split = line.split("/")
            match split[3]:
                case "works":
                    workID = int(split[4].split("?")[0])
                    common.logger.debug(f"Parsed [{line}] into workID [{workID}]")
                    workIDs.add(workID)
                case "series":
                    seriesID = int(split[4].split("?")[0])
                    common.logger.debug(f"Parsed [{line}] into seriesID [{seriesID}]")
                    seriesIDs.add(seriesID)
                case "users":
                    author = str(split[4].split("?")[0])
                    common.logger.debug(f"Parsed [{line}] into author [{author}]")
                    authors.add(author)
        else:
            authors.add(line)
            common.logger.warning(
                f"Plain string [{line}] passed to parser.py, assuming it is an author username."
            )

    futuresSeries = {}
    futuresAuthors = {}

    session = download.login(config, common.logger)

    with concurrent.futures.ThreadPoolExecutor(
        max_workers=10, thread_name_prefix=common.threadNameBulk
    ) as pool1:
        for seriesID in seriesIDs:
            futuresSeries[seriesID] = pool1.submit(
                getAllWorksFromSeries,
                seriesID,
                session=session,
                logger=common.logger,
            )
        for author in authors:
            futuresAuthors[author] = pool1.submit(
                getAllWorksFromAuthor,
                author,
                session=session,
                logger=common.logger,
            )

    common.logger.info("Retreiving from futures...")

    seriesContents = []
    authorContents = []
    for series in futuresSeries:
        try:
            result = futuresSeries[series].result()
        except (Exception,) as ex:
            common.logger.error(
                f"Series [{series}] raised [{type(ex).__name__}]: [{ex.args}]"
            )
            traceback.print_exception(ex, file=common.args.outfile)
        else:
            seriesContents.append(result)
    for author in futuresAuthors:
        try:
            result = futuresAuthors[author].result()
        except (Exception,) as ex:
            common.logger.error(
                f"Author [{author}] raised [{type(ex).__name__}]: [{ex.args}]"
            )
            traceback.print_exception(ex, file=common.args.outfile)
        else:
            authorContents.append(result)

    for series in seriesContents:
        for workID in series:
            workIDs.add(workID)
    for author in authorContents:
        for workID in author:
            workIDs.add(workID)
    for workID in sorted(workIDs):
        common.args.batch_out.write(f"{workID}\n")