main.py

# /// script
# dependencies = [
#   "requests<3",
#   "google-api-python-client",
#   "youtube-transcript-api",
#   "openai",
#   "python-dotenv",
# ]
# ///

## NEW NAME: FlashCraft

import argparse
import email
import imaplib
import json
import logging
import math
import os
import re
import subprocess
import time
from email.header import decode_header

import json_repair
import requests
import tiktoken
from dotenv import load_dotenv
from openai import OpenAI

from content_types.books import process_book_attachment
from content_types.youtube import (
    extract_transcript_from_youtube,
    get_youtube_video_details,
)

# Load environment variables from .env file
load_dotenv(".env")
# Get environment variables
IMAP_SERVER = os.getenv("IMAP_SERVER")
EMAIL = os.getenv("EMAIL")
PASSWORD = os.getenv("PASSWORD")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
YOUTUBE_DATA_API_KEY = os.getenv("YOUTUBE_DATA_API_KEY")
ANKI_CONNECT_URL = os.getenv("ANKI_CONNECT_URL")
ANKI_API_KEY = os.getenv("ANKI_API_KEY")

# Parse command-line arguments
parser = argparse.ArgumentParser(description="FlashCraft: Generate flashcards content.")
parser.add_argument(
    "-m",
    "--model",
    type=str,
    default="gpt-4o",
    help="The LLM model to use for OpenAI API calls.",
    choices=["gpt-4o", "gpt-4o-mini"],
)
args = parser.parse_args()

# Use the model argument in the openai_call function
LLM_MODEL = args.model

# Create an instance of the OpenAI client
OPENAI_CLIENT = OpenAI(api_key=OPENAI_API_KEY)

# Configure logging
logging.basicConfig(
    level=logging.INFO,  # Set the logging level
    format="%(asctime)s - %(levelname)s - %(message)s",  # Log format
)

# Define the list of supported ebook file extensions
# SUPPORTED_TEXT_FILE_EXTENSIONS = (
#     ".azw",
#     ".azw3",
#     ".cbz",
#     ".cbr",
#     ".cbc",
#     ".chm",
#     ".docx",
#     ".epub",
#     ".fb2",
#     ".html",
#     ".htmlz",
#     ".lit",
#     ".lrf",
#     ".mobi",
#     ".odt",
#     # ".pdf",
#     ".pdb",
#     ".pml",
#     ".prc",
#     ".rb",
#     ".rtf",
#     ".snb",
#     ".tcr",
#     ".txt",
#     ".txtz",
# )


def check_email():
    """Check the email inbox for unread emails containing YouTube links or book attachments.

    Returns:
        dict: A dictionary with keys 'youtube' and 'books'.
            'youtube' is a list of YouTube video IDs extracted from unread emails.
            'books' is a list of file paths to downloaded .epub or .mobi files.
    """
    logging.info("Checking email for content...")

    # Connect to the email server
    mail = imaplib.IMAP4_SSL(IMAP_SERVER)
    mail.login(EMAIL, PASSWORD)
    mail.select("inbox")

    # Search for all unread emails
    status, messages = mail.search(None, "UNSEEN")
    email_ids = messages[0].split()

    # Pattern to extract YouTube video ID from either full URL or shortened URL
    pattern = (
        r"(?:https?://(?:www\.)?youtube\.com/watch\?v=|https?://youtu\.be/)([\w-]+)"
    )

    youtube_video_ids = []
    book_file_paths = []
    documents_file_paths = []

    for email_id in email_ids:
        # Fetch the email by ID
        status, msg_data = mail.fetch(email_id, "(RFC822)")
        msg = email.message_from_bytes(msg_data[0][1])

        # Get the email subject
        subject, encoding = decode_header(msg["Subject"])[0]
        if isinstance(subject, bytes):
            subject = subject.decode(encoding if encoding else "utf-8")

        # Check if the email contains a YouTube link or book attachments
        if msg.is_multipart():
            for part in msg.walk():
                content_type = part.get_content_type()
                if content_type == "text/plain":
                    body = part.get_payload(decode=True).decode()
                    video_ids = re.findall(pattern, body)
                    youtube_video_ids.extend(video_ids)
                elif (
                    part.get_content_maintype() == "application" and part.get_filename()
                ):
                    filename = part.get_filename()
                    if filename.lower().endswith(("epub", "mobi", ".pdf")):
                        filepath = os.path.join("downloads", filename)
                        if not os.path.exists("downloads"):
                            os.makedirs("downloads")
                        with open(filepath, "wb") as f:
                            f.write(part.get_payload(decode=True))
                        # Check if the file is a book or a document
                        if filename.lower().endswith((".epub", ".mobi")):
                            book_file_paths.append(filepath)
                        else:
                            documents_file_paths.append(filepath)
                    else:
                        logging.info(
                            f"Unsupported file format for attachment: {filename}"
                        )
        else:
            body = msg.get_payload(decode=True).decode()
            video_ids = re.findall(pattern, body)
            youtube_video_ids.extend(video_ids)

    # Close the connection and logout
    mail.close()
    mail.logout()

    logging.info(
        f"Found {len(email_ids)} unread emails with:"
        f"\n\t{len(youtube_video_ids)} YouTube links."
        f"\n\t{len(book_file_paths)} books."
        f"\n\t{len(documents_file_paths)} documents."
    )
    # Remove duplicate video IDs
    unique_youtube_video_ids = list(set(youtube_video_ids))

    # If there were duplicate video IDs, log the number of duplicates removed
    if len(unique_youtube_video_ids) < len(youtube_video_ids):
        logging.info(
            f"Removed {len(youtube_video_ids) - len(unique_youtube_video_ids)} duplicate video IDs."
        )
    return {
        "youtube": unique_youtube_video_ids,
        "books": book_file_paths,
        "documents": documents_file_paths,
    }


def openai_call(prompt, model=LLM_MODEL):
    """Call the OpenAI API with a given prompt.

    Args:
        prompt (str): The prompt to send to the OpenAI API.
        model (str): The model to use for the API call.

    Returns:
        str: The response content from the OpenAI API.
    """
    chat_completion = OPENAI_CLIENT.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model=model,
    )

    return chat_completion.choices[0].message.content


def summarize_transcript(transcript):
    """Summarize a YouTube video transcript using the OpenAI API.

    Args:
        transcript (str): The transcript of the YouTube video.

    Returns:
        str: The summary of the transcript.
    """
    logging.info("Summarizing the YouTube video transcript with OpenAI API...")

    # Define the prompt with the preparatory instruction and append the transcript
    with open("prompts/summarization.txt", "r") as file:
        prompt = file.read().strip()

    full_prompt = f"{prompt}: {transcript}"

    message = openai_call(full_prompt.replace("\n", ""))

    # Return the summary
    return message


def calc_num_questions(num_tokens):
    """Calculate the number of questions to generate based on the number of tokens.

    Args:
        num_tokens (int): The number of tokens in the text.

    Returns:
        int: The number of questions to generate.
    """
    return round(math.sqrt(num_tokens) / 25)


def generate_flashcards(text, language="english"):
    """Generate flashcards from a text.

    Args:
        text (str): The input text.
        language (str): The language for the flashcards.

    Returns:
        dict: A JSON object containing the generated flashcards.
    """
    logging.info("Generating flashcards from the input text...")

    with open("prompts/flashcard_generation.txt", "r") as file:
        prompt = file.read().strip()

    flashcards_prompt = f"{prompt}\n\n{text}"

    # Insert the selected language into the prompt
    flashcards_prompt = flashcards_prompt.replace("[LANGUAGE]", language)

    # Count the number of tokens in the flashcards_prompt
    num_tokens = len(flashcards_prompt.split())

    encoding = tiktoken.encoding_for_model(LLM_MODEL)
    num_tokens = len(encoding.encode(flashcards_prompt))
    logging.info(f"Number of tokens: {num_tokens}")

    # Check if the number of tokens exceeds 200000 (limit for gpt-4o-mini)
    if num_tokens > 200_000:
        logging.error("The number of tokens exceeds the limit for the selected model.")

    # Adjust number of questions based on the number of tokens
    num_questions = calc_num_questions(num_tokens)
    logging.info(f"Number of questions to generate: {num_questions}")
    flashcards_prompt = flashcards_prompt.replace("[NUM_QUESTIONS]", str(num_questions))

    # Call the OpenAI API to generate flashcards with the appropriate model
    flashcards = openai_call(flashcards_prompt, model="gpt-4o-mini")

    # Clean the output
    flashcards = (
        flashcards.replace("```json\n", "").replace("```", "").replace("\n", "").strip()
    )

    # Convert to json
    flashcards_json = json.loads(flashcards)

    # Placeholder for your OpenAI API code to generate flashcards
    return flashcards_json


def process_flashcards(flashcards):
    """Process the generated flashcards. It send them to gpt-4o for verification and correction.
    It ensures that the questions are not repeated and that no information is overlapped between questions.
    It also makes sure that the essential information is kept.

    Args:
        flashcards (dict): The generated flashcards.

    Returns:
        dict: The processed flashcards.
    """
    # TODO Placeholder for your code to process the flashcards
    return flashcards


def generate_tags(text):
    """Generate tags from a given text using the OpenAI API.

    Args:
        text (str): The text to generate tags from.

    Returns:
        list: A list of generated tags.
    """
    with open("prompts/tags_generation.txt", "r") as file:
        prompt = file.read().strip()

    tags_prompt = f"{prompt}\n\n{text}"

    # Call the OpenAI API to generate tags
    tags_response = openai_call(tags_prompt)
    print(tags_response)

    # Convert the response to a list of tags
    tags = tags_response.split(" ")

    return tags


def save_to_file(content: str, path: str = ""):
    """Save content to a file with a title derived from the content.

    Args:
        content (str): The content to save.
        path (str): The directory path to save the file in.
    """
    # Save the summarized transcript to a file
    # Get the title from first line
    title = (
        content.split("\n")[0]
        .replace("#", "")
        .replace("*", "")
        .strip()
        .lower()
        .replace(" ", "_")
    )
    with open(os.path.join(path, f"{title}.md"), "w") as file:
        file.write(content)


def deck_exists(deck_name):
    """Check if a specified Anki deck exists.

    Args:
        deck_name (str): The name of the deck to check.

    Returns:
        bool: True if the deck exists, False otherwise.
    """
    payload = {
        "action": "deckNames",
        "version": 6,
        "key": ANKI_API_KEY,
    }
    response = requests.post("http://localhost:8765", json=payload)
    if response.status_code != 200:
        raise Exception("Failed to fetch deck names from AnkiConnect")
    deck_names = response.json().get("result", [])
    return deck_name in deck_names


def create_deck(deck_name):
    """Create a new Anki deck.

    Args:
        deck_name (str): The name of the deck to create.
    """
    payload = {
        "action": "createDeck",
        "version": 6,
        "params": {"deck": deck_name},
        "key": ANKI_API_KEY,
    }
    response = requests.post("http://localhost:8765", json=payload)
    if response.status_code != 200:
        raise Exception("Failed to create deck in AnkiConnect")


def add_anki_card(deck_name, note_type, front, back, tags=None):
    """Add a new card to an Anki deck.

    Args:
        deck_name (str): The name of the deck to add the card to.
        note_type (str): The type of the note (card).
        front (str): The front content of the card.
        back (str): The back content of the card.
        tags (list, optional): A list of tags for the card.
    """

    if not deck_exists(deck_name):
        create_deck(deck_name)

    # Define the card note structure
    note = {
        "deckName": deck_name,
        "modelName": note_type,
        "fields": {"Front": front, "Back": back},
        "tags": tags or [],
        "options": {"allowDuplicate": False},
        "audio": [],
        "video": [],
        "picture": [],
    }

    send_anki_request("addNote", {"note": note})


def send_anki_request(action, params=None):
    """Send a request to AnkiConnect.

    Args:
        action (str): The action to perform.
        params (dict, optional): The parameters for the action.
    """
    # Prepare the request payload
    payload = {
        "action": action,
        "version": 6,
        "params": params or {},
        "key": ANKI_API_KEY,
    }

    # Send the request to AnkiConnect
    response = requests.post(ANKI_CONNECT_URL, json=payload)

    # Check the response
    if response.status_code == 200:
        result = response.json()
        if not ("error" in result and result["error"] is None):
            print(f"Error syncing media: {result['error']}")
    else:
        print(f"HTTP Error: {response.status_code}")


def improve_flashcards(flashcards):
    """Improve the generated flashcards by removing duplicates and overlapping information.

    Args:
        flashcards (dict): The generated flashcards.

    Returns:
        dict: The improved flashcards.
    """
    # Placeholder for your code to improve the flashcards

    # TODO An idea I had is to here do a second pass on the flashcards to make sure that the
    # questions are not repeated or that no information is overlapped between questions which
    # sometimes happens. Prompt GPT4 again with the flashcards and ask it to check for repeated
    # information and modify them to keep the essential information.

    return flashcards


def process_youtube_videos(video_ids):
    """Process YouTube videos by extracting the transcript, generating flashcards, and uploading them to Anki.

    Args:
        video_ids (list): The list of YouTube video IDs to process.
    """
    logging.info(f"Processing {len(video_ids)} YouTube videos...")
    for video_id in video_ids:
        logging.info("----------------------------------")
        logging.info(f"Processing video with ID: {video_id}")
        # Get the title and channel name of the YouTube video
        video_title, channel_name = get_youtube_video_details(
            video_id, YOUTUBE_DATA_API_KEY
        )
        logging.info(f"Channel name: {channel_name}")
        logging.info(f"Video title: {video_title}")
        # Get the transcript for the YouTube video
        transcript = extract_transcript_from_youtube(video_id)

        # Generate flashcards from the transcript
        flashcards = generate_flashcards(transcript)
        print(flashcards)
        logging.info(f"Created {len(flashcards)} flashcards for the video.")

        # Generate tags for the flashcards
        tags = generate_tags(flashcards)
        logging.info(f"Generated {len(tags)} tags: {tags}.")

        # Upload the flashcards to Anki
        for card in flashcards:
            # Add channel name and video title as a header to the card
            front = (
                f"<h1>{channel_name}</h1><h2>{video_title}</h2><br>{card['question']}"
            )
            add_anki_card(
                f"YouTube::{channel_name}",
                "Basic",
                front,
                card["answer"],
                tags=tags,
            )
        logging.info("Uploaded flashcards to Anki.")


def process_books(books):
    """
    Process a book files by converting them to text, generating flashcards, and uploading them to Anki.

    Args:
        books (list): The list of paths to the book files.
    """
    logging.info(f"Processing {len(books)} new books...")
    for ebook_path in books:
        try:
            # Convert ebook to text
            book_content = process_book_attachment(ebook_path)

            # Generate flashcards from the text
            flashcards = generate_flashcards(book_content)
            logging.info(f"Created {len(flashcards)} flashcards from the book.")

            # Generate tags for the flashcards
            tags = generate_tags(flashcards)
            logging.info(f"Generated tags: {tags}.")

            # Use the book filename to infer the author and title
            author_name = openai_call(
                f"Return just the author name of the book inferred from this filename: {ebook_path}. The answer should ONLY contain the name of the author and nothing else.",
                model="gpt-4o-mini",
            )
            logging.info(f"Author name: {author_name}")
            book_title = openai_call(
                f"Return just the title of the book inferred from this filename: {ebook_path}. The answer should ONLY contain the name of the book and nothing else.",
                model="gpt-4o-mini",
            )
            logging.info(f"Book title: {book_title}")

            # Upload the flashcards to Anki
            for card in flashcards:
                front = (
                    f"<h1>{author_name}</h1><h2>{book_title}</h2><br>{card['question']}"
                )
                add_anki_card(
                    f"Books::{author_name}",
                    "Basic",
                    front,
                    card["answer"],
                    tags=tags,
                )
            logging.info(
                f"Uploaded flashcards for book '{author_name} - {book_title}' to Anki."
            )

        except Exception as e:
            logging.error(f"Error processing book '{ebook_path}': {e}")


def process_documents(documents):
    """
    Process a document files by converting them to text, generating flashcards, and uploading them to Anki.

    Args:
        documents (list): The list of paths to the document files.
    """
    logging.info(f"Processing {len(documents)} new documents...")
    for document_path in documents:
        try:
            # Convert document to text
            text_file_path = document_path.rsplit(".", 1)[0] + ".txt"
            subprocess.run(["pdftotext", document_path, text_file_path], check=True)
            with open(text_file_path, "r") as file:
                document_content = file.read()

            # Generate flashcards from the text
            flashcards = generate_flashcards(document_content)
            logging.info(f"Created {len(flashcards)} flashcards from the document.")

            # Generate tags for the flashcards
            tags = generate_tags(flashcards)
            logging.info(f"Generated tags: {tags}.")

            # Use the document filename to infer the topic
            topic = openai_call(
                f"Return the topic or title of this document inferred from this filename: {document_path}."
                f"\nAnd these flashcards:\n{flashcards}."
                "\n\nThe answer should ONLY contain the topic or title and nothing else.",
                model="gpt-4o-mini",
            )
            logging.info(f"Topic: {topic}")

            # Upload the flashcards to Anki
            for card in flashcards:
                front = f"<h1>{topic}</h1><br>{card['question']}"
                add_anki_card(
                    f"Documents::{topic}",
                    "Basic",
                    front,
                    card["answer"],
                    tags=tags,
                )
            logging.info(f"Uploaded flashcards for document '{topic}' to Anki.")

        except Exception as e:
            logging.error(f"Error processing document '{document_path}': {e}")


def main():
    """Main function to execute the application logic."""
    # Check email for unread YouTube links or book attachments
    content = check_email()
    video_ids = content.get("youtube", [])
    books = content.get("books", [])
    documents = content.get("documents", [])

    if not any(content.values()):
        logging.info("No new YouTube links or book attachments found.")
        return

    # Open Anki
    with open("anki_output.log", "w") as f:
        anki_process = subprocess.Popen(["anki"], stdout=f, stderr=f)

    # Wait for Anki to be ready
    start_time = time.time()
    while time.time() - start_time < 30:  # Maximum wait time of 30 seconds
        try:
            response = requests.post(
                ANKI_CONNECT_URL, json={"action": "version", "version": 6}, timeout=5
            )
            if response.status_code == 200:
                break
        except requests.RequestException:
            pass
        time.sleep(1)  # Check every 1 second

    if time.time() - start_time >= 30:
        logging.error("Anki is not responding.")
        anki_process.kill()
        return

    logging.info("Anki is ready.")
    # Sync Anki
    send_anki_request("sync")

    # Process YouTube videos if found
    if video_ids:
        process_youtube_videos(video_ids)

    # Process books if found
    if books:
        process_books(books)

    # Process documents if found
    if documents:
        process_documents(documents)

    # Sync the media files with Anki
    send_anki_request("sync")
    logging.info("Sync completed!")
    # Close Anki
    anki_process.kill()


if __name__ == "__main__":
    main()