scripts/append_db_functions.py

"""
Helper script: appends new DB functions to modules/db.py for actual-trade tracking.
Run from project root: python scripts/append_db_functions.py
"""
import os

DB_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "modules", "db.py")

APPEND = r'''

def snapshot_holdings(owner_id, snapshot_date: str):
    """
    Saves current portfolio_holdings as a holdings_snapshot before they are overwritten.
    Idempotent — only snapshots once per date.
    """
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                "SELECT COUNT(*) FROM holdings_snapshots WHERE user_id = %s AND snapshot_date = %s;",
                (owner_id, snapshot_date)
            )
            count = cur.fetchone()[0]
            if count > 0:
                print(f"[db] Snapshot for {snapshot_date} already exists, skipping.")
                return
            cur.execute(
                "SELECT symbol, qty, avg_cost, current_price FROM portfolio_holdings WHERE user_id = %s;",
                (owner_id,)
            )
            rows = cur.fetchall()
            if not rows:
                print("[db] No holdings to snapshot.")
                return
            for r in rows:
                cur.execute(
                    """
                    INSERT INTO holdings_snapshots (user_id, snapshot_date, symbol, qty, avg_cost, ltp)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (user_id, snapshot_date, symbol) DO NOTHING;
                    """,
                    (owner_id, snapshot_date, r[0], r[1], r[2], r[3])
                )
            conn.commit()
            print(f"[db] Saved holdings snapshot for {snapshot_date} ({len(rows)} symbols).")
    except Exception as e:
        conn.rollback()
        print(f"[db] snapshot_holdings failed: {e}")
    finally:
        conn.close()


def detect_and_save_actual_trades(owner_id, new_holdings: list, trade_date: str):
    """
    Compares new_holdings against the most recent holdings snapshot.
    Symbols where new_qty > snapshot_qty = new purchase. Saves deltas to actual_trades.
    new_holdings: list of dicts with keys: symbol, qty, avg_cost, ltp
    """
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT DISTINCT ON (symbol) symbol, qty
                FROM holdings_snapshots
                WHERE user_id = %s
                ORDER BY symbol, snapshot_date DESC;
                """,
                (owner_id,)
            )
            snap_rows = cur.fetchall()
            if not snap_rows:
                print("[db] No snapshot found - cannot detect trades. Snapshot saved next upload.")
                return 0
            snap_map = {r[0]: float(r[1]) for r in snap_rows}
            new_map = {
                h["symbol"]: (float(h["qty"]), float(h.get("avg_cost", 0)), float(h.get("ltp") or 0))
                for h in new_holdings
            }
            nifty_price = None
            try:
                from modules.price_fetcher import fetch_nifty50_current
                nifty_price = fetch_nifty50_current()
            except Exception as nifty_err:
                print(f"[db] Could not fetch Nifty: {nifty_err}")
            trades_inserted = 0
            for symbol, (new_qty, avg_cost, ltp) in new_map.items():
                old_qty = snap_map.get(symbol, 0.0)
                delta_qty = new_qty - old_qty
                if delta_qty > 0.001 and avg_cost > 0:
                    cur.execute(
                        """
                        INSERT INTO actual_trades (user_id, trade_date, symbol, qty_bought, buy_price, nifty_on_trade_day)
                        VALUES (%s, %s, %s, %s, %s, %s)
                        ON CONFLICT (user_id, trade_date, symbol) DO UPDATE
                        SET qty_bought = EXCLUDED.qty_bought,
                            buy_price = EXCLUDED.buy_price,
                            nifty_on_trade_day = EXCLUDED.nifty_on_trade_day;
                        """,
                        (owner_id, trade_date, symbol, delta_qty, avg_cost, nifty_price)
                    )
                    trades_inserted += 1
                    print(f"[db] Actual trade: {symbol} +{delta_qty:.2f} @ Rs {avg_cost:.2f}")
            conn.commit()
            print(f"[db] Saved {trades_inserted} actual trade(s) for {trade_date}.")
            return trades_inserted
    except Exception as e:
        conn.rollback()
        print(f"[db] detect_and_save_actual_trades failed: {e}")
        return 0
    finally:
        conn.close()


def load_actual_trades_from_db(owner_id) -> list:
    """Returns all actual trades for the owner, newest first."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT trade_date::text, symbol, qty_bought, buy_price, nifty_on_trade_day
                FROM actual_trades
                WHERE user_id = %s
                ORDER BY trade_date DESC, symbol ASC;
                """,
                (owner_id,)
            )
            rows = cur.fetchall()
            return [
                {
                    "trade_date": r[0],
                    "symbol": r[1],
                    "qty_bought": float(r[2]),
                    "buy_price": float(r[3]),
                    "nifty_on_trade_day": float(r[4]) if r[4] is not None else None,
                }
                for r in rows
            ]
    except Exception as e:
        print(f"[db] load_actual_trades_from_db failed: {e}")
        return []
    finally:
        conn.close()


def get_ai_suggestions_for_date(owner_id, trade_date: str) -> list:
    """Returns AI BUY signals from the briefing on trade_date for comparison table."""
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT s.symbol, s.signal, s.price_on_day
                FROM stock_signals s
                JOIN daily_briefings db ON s.briefing_id = db.id
                WHERE db.user_id = %s AND db.date = %s
                ORDER BY s.signal, s.symbol;
                """,
                (owner_id, trade_date)
            )
            rows = cur.fetchall()
            return [
                {"symbol": r[0], "signal": r[1], "price_on_day": float(r[2])}
                for r in rows
            ]
    except Exception as e:
        print(f"[db] get_ai_suggestions_for_date failed: {e}")
        return []
    finally:
        conn.close()
'''

with open(DB_PATH, "r", encoding="utf-8") as f:
    existing = f.read()

# Only append if functions not already present
if "def snapshot_holdings" not in existing:
    with open(DB_PATH, "a", encoding="utf-8") as f:
        f.write(APPEND)
    print(f"SUCCESS: Appended new functions to {DB_PATH}")
else:
    print("SKIPPED: Functions already present in db.py")