main.py

import requests
from requests.auth import HTTPBasicAuth
import os
from db import db, Post, Comment
from vars import (
    REDDIT_CLIENT_ID,
    REDDIT_CLIENT_SECRET,
    REDDIT_USERNAME,
    REDDIT_PASSWORD,
    SUBREDDITS_FILE,
    SCRAPE_LIMIT,
    OPENROUTER_API_KEY,
    OPENROUTER_MODEL,
    HUGGINGFACE_TOKEN,
    HUGGINGFACE_MODEL,
)
import datetime
import numpy as np
import json
from utils import HuggingFaceClient


class Scraper:
    def __init__(
        self,
        reddit_client_id: str,
        reddit_client_secret: str,
        reddit_username: str,
        reddit_password: str,
    ):
        auth = HTTPBasicAuth(reddit_client_id, reddit_client_secret)
        data = {
            "grant_type": "password",
            "username": reddit_username,
            "password": reddit_password,
        }
        headers = {"User-Agent": "MyRedditApp/0.1 by " + reddit_username}

        response = requests.post(
            "https://www.reddit.com/api/v1/access_token",
            auth=auth,
            data=data,
            headers=headers,
        )

        if response.status_code == 200:
            self.access_token: str = response.json().get("access_token")

    def scrape_subreddit(self, subreddit: str, limit: int = 10):
        headers = {
            "Authorization": f"bearer {self.access_token}",
            "User-Agent": "Nore/0.1 by " + os.environ["REDDIT_USERNAME"],
        }

        # Reddit API has a maximum limit of 100 posts per request
        # To get more than 100 posts, we need to use pagination with the 'after' token
        all_posts = []
        after_token = None
        remaining_limit = limit
        page_count = 0

        while remaining_limit > 0:
            page_count += 1
            # Request up to 100 posts at a time (Reddit's max)
            current_limit = min(remaining_limit, 100)
            url = f"https://oauth.reddit.com/r/{subreddit}/new?limit={current_limit}"

            # Add after token if we have one (for pagination)
            if after_token:
                url += f"&after={after_token}"

            response = requests.get(url, headers=headers)

            if response.status_code != 200:
                print(f"Failed to fetch posts: {response.status_code}")
                break

            response_data = response.json()
            posts = response_data.get("data", {}).get("children", [])

            if not posts:
                break

            all_posts.extend(posts)

            # Get the after token for the next page
            after_token = response_data.get("data", {}).get("after")

            # If there's no after token, we've reached the end
            if not after_token:
                break

            remaining_limit -= len(posts)

        return all_posts

    def scrape_comments(self, post_id: str):
        headers = {
            "Authorization": f"bearer {self.access_token}",
            "User-Agent": "Nore/0.1 by " + os.environ["REDDIT_USERNAME"],
        }
        url = f"https://oauth.reddit.com/comments/{post_id}"
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            comments = response.json()[1].get("data", {}).get("children", [])
            return comments
        else:
            print(
                "Failed to fetch comments:",
                response.status_code,
                response.text,
            )


def process_comment_node(node, post_obj, parent_comment=None):
    """Recursively process a Reddit comment node and store it in the DB.

    node: a dict with keys 'kind' and 'data' matching Reddit API response
    post_obj: Peewee Post instance
    parent_comment: Peewee Comment instance or None
    """
    kind = node.get("kind")
    data = node.get("data", {})

    if kind == "t1":
        # It's a comment
        comment_id = data.get("id")
        if not comment_id:
            return

        content = data.get("body", "")
        created_at = datetime.datetime.fromtimestamp(data.get("created_utc", 0))
        embedding = hf_client.generate_embedding(content)
        if not embedding or len(embedding) == 0 or type(embedding) == dict:
            return

        # Use get_or_create to avoid duplicate unique constraint errors
        comment_obj, created = Comment.get_or_create(
            comment_id=comment_id,
            defaults={
                "post": post_obj,
                "parent": parent_comment,
                "author": data.get("author", "[deleted]"),
                "created_at": created_at,
                "content": content,
                "score": data.get("score", 0),
                "embedding": embedding,
            },
        )

        # If the comment existed but parent or content might be missing, ensure parent is set
        if not created and parent_comment is not None and comment_obj.parent is None:
            comment_obj.parent = parent_comment
            comment_obj.save()

        # Recurse into replies if present
        replies = data.get("replies")
        if isinstance(replies, dict):
            children = replies.get("data", {}).get("children", [])
            for child in children:
                process_comment_node(child, post_obj, parent_comment=comment_obj)

    elif kind == "more":
        # 'more' contains additional comment IDs that would need extra API requests.
        # Skipping for now to keep this implementation simple.
        return
    else:
        # t3 (post), t2 (user), t4 (message) etc. - ignore
        return


with open(SUBREDDITS_FILE, "r") as f:
    SUBREDDITS = f.readlines()

scraper = Scraper(
    REDDIT_CLIENT_ID,
    REDDIT_CLIENT_SECRET,
    REDDIT_USERNAME,
    REDDIT_PASSWORD,
)

hf_client = HuggingFaceClient(
    hf_token=HUGGINGFACE_TOKEN,
    model=HUGGINGFACE_MODEL,
)


def initialize_database():
    db.connect()
    db.create_tables([Post, Comment])


def main():
    initialize_database()

    if not scraper.access_token:
        print("Failed to authenticate with Reddit API.")
        return

    for subreddit in SUBREDDITS:
        subreddit = subreddit.strip()
        print(f"Scraping subreddit: {subreddit}")
        posts = scraper.scrape_subreddit(subreddit, limit=SCRAPE_LIMIT)
        if not posts:
            print(f"No posts found or failed to fetch posts for r/{subreddit}")
            continue

        print(f"Found {len(posts)} posts in r/{subreddit}")
        for post in posts:
            post_data = post["data"]
            # Skip posts already in DB
            post_obj = Post.get_or_none(post_id=post_data["id"])
            if post_obj:
                continue

            if post_data.get("selftext"):
                embedding = hf_client.generate_embedding(post_data.get("selftext", ""))
            else:
                embedding = hf_client.generate_embedding(post_data.get("title", ""))
            if not embedding or len(embedding) == 0 or type(embedding) == dict:
                continue
            post_obj, _ = Post.get_or_create(
                post_id=post_data["id"],
                defaults={
                    "title": post_data.get("title", ""),
                    "author": post_data.get("author", ""),
                    "created_at": datetime.datetime.fromtimestamp(
                        post_data.get("created_utc", 0)
                    ),
                    "content": post_data.get("selftext", ""),
                    "score": post_data.get("score", 0),
                    "embedding": embedding,
                },
            )

            print(f"Scraping comments for post: {post_data['id']}")
            comments = scraper.scrape_comments(post_data["id"])
            if not comments:
                print(
                    f"No comments found or failed to fetch comments for post {post_data['id']}"
                )
                continue

            print(
                f"Found {len(comments)} top-level comments for post {post_data['id']}"
            )
            for comment in comments:
                # Recursively process each top-level comment node
                process_comment_node(comment, post_obj)

        """post_id = "1ojlkru"
        post_obj = Post.get_or_none(post_id=post_id)
        comments = scraper.scrape_comments(post_id)
        if not comments:
            print(f"No comments found or failed to fetch comments for post {post_id}")
            continue
        print(f"Found {len(comments)} top-level comments for post {post_id}")
        for comment in comments:
            # Recursively process each top-level comment node
            process_comment_node(comment, post_obj)"""

    db.close()


if __name__ == "__main__":
    main()