scripts/migrate_to_postgres.py

"""
Database Migration Script for PortIQ.
Loads local JSON history files into the PostgreSQL database.

Requirements:
    pip install psycopg2-binary python-dotenv pandas

Run from the project root:
    python scripts/migrate_to_postgres.py
"""

import os
import sys
import json
import re
from datetime import datetime
import pandas as pd
from dotenv import load_dotenv

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

# Ensure we can load packages and configs from project root
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Load env variables
load_dotenv()

# Verify psycopg2 installation
try:
    import psycopg2
    from psycopg2.extras import Json
except ImportError:
    print("❌ Error: psycopg2 is not installed.")
    print("👉 Please run: pip install psycopg2-binary")
    sys.exit(1)


HISTORY_DIR = "history"
SCHEMA_FILE = os.path.join("config", "schema.sql")
HOLDINGS_FILE = "holdings.csv"


def get_db_connection():
    """Establishes connection to the PostgreSQL database."""
    db_url = os.getenv("DATABASE_URL")
    if db_url:
        print(f"Connecting to database using DATABASE_URL...")
        return psycopg2.connect(db_url)
    
    # Try connecting using individual variables
    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")
    
    print(f"Connecting to database {dbname} on {host}:{port}...")
    return psycopg2.connect(
        host=host,
        port=port,
        dbname=dbname,
        user=user,
        password=password
    )


def initialize_schema(conn):
    """Executes the DDL schema in config/schema.sql to create tables."""
    if not os.path.exists(SCHEMA_FILE):
        print(f"❌ Schema file {SCHEMA_FILE} not found!")
        return False
        
    print("🛠️ Initializing database tables...")
    with open(SCHEMA_FILE, "r", encoding="utf-8") as f:
        schema_ddl = f.read()
        
    with conn.cursor() as cur:
        cur.execute(schema_ddl)
        
        # Schema upgrade guard: Ensure new columns exist on existing databases
        try:
            cur.execute("ALTER TABLE daily_briefings ADD COLUMN IF NOT EXISTS portfolio JSONB DEFAULT '[]'::jsonb;")
            cur.execute("ALTER TABLE daily_briefings ADD COLUMN IF NOT EXISTS portfolio_snapshot JSONB DEFAULT '[]'::jsonb;")
        except Exception as alter_err:
            print(f"⚠️ Note: daily_briefings alteration skipped: {alter_err}")

        # Schema upgrade guard: portfolio_holdings price columns (added in v2)
        for col_ddl in [
            "ALTER TABLE portfolio_holdings ADD COLUMN IF NOT EXISTS current_price NUMERIC(12,4);",
            "ALTER TABLE portfolio_holdings ADD COLUMN IF NOT EXISTS close_price_prev NUMERIC(12,4);",
            "ALTER TABLE portfolio_holdings ADD COLUMN IF NOT EXISTS csv_price NUMERIC(12,4);",
            "ALTER TABLE portfolio_holdings ADD COLUMN IF NOT EXISTS csv_close_prev NUMERIC(12,4);",
        ]:
            try:
                cur.execute(col_ddl)
            except Exception as col_err:
                print(f"⚠️ Column alteration skipped: {col_err}")
            
    conn.commit()
    print("✅ Database schema initialized successfully.")
    return True


def get_or_create_owner_user(conn):
    """Creates a default owner user account matching the env credentials."""
    owner_hash = os.getenv("OWNER_PASSWORD_HASH")
    if not owner_hash:
        # Fallback owner hash of 'tcr-owner'
        owner_hash = "084f7fa87d1dfcbb3965db0183b544b60a3cc180c59800a6e30018f70094770e"
        
    email = "owner@portiq.com"
    
    with conn.cursor() as cur:
        # Check if owner already exists
        cur.execute("SELECT id FROM users WHERE email = %s;", (email,))
        row = cur.fetchone()
        if row:
            print(f"👤 Found existing Owner user: {email} (ID: {row[0]})")
            return row[0]
            
        # Create default owner user
        print(f"👤 Creating default Owner user profile: {email}...")
        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()
        print(f"👤 Created Owner user with UUID: {owner_id}")
        return owner_id


def parse_nifty_from_market_summary(market_summary) -> float | None:
    """Helper to parse nifty 50 numeric value from market summary."""
    if not isinstance(market_summary, list):
        return None
    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:
                    return float(match.group())
                except ValueError:
                    pass
    return None


def get_signal_prices_from_snapshot(data: dict) -> dict:
    """Derives signal prices if _signal_prices field is missing from JSON."""
    snapshot = data.get("_portfolio_snapshot", [])
    portfolio_signals = data.get("portfolio", [])
    market_summary = data.get("market_summary", [])

    ltp_map = {}
    qty_map = {}
    for item in snapshot:
        sym = item.get("symbol")
        if sym:
            ltp_map[sym] = float(item.get("ltp") or 0)
            qty_map[sym] = float(item.get("qty") or 0)

    nifty_on_day = parse_nifty_from_market_summary(market_summary)

    result = {}
    for entry in portfolio_signals:
        sym = entry.get("symbol")
        sig = entry.get("signal", "")
        # Capture all signals to record history
        if sym and sig in ["BUY", "AVOID", "HOLD", "WATCH"] and sym in ltp_map and ltp_map[sym] > 0:
            result[sym] = {
                "signal": sig,
                "price_on_day": ltp_map[sym],
                "qty": qty_map.get(sym, 1.0),
                "nifty_on_day": nifty_on_day,
            }
    return result


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 migrate_history_briefings(conn, owner_id):
    """Reads history JSON files and inserts them into DB."""
    if not os.path.exists(HISTORY_DIR):
        print(f"⚠️ History directory '{HISTORY_DIR}' not found. Skipping briefings migration.")
        return
        
    files = sorted([f for f in os.listdir(HISTORY_DIR) if f.endswith(".json") and f != "performance_cache.json"])
    if not files:
        print("ℹ️ No JSON briefings found in history folder.")
        return
        
    print(f"🔄 Found {len(files)} briefing files. Migrating to database...")
    
    migrated_count = 0
    skipped_count = 0
    
    for filename in files:
        date_str = filename.replace(".json", "")
        filepath = os.path.join(HISTORY_DIR, filename)
        
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                data = json.load(f)
        except Exception as e:
            print(f"  ❌ Error reading {filename}: {e}")
            continue
            
        # Parse fields from the briefing json and clean NaN values
        mood = clean_nans(data.get("mood", {}))
        news_summary = clean_nans(data.get("news_summary", []))
        portfolio = clean_nans(data.get("portfolio", []))
        top_picks = clean_nans(data.get("top_picks", []))
        avoid_today = clean_nans(data.get("avoid_today", []))
        market_summary = clean_nans(data.get("market_summary", []))
        news = clean_nans(data.get("news", []))
        dividends = clean_nans(data.get("dividends", []))
        portfolio_snapshot = clean_nans(data.get("_portfolio_snapshot", []))
        
        # Overwrite existing records to ensure new schema columns (portfolio & snapshot) are backfilled
        with conn.cursor() as cur:
            cur.execute(
                "DELETE FROM daily_briefings WHERE user_id = %s AND date = %s RETURNING id;",
                (owner_id, date_str)
            )
            overwritten = cur.fetchone()
            if overwritten:
                print(f"  🔄 Overwriting briefing for {date_str} to update schema fields...")
                
            # Insert daily briefing
            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]
            
            # Extract and migrate signal history
            signal_prices = data.get("_signal_prices")
            if not signal_prices:
                signal_prices = get_signal_prices_from_snapshot(data)
                
            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)
                    )
            
            conn.commit()
            print(f"  ✅ Migrated briefing for {date_str} (Created {len(signal_prices)} signals)")
            migrated_count += 1
            
    print(f"🎉 Briefings migration completed: {migrated_count} inserted, {skipped_count} skipped.")


def migrate_holdings(conn, owner_id):
    """Loads CSV holdings (if present) and updates the portfolio_holdings table."""
    if not os.path.exists(HOLDINGS_FILE):
        print(f"⚠️ CSV Holdings file '{HOLDINGS_FILE}' not found. Skipping holdings migration.")
        return
        
    try:
        df = pd.read_csv(HOLDINGS_FILE)
        # Clean columns: rename typical Zerodha columns to lowercase keys
        df.columns = [c.strip().lower() for c in df.columns]
    except Exception as e:
        print(f"❌ Error loading holdings CSV: {e}")
        return
        
    synonyms = {
        'instrument': 'symbol',
        'quantity': 'qty',
        'qty.': 'qty',
        'average cost': 'avg_cost',
        'avg. cost': 'avg_cost',
        'buy average': 'avg_cost',
        'ltp': 'ltp',
        'day chg.': 'day_chg',
        'day chg': 'day_chg'
    }
    df = df.rename(columns=synonyms)
    
    required_cols = {'symbol', 'qty', 'avg_cost'}
    if not required_cols.issubset(df.columns):
        print(f"❌ Holdings CSV is missing required columns (symbol, qty, avg_cost). Found: {list(df.columns)}")
        return
            
    print(f"📋 Found holdings CSV with {len(df)} stocks. Migrating to database...")
    
    inserted_count = 0
    with conn.cursor() as cur:
        for _, row in df.iterrows():
            symbol = str(row['symbol']).strip().upper()
            qty = float(row['qty'])
            avg_cost = float(row['avg_cost'])
            ltp = float(row.get('ltp', 0.0))
            day_chg = str(row.get('day_chg', '0.0')).replace('%', '').strip()
            try:
                day_chg = float(day_chg)
            except ValueError:
                day_chg = 0.0
                
            # Calculate close_price_prev
            if ltp > 0:
                close_prev = ltp / (1.0 + day_chg / 100.0)
            else:
                close_prev = 0.0
            
            if not symbol or qty <= 0:
                continue
                
            cur.execute(
                """
                INSERT INTO portfolio_holdings (user_id, symbol, qty, avg_cost, current_price, close_price_prev, last_updated)
                VALUES (%s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP)
                ON CONFLICT (user_id, symbol) DO UPDATE
                SET qty = EXCLUDED.qty,
                    avg_cost = EXCLUDED.avg_cost,
                    current_price = EXCLUDED.current_price,
                    close_price_prev = EXCLUDED.close_price_prev,
                    last_updated = CURRENT_TIMESTAMP;
                """,
                (owner_id, symbol, qty, avg_cost, ltp, close_prev)
            )
            inserted_count += 1
            
    conn.commit()
    print(f"✅ Successfully migrated {inserted_count} holding assets to database portfolio.")


def main():
    print("🚀 Starting PortIQ PostgreSQL Database Migration...")
    
    conn = None
    try:
        conn = get_db_connection()
    except Exception as e:
        print(f"❌ Database Connection Failed: {e}")
        print("👉 Make sure your PostgreSQL server is running and .env configuration is correct.")
        sys.exit(1)
        
    try:
        # 1. Initialize schema tables
        if not initialize_schema(conn):
            print("❌ Failed to initialize schema. Exiting.")
            sys.exit(1)
            
        # 2. Setup user role
        owner_id = get_or_create_owner_user(conn)
        
        # 3. Migrate historical daily reports
        migrate_history_briefings(conn, owner_id)
        
        # 4. Migrate current holdings csv
        migrate_holdings(conn, owner_id)
        
        print("\n🏆 Migration script completed successfully!")
        
    except Exception as e:
        print(f"\n❌ Error occurred during migration: {e}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()
            print("🔌 Database connection closed.")


if __name__ == "__main__":
    main()