worker/worker.py

import sys
import os
import json
import asyncio
import datetime
from typing import Dict, List, Optional
from pathlib import Path
import yfinance as yf    # type: ignore
from aiokafka import AIOKafkaProducer     # type: ignore

# Add the project root directory to the Python path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from backend.app.config.settings import settings
from backend.app.database.session import SessionLocal
from backend.app.database import crud
from backend.app.schemas import schemas

KAFKA_BOOTSTRAP_SERVERS = settings.KAFKA_BOOTSTRAP_SERVERS

# Tickers are dynamically fetched from DB each cycle — no hardcoded list

KAFKA_TOPIC = "stock-ticks"

# In-memory buffer to aggregate ticks into 1-minute candles
aggregation_buffer: Dict[str, Dict[str, List]] = {}

# Global rate-limiting cooldown trackers
COOLDOWN_UNTIL = None
COOLDOWN_BACKOFF_SEC = 60

async def get_latest_stock_data(ticker: str) -> Optional[dict]:
    """
    Fetch current stock price and cumulative volume using yfinance fast_info.
    Respects global rate-limiting cooldown if Yahoo Finance blocks the IP.
    """
    global COOLDOWN_UNTIL, COOLDOWN_BACKOFF_SEC
    
    now_utc = datetime.datetime.now(datetime.timezone.utc)
    if COOLDOWN_UNTIL and now_utc < COOLDOWN_UNTIL:
        return None

    try:
        # Run synchronous yfinance operations in a thread pool to avoid blocking the event loop
        loop = asyncio.get_event_loop()
        yf_ticker = yf.Ticker(ticker)
        
        # fast_info retrieves basic metrics quickly
        info = await loop.run_in_executor(None, lambda: yf_ticker.fast_info)
        
        price = info.last_price
        volume = info.last_volume
        
        if price is None or volume is None:
            return None
            
        # Reset backoff on success
        COOLDOWN_BACKOFF_SEC = 60
        COOLDOWN_UNTIL = None
        
        return {
            "price": float(price),
            "volume": int(volume),
            "timestamp": now_utc
        }
    except Exception as e:
        err_str = str(e).lower()
        if "too many requests" in err_str or "rate limit" in err_str or "429" in err_str:
            COOLDOWN_UNTIL = now_utc + datetime.timedelta(seconds=COOLDOWN_BACKOFF_SEC)
            print(f"Worker rate limited. Entering yfinance cooldown for {COOLDOWN_BACKOFF_SEC}s. Error: {e}")
            COOLDOWN_BACKOFF_SEC = min(COOLDOWN_BACKOFF_SEC * 2, 900)  # Double backoff up to 15 mins
        else:
            print(f"Error fetching data for {ticker}: {str(e)}")
        return None


async def aggregate_and_save_candle(db, ticker: str, timestamp: datetime.datetime):
    """
    Takes buffered ticks for a ticker, calculates OHLCV, and saves it to NeonDB.
    """
    data = aggregation_buffer[ticker]
    prices = data["prices"]
    volumes = data["volumes"]
    
    if not prices:
        return

    # Calculate candlestick values
    open_price = prices[0]
    high_price = max(prices)
    low_price = min(prices)
    close_price = prices[-1]
    
    # Volume for the minute is the difference between last and first cumulative volumes
    volume_diff = 0
    if len(volumes) > 1:
        volume_diff = max(0, volumes[-1] - volumes[0])
    else:
        volume_diff = volumes[0] if volumes else 0

    # Skip flat, inactive candles when the market is closed (no price change and zero volume)
    # This prevents the DB from filling with flat stale lines during closed hours.
    # Crypto tickers are exempted as they trade 24/7.
    is_crypto = ticker.endswith("-USD") or ticker.endswith("-BTC")
    if not is_crypto and open_price == close_price == high_price == low_price and volume_diff == 0:
        return

    candle = schemas.StockHistoryBase(
        ticker=ticker,
        timestamp=timestamp.replace(second=0, microsecond=0),  # Round to the minute
        open=open_price,
        high=high_price,
        low=low_price,
        close=close_price,
        volume=volume_diff
    )


    try:
        await crud.insert_stock_candle(db, candle)
        print(f"Aggregated & Saved candle for {ticker} at {candle.timestamp}: O={open_price} C={close_price}")
    except Exception as e:
        print(f"Failed to save candle for {ticker}: {str(e)}")
        
    data["prices"].clear()
    data["volumes"].clear()

async def main():
    print(f"Starting Ingestion Worker...")
    print(f"Connecting to Redpanda at {KAFKA_BOOTSTRAP_SERVERS}...")
    
    # 1. Initialize Redpanda/Kafka Producer
    producer = AIOKafkaProducer(
        bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS,
        value_serializer=lambda v: json.dumps(v).encode("utf-8")
    )
    
    await producer.start()
    print("Redpanda Producer connected.")

    last_minute = datetime.datetime.now(datetime.timezone.utc).minute

    try:
        while True:
            start_time = asyncio.get_event_loop().time()
            now = datetime.datetime.now(datetime.timezone.utc)
            
            # Check if a new minute has started
            minute_changed = now.minute != last_minute
            
            try:
                # Create a database session for candle aggregation
                async with SessionLocal() as db:
                    from sqlalchemy import select
                    from backend.app.database.models import Watchlist 
                    result= await db.execute(select(Watchlist.ticker).distinct())
                    dynamic_tickers= [row[0] for row in result.fetchall()]
                    
                    for t in dynamic_tickers:
                        if t not in aggregation_buffer:
                            aggregation_buffer[t]= {"prices": [], "volumes": []}
                    
                    for ticker in dynamic_tickers:
                        tick = await get_latest_stock_data(ticker)
                        
                        if tick:
                            aggregation_buffer[ticker]["prices"].append(tick["price"])
                            aggregation_buffer[ticker]["volumes"].append(tick["volume"])
                            
                            payload = {
                                "ticker": ticker,
                                "price": tick["price"],
                                "volume": tick["volume"],
                                "timestamp": tick["timestamp"].isoformat()
                            }
                            
                            await producer.send(KAFKA_TOPIC, payload)
                            print(f"Tick published: {ticker} = {tick['price']}")
                        
                        if minute_changed:
                            await aggregate_and_save_candle(db, ticker, now)
                            
                        # Stagger calls to yfinance to prevent rate limiting
                        await asyncio.sleep(1.5)
                
                if minute_changed:
                    last_minute = now.minute
            except Exception as loop_err:
                print(f"Error in poll cycle (database/network drop): {str(loop_err)}")

            # Sleep for 30 seconds (must run outside the loop and minute check)
            elapsed_time = asyncio.get_event_loop().time() - start_time
            sleep_time = max(1.0, 30.0 - elapsed_time)
            await asyncio.sleep(sleep_time)

            
    except asyncio.CancelledError:
        print("Worker stopped.")
    finally:
        await producer.stop()
        print("Redpanda connection closed.")

if __name__ == "__main__":
    asyncio.run(main())