import os
import re
import sys
import argparse
from pathlib import Path

import AO3
from dotenv import load_dotenv

load_dotenv()


def sanitize_filename(name: str) -> str:
    """Remove characters that are invalid in filenames."""
    return re.sub(r'[\\/*?:"<>|]', "", name).strip()


def get_session() -> AO3.Session | None:
    username = os.getenv("AO3_USERNAME")
    password = os.getenv("AO3_PASSWORD")

    if not username or not password:
        print("No AO3_USERNAME/AO3_PASSWORD found in .env — proceeding without login.")
        return None

    try:
        session = AO3.Session(username, password)
        print(f"Logged in as {username}")
        return session
    except Exception as e:
        print(f"Login failed: {e}")
        return None


def download_work(workid: int, filetype: str = "EPUB", output_dir: str = ".", session=None):
    work = AO3.Work(workid, session=session)

    print(f"Title: {work.title}")
    print(f"Author(s): {', '.join(a.username for a in work.authors)}")
    print(f"Summary:\n{work.summary}\n")

    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    filename = out_path / f"{sanitize_filename(work.title)}.{filetype.lower()}"

    with open(filename, "wb") as file:
        file.write(work.download(filetype))

    print(f"Downloaded: {filename}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Download an AO3 work.")
    parser.add_argument("workid", type=int, nargs="?", help="AO3 work ID")
    parser.add_argument("-t", "--type", default="EPUB", help="File type (EPUB, PDF, MOBI, HTML, AZW3)")
    parser.add_argument("-o", "--output", default=os.getenv("AO3_OUTPUT_DIR", "."), help="Output directory")
    args = parser.parse_args()

    workid = args.workid or int(input("Enter AO3 work ID: "))

    session = get_session()
    print(session)
    download_work(workid, filetype=args.type, output_dir=args.output, session=session)