reddit_stock_analyzer/rstat_tool/main.py

# rstat_tool/main.py

import argparse
import json
import os
import time
import sys
import subprocess
from dotenv import load_dotenv
from pathlib import Path

import praw

from . import database
from .ticker_extractor import extract_tickers
from .sentiment_analyzer import get_sentiment_score
from .logger_setup import setup_logging, logger as log

def load_subreddits(filepath):
    """Loads a list of subreddits from a JSON file."""
    try:
        with open(filepath, 'r') as f:
            return json.load(f).get("subreddits", [])
    except (FileNotFoundError, json.JSONDecodeError) as e:
        log.error(f"Error loading config file '{filepath}': {e}")
        return None

def get_reddit_instance():
    """Initializes and returns a PRAW Reddit instance."""
    env_path = Path(__file__).parent.parent / '.env'
    load_dotenv(dotenv_path=env_path)
    
    client_id = os.getenv("REDDIT_CLIENT_ID")
    client_secret = os.getenv("REDDIT_CLIENT_SECRET")
    user_agent = os.getenv("REDDIT_USER_AGENT")
    if not all([client_id, client_secret, user_agent]):
        log.error("Error: Reddit API credentials not found in .env file.")
        return None
    return praw.Reddit(client_id=client_id, client_secret=client_secret, user_agent=user_agent)

def get_financial_data_via_fetcher(ticker_symbol):
    """
    Calls two separate, isolated fetcher scripts to get market cap and closing price,
    bypassing the internal library conflict.
    """
    financials = {"market_cap": None, "closing_price": None}
    project_root = Path(__file__).parent.parent

    # --- Call 1: Get Market Cap ---
    try:
        mc_script_path = project_root / 'fetch_market_cap.py'
        command_mc = [sys.executable, str(mc_script_path), ticker_symbol]
        result_mc = subprocess.run(command_mc, capture_output=True, text=True, check=True, timeout=30)
        financials.update(json.loads(result_mc.stdout))
    except Exception as e:
        log.warning(f"Market cap fetcher failed for {ticker_symbol}: {e}")

    # --- Call 2: Get Closing Price ---
    try:
        cp_script_path = project_root / 'fetch_close_price.py'
        command_cp = [sys.executable, str(cp_script_path), ticker_symbol]
        result_cp = subprocess.run(command_cp, capture_output=True, text=True, check=True, timeout=30)
        financials.update(json.loads(result_cp.stdout))
    except Exception as e:
        log.warning(f"Closing price fetcher failed for {ticker_symbol}: {e}")
        
    return financials

# --- HELPER FUNCTION: Contains all the optimized logic for one post ---
def _process_submission(submission, subreddit_id, conn, comment_limit, fetch_financials):
    """
    Processes a single Reddit submission with optimized logic.
    - Uses a single loop over comments.
    - Caches ticker IDs to reduce DB lookups.
    """
    current_time = time.time()
    
    # 1. Initialize data collectors for this post
    tickers_in_title = set(extract_tickers(submission.title))
    all_tickers_found_in_post = set(tickers_in_title)
    all_comment_sentiments = []
    ticker_id_cache = {} # In-memory cache for ticker IDs for this post

    submission.comments.replace_more(limit=0)
    all_comments = submission.comments.list()[:comment_limit]

    # 2. --- SINGLE LOOP OVER COMMENTS ---
    # We gather all necessary information in one pass.
    for comment in all_comments:
        comment_sentiment = get_sentiment_score(comment.body)
        all_comment_sentiments.append(comment_sentiment) # For the deep dive
        
        tickers_in_comment = set(extract_tickers(comment.body))
        if not tickers_in_comment:
            continue

        all_tickers_found_in_post.update(tickers_in_comment)

        # Apply the hybrid logic
        if tickers_in_title:
            # If the title has tickers, every comment is a mention for them
            for ticker_symbol in tickers_in_title:
                if ticker_symbol not in ticker_id_cache:
                    ticker_id_cache[ticker_symbol] = database.get_or_create_entity(conn, 'tickers', 'symbol', ticker_symbol)
                ticker_id = ticker_id_cache[ticker_symbol]
                database.add_mention(conn, ticker_id, subreddit_id, submission.id, 'comment', int(comment.created_utc), comment_sentiment)
        else:
            # If no title tickers, only direct mentions in comments count
            for ticker_symbol in tickers_in_comment:
                if ticker_symbol not in ticker_id_cache:
                    ticker_id_cache[ticker_symbol] = database.get_or_create_entity(conn, 'tickers', 'symbol', ticker_symbol)
                ticker_id = ticker_id_cache[ticker_symbol]
                database.add_mention(conn, ticker_id, subreddit_id, submission.id, 'comment', int(comment.created_utc), comment_sentiment)

    # 3. Process title mentions (if any)
    if tickers_in_title:
        log.info(f"  -> Title Mention(s): {', '.join(tickers_in_title)}. Attributing all comments.")
        post_sentiment = get_sentiment_score(submission.title)
        for ticker_symbol in tickers_in_title:
            if ticker_symbol not in ticker_id_cache:
                ticker_id_cache[ticker_symbol] = database.get_or_create_entity(conn, 'tickers', 'symbol', ticker_symbol)
            ticker_id = ticker_id_cache[ticker_symbol]
            database.add_mention(conn, ticker_id, subreddit_id, submission.id, 'post', int(submission.created_utc), post_sentiment)

    # 4. Fetch financial data if enabled
    if fetch_financials:
        for ticker_symbol in all_tickers_found_in_post:
            ticker_id = ticker_id_cache[ticker_symbol] # Guaranteed to be in cache
            ticker_info = database.get_ticker_info(conn, ticker_id)
            if not ticker_info['last_updated'] or (current_time - ticker_info['last_updated'] > database.MARKET_CAP_REFRESH_INTERVAL):
                log.info(f"    -> Fetching financial data for {ticker_symbol}...")
                financials = get_financial_data_via_fetcher(ticker_symbol)
                database.update_ticker_financials(conn, ticker_id, financials.get('market_cap'), financials.get('closing_price'))
    
    # 5. Save deep dive analysis
    avg_sentiment = sum(all_comment_sentiments) / len(all_comment_sentiments) if all_comment_sentiments else 0
    post_analysis_data = {
        "post_id": submission.id, "title": submission.title,
        "post_url": f"https://reddit.com{submission.permalink}", "subreddit_id": subreddit_id,
        "post_timestamp": int(submission.created_utc), "comment_count": len(all_comments),
        "avg_comment_sentiment": avg_sentiment
    }
    database.add_or_update_post_analysis(conn, post_analysis_data)

def scan_subreddits(reddit, subreddits_list, post_limit=100, comment_limit=100, days_to_scan=1, fetch_financials=True):
    conn = database.get_db_connection()
    post_age_limit = days_to_scan * 86400
    current_time = time.time()
    
    log.info(f"Scanning {len(subreddits_list)} subreddit(s) for NEW posts in the last {days_to_scan} day(s)...")
    if not fetch_financials:
        log.warning("NOTE: Financial data fetching is disabled for this run.")

    for subreddit_name in subreddits_list:
        try:
            normalized_sub_name = subreddit_name.lower()
            subreddit_id = database.get_or_create_entity(conn, 'subreddits', 'name', normalized_sub_name)
            subreddit = reddit.subreddit(normalized_sub_name)
            log.info(f"Scanning r/{normalized_sub_name}...")
            
            for submission in subreddit.new(limit=post_limit):
                if (current_time - submission.created_utc) > post_age_limit:
                    log.info(f"  -> Reached posts older than the {days_to_scan}-day limit.")
                    break
                
                # Call the new helper function for each post
                _process_submission(submission, subreddit_id, conn, comment_limit, fetch_financials)
                
        except Exception as e:
            log.error(f"Could not scan r/{normalized_sub_name}. Error: {e}", exc_info=True)
            
    conn.close()
    log.critical("\n--- Scan Complete ---")

def main():
    """Main function to run the Reddit stock analysis tool."""
    parser = argparse.ArgumentParser(description="Analyze stock ticker mentions on Reddit.", formatter_class=argparse.RawTextHelpFormatter)
    
    parser.add_argument("-f", "--config", default="subreddits.json", help="Path to the JSON file for scanning. (Default: subreddits.json)")
    parser.add_argument("-s", "--subreddit", help="Scan a single subreddit, ignoring the config file.")
    parser.add_argument("-d", "--days", type=int, default=1, help="Number of past days to scan for new posts. (Default: 1)")
    parser.add_argument("-p", "--posts", type=int, default=200, help="Max posts to check per subreddit. (Default: 200)")
    parser.add_argument("-c", "--comments", type=int, default=100, help="Number of comments to scan per post. (Default: 100)")
    parser.add_argument("-n", "--no-financials", action="store_true", help="Disable fetching of financial data during the Reddit scan.")
    parser.add_argument("--update-top-tickers", action="store_true", help="Update financial data only for tickers currently in the Top 10 daily/weekly dashboards.")
    parser.add_argument(
        "-u", "--update-financials-only",
        nargs='?',
        const="ALL_TICKERS", # A special value to signify "update all"
        default=None,
        metavar='TICKER',
        help="Update financials. Provide a ticker symbol to update just one,\nor use the flag alone to update all tickers in the database."
    )
    parser.add_argument("--debug", action="store_true", help="Enable detailed debug logging to the console.")
    parser.add_argument("--stdout", action="store_true", help="Print all log messages to the console.")
    
    args = parser.parse_args()
    setup_logging(console_verbose=args.stdout, debug_mode=args.debug)

    database.initialize_db()
    
    if args.update_top_tickers:
        # --- Mode 1: Update Top Tickers ---
        log.critical("--- Starting Financial Data Update for Top Tickers ---")
        top_daily = database.get_top_daily_ticker_symbols()
        top_weekly = database.get_top_weekly_ticker_symbols()
        unique_top_tickers = sorted(list(set(top_daily + top_weekly)))
        
        if not unique_top_tickers:
            log.info("No top tickers found in the last week. Nothing to update.")
        else:
            log.info(f"Found {len(unique_top_tickers)} unique tickers to update: {', '.join(unique_top_tickers)}")
            conn = database.get_db_connection()
            for ticker_symbol in unique_top_tickers:
                ticker_info = database.get_ticker_by_symbol(ticker_symbol)
                if ticker_info:
                    log.info(f"  -> Updating financials for {ticker_info['symbol']}...")
                    financials = get_financial_data_via_fetcher(ticker_info['symbol'])
                    database.update_ticker_financials(
                        conn, ticker_info['id'],
                        financials.get('market_cap'),
                        financials.get('closing_price')
                    )
            conn.close()
        log.critical("--- Top Ticker Financial Data Update Complete ---")

    elif args.update_financials_only:
        # --- Mode 2: Update All or a Single Ticker ---
        update_mode = args.update_financials_only
        if update_mode == "ALL_TICKERS":
            log.critical("--- Starting Financial Data Update for ALL tickers ---")
            all_tickers = database.get_all_tickers()
            log.info(f"Found {len(all_tickers)} tickers in the database to update.")
            conn = database.get_db_connection()
            for ticker in all_tickers:
                symbol = ticker['symbol']
                log.info(f"  -> Updating financials for {symbol}...")
                financials = get_financial_data_via_fetcher(symbol)
                database.update_ticker_financials(
                    conn, ticker['id'],
                    financials.get('market_cap'),
                    financials.get('closing_price')
                )
            conn.close()
        else:
            ticker_symbol_to_update = update_mode
            log.critical(f"--- Starting Financial Data Update for single ticker: {ticker_symbol_to_update} ---")
            ticker_info = database.get_ticker_by_symbol(ticker_symbol_to_update)
            if ticker_info:
                conn = database.get_db_connection()
                log.info(f"  -> Updating financials for {ticker_info['symbol']}...")
                financials = get_financial_data_via_fetcher(ticker_info['symbol'])
                database.update_ticker_financials(
                    conn, ticker_info['id'],
                    financials.get('market_cap'),
                    financials.get('closing_price')
                )
                conn.close()
            else:
                log.error(f"Ticker '{ticker_symbol_to_update}' not found in the database.")
        log.critical("--- Financial Data Update Complete ---")

    else:
        # --- Mode 3: Default Reddit Scan ---
        log.critical("--- Starting Reddit Scan Mode ---")
        if args.subreddit:
            subreddits_to_scan = [args.subreddit]
            log.info(f"Targeted Scan Mode: Focusing on r/{args.subreddit}")
        else:
            log.info(f"Config Scan Mode: Loading subreddits from {args.config}")
            subreddits_to_scan = load_subreddits(args.config) 

        if not subreddits_to_scan:
            log.error("Error: No subreddits to scan.")
            return

        reddit = get_reddit_instance()
        if not reddit: return

        scan_subreddits(
            reddit,
            subreddits_to_scan,
            post_limit=args.posts,
            comment_limit=args.comments,
            days_to_scan=args.days,
            fetch_financials=(not args.no_financials)
        )

if __name__ == "__main__":
    main()