modules/db.py

"""
PortIQ Database Utility Module.
Handles database connections and CRUD operations for daily briefings,
portfolio holdings, and stock signals.
"""

import os
import re
import json
import psycopg2
from psycopg2.extras import RealDictCursor, Json
from dotenv import load_dotenv

load_dotenv()

# Global Connection Pool or connection factory
def get_db_connection():
    """
    Creates and returns a connection to the PostgreSQL database.
    Caller must close the connection.
    """
    db_url = os.getenv("DATABASE_URL")
    if db_url:
        return psycopg2.connect(db_url)
        
    host = os.getenv("DB_HOST", "localhost")
    port = os.getenv("DB_PORT", "5432")
    dbname = os.getenv("DB_NAME", "portiq_db")
    user = os.getenv("DB_USER", "postgres")
    password = os.getenv("DB_PASSWORD", "postgres")
    
    return psycopg2.connect(
        host=host,
        port=port,
        dbname=dbname,
        user=user,
        password=password
    )


def get_owner_user_id():
    """
    Returns the UUID of the owner user (owner@portiq.com).
    Creates it if it does not exist.
    """
    email = "owner@portiq.com"
    owner_hash = os.getenv("OWNER_PASSWORD_HASH")
    if not owner_hash:
        owner_hash = "084f7fa87d1dfcbb3965db0183b544b60a3cc180c59800a6e30018f70094770e" # fallback
        
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute("SELECT id FROM users WHERE email = %s;", (email,))
            row = cur.fetchone()
            if row:
                return row[0]
            
            # If not found, create it
            cur.execute(
                """
                INSERT INTO users (email, password_hash, role)
                VALUES (%s, %s, 'owner')
                RETURNING id;
                """,
                (email, owner_hash)
            )
            owner_id = cur.fetchone()[0]
            conn.commit()
            return owner_id
    finally:
        conn.close()


def clean_nans(val):
    """Recursively replaces float('nan')/NaN values with None (standard null)."""
    import math
    if isinstance(val, float) and math.isnan(val):
        return None
    elif isinstance(val, dict):
        return {k: clean_nans(v) for k, v in val.items()}
    elif isinstance(val, list):
        return [clean_nans(v) for v in val]
    return val


def save_briefing_to_db(date_str, data_dict, portfolio_df=None):
    """
    Saves a daily briefing JSON and its associated stock signals to the database.
    Overwrites if a briefing already exists for this date.
    """
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            # 1. Delete existing briefing for this date to support overwrites
            cur.execute(
                "DELETE FROM daily_briefings WHERE user_id = %s AND date = %s RETURNING id;",
                (owner_id, date_str)
            )
            cur.fetchone()  # clear result if any
            
            # 2. Extract fields from briefing JSON and clean NaN values
            mood = clean_nans(data_dict.get("mood", {}))
            news_summary = clean_nans(data_dict.get("news_summary", []))
            top_picks = clean_nans(data_dict.get("top_picks", []))
            avoid_today = clean_nans(data_dict.get("avoid_today", []))
            market_summary = clean_nans(data_dict.get("market_summary", []))
            news = clean_nans(data_dict.get("news", []))
            dividends = clean_nans(data_dict.get("dividends", []))
            
            # 3. Insert new daily briefing record
            portfolio = clean_nans(data_dict.get("portfolio", []))
            portfolio_snapshot = clean_nans(data_dict.get("_portfolio_snapshot", []))
            if not portfolio_snapshot and portfolio_df is not None and not portfolio_df.empty:
                portfolio_snapshot = clean_nans(portfolio_df.to_dict(orient="records"))

            cur.execute(
                """
                INSERT INTO daily_briefings (
                    user_id, date, mood, news_summary, portfolio, top_picks, 
                    avoid_today, market_summary, news, dividends, portfolio_snapshot
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                RETURNING id;
                """,
                (
                    owner_id,
                    date_str,
                    Json(mood),
                    Json(news_summary),
                    Json(portfolio),
                    Json(top_picks),
                    Json(avoid_today),
                    Json(market_summary),
                    Json(news),
                    Json(dividends),
                    Json(portfolio_snapshot)
                )
            )
            briefing_id = cur.fetchone()[0]
            
            # 4. Process and insert stock signals
            # Look at _signal_prices or derive from snapshot/portfolio
            signal_prices = data_dict.get("_signal_prices")
            if not signal_prices:
                signal_prices = _derive_signals(data_dict, portfolio_df)
                
            for symbol, sp in signal_prices.items():
                sig = sp.get("signal")
                price = sp.get("price_on_day", 0)
                qty = sp.get("qty", 1.0)
                nifty = sp.get("nifty_on_day")
                
                if sig and price > 0:
                    cur.execute(
                        """
                        INSERT INTO stock_signals (
                            briefing_id, symbol, signal, price_on_day, qty, nifty_on_day
                        ) VALUES (%s, %s, %s, %s, %s, %s);
                        """,
                        (briefing_id, symbol, sig, price, qty, nifty)
                    )
            
            # 5. Update portfolio_holdings in the DB if portfolio_df is provided
            if portfolio_df is not None and not portfolio_df.empty:
                for _, row in portfolio_df.iterrows():
                    sym = str(row.get("symbol", "")).strip().upper()
                    qty = float(row.get("qty", 0))
                    avg_cost = float(row.get("avg_cost", 0))
                    if sym and qty > 0:
                        cur.execute(
                            """
                            INSERT INTO portfolio_holdings (user_id, symbol, qty, avg_cost, last_updated)
                            VALUES (%s, %s, %s, %s, CURRENT_TIMESTAMP)
                            ON CONFLICT (user_id, symbol) DO UPDATE
                            SET qty = EXCLUDED.qty,
                                avg_cost = EXCLUDED.avg_cost,
                                last_updated = CURRENT_TIMESTAMP;
                            """,
                            (owner_id, sym, qty, avg_cost)
                        )
                        
            conn.commit()
            return True
    except Exception as e:
        conn.rollback()
        print(f"[db] Failed to save briefing: {e}")
        return False
    finally:
        conn.close()


def get_briefing_from_db(date_str):
    """
    Fetches a daily briefing JSON and its portfolio snapshot (if available)
    from the database and returns it as a dict.
    """
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                SELECT id, mood, news_summary, portfolio, top_picks, avoid_today, 
                       market_summary, news, dividends, portfolio_snapshot
                FROM daily_briefings
                WHERE user_id = %s AND date = %s;
                """,
                (owner_id, date_str)
            )
            row = cur.fetchone()
            if not row:
                return None
                
            # Convert RealDictRow to standard dict
            briefing_data = dict(row)
            briefing_id = briefing_data.pop("id")
            
            # Map database keys to expected dictionary keys
            if "portfolio_snapshot" in briefing_data:
                briefing_data["_portfolio_snapshot"] = briefing_data.pop("portfolio_snapshot")
            
            # Fetch the associated signal prices to rebuild _signal_prices
            cur.execute(
                """
                SELECT symbol, signal, price_on_day, qty, nifty_on_day
                FROM stock_signals
                WHERE briefing_id = %s;
                """,
                (briefing_id,)
            )
            signals_rows = cur.fetchall()
            
            signal_prices = {}
            for s in signals_rows:
                signal_prices[s["symbol"]] = {
                    "signal": s["signal"],
                    "price_on_day": float(s["price_on_day"]),
                    "qty": float(s["qty"]),
                    "nifty_on_day": float(s["nifty_on_day"]) if s["nifty_on_day"] is not None else None
                }
            briefing_data["_signal_prices"] = signal_prices
            
            # Reconstruct _meta
            briefing_data["_meta"] = {
                "timestamp": date_str + "T00:00:00.000000",
                "date": date_str
            }
            briefing_data["date"] = date_str
            
            return briefing_data
    except Exception as e:
        print(f"[db] Failed to fetch briefing: {e}")
        return None
    finally:
        conn.close()


def get_available_briefing_dates():
    """
    Returns a list of date strings (sorted descending)
    representing all available briefings in the database.
    """
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT DISTINCT date::text 
                FROM daily_briefings 
                WHERE user_id = %s 
                ORDER BY date::text DESC;
                """,
                (owner_id,)
            )
            rows = cur.fetchall()
            return [r[0] for r in rows]
    except Exception as e:
        print(f"[db] Failed to fetch dates: {e}")
        return []
    finally:
        conn.close()


def load_all_buy_signals_from_db():
    """
    Fetches all historical BUY signals from the database.
    Returns a list of dicts.
    """
    owner_id = get_owner_user_id()
    conn = get_db_connection()
    try:
        with conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(
                """
                SELECT db.date::text as date, s.symbol, s.signal, 
                       s.price_on_day, s.qty, s.nifty_on_day
                FROM stock_signals s
                JOIN daily_briefings db ON s.briefing_id = db.id
                WHERE db.user_id = %s AND s.signal = 'BUY'
                ORDER BY db.date ASC;
                """,
                (owner_id,)
            )
            rows = cur.fetchall()
            cleaned_rows = []
            for r in rows:
                row_dict = dict(r)
                row_dict["price_on_day"] = float(row_dict["price_on_day"]) if row_dict["price_on_day"] is not None else 0.0
                row_dict["qty"] = float(row_dict["qty"]) if row_dict["qty"] is not None else 0.0
                row_dict["nifty_on_day"] = float(row_dict["nifty_on_day"]) if row_dict["nifty_on_day"] is not None else None
                cleaned_rows.append(row_dict)
            return cleaned_rows
    except Exception as e:
        print(f"[db] Failed to fetch signals: {e}")
        return []
    finally:
        conn.close()


def _derive_signals(data_dict, portfolio_df):
    """Helper to derive signals from snapshot / data_dict if _signal_prices is missing."""
    import re
    if portfolio_df is None or portfolio_df.empty:
        return {}
        
    nifty_on_day = None
    market_summary = data_dict.get("market_summary", [])
    for item in market_summary:
        label = item.get("label", "").lower()
        if "nifty 50" in label or "nifty50" in label:
            raw = item.get("value", "").replace(",", "")
            match = re.search(r"[\d]+\.?\d*", raw)
            if match:
                try:
                    nifty_on_day = float(match.group())
                except ValueError:
                    pass
                break
                
    ltp_map = {}
    qty_map = {}
    for _, row in portfolio_df.iterrows():
        sym = row.get("symbol")
        if sym:
            ltp_map[sym] = float(row.get("ltp") or 0)
            qty_map[sym] = float(row.get("qty") or 0)

    signal_prices = {}
    portfolio_signals = data_dict.get("portfolio", [])
    for entry in portfolio_signals:
        sym = entry.get("symbol")
        sig = entry.get("signal", "")
        if sym and sig in ["BUY", "AVOID", "HOLD", "WATCH"] and sym in ltp_map and ltp_map[sym] > 0:
            signal_prices[sym] = {
                "signal": sig,
                "price_on_day": ltp_map[sym],
                "qty": qty_map.get(sym, 1.0),
                "nifty_on_day": nifty_on_day
            }
    return signal_prices


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()