#!/usr/bin/env python3
"""Investment tracker — unified 401k + Robinhood trade ledger."""

import argparse
import csv
import sqlite3
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path

DB_PATH = Path(__file__).resolve().parent / "data" / "invest.sqlite3"
FIDELITY_DEFAULT_CSV = Path("/root/.hermes/document_cache/doc_82e56d222fa1_collated_holdings.csv")
ROBINHOOD_DEFAULT_CSV = Path("/root/.hermes/document_cache/doc_9bfbb0458a57_robinhood_combined_transactions.csv")

SCHEMA = """
CREATE TABLE IF NOT EXISTS accounts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL UNIQUE,
    provider TEXT NOT NULL,
    account_type TEXT NOT NULL,  -- '401k', 'ira', 'taxable'
    status TEXT NOT NULL DEFAULT 'active',  -- 'active', 'closed', 'rolled_over'
    notes TEXT
);

CREATE TABLE IF NOT EXISTS trades (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id INTEGER NOT NULL REFERENCES accounts(id),
    activity_date TEXT NOT NULL,
    settle_date TEXT,
    symbol TEXT NOT NULL,
    asset_class TEXT NOT NULL,  -- 'equity', 'option', 'crypto', 'etf'
    side TEXT NOT NULL,         -- 'buy', 'sell', 'split', 'conversion'
    quantity REAL NOT NULL,
    price REAL,
    amount_usd REAL,
    fee_usd REAL DEFAULT 0,
    cost_basis REAL,
    description TEXT,
    source TEXT,                -- 'fidelity_csv', 'robinhood_txn', 'inferred'
    source_id TEXT,             -- original row id for dedup
    notes TEXT
);
CREATE INDEX IF NOT EXISTS idx_trades_date ON trades(activity_date DESC);
CREATE INDEX IF NOT EXISTS idx_trades_symbol ON trades(symbol);
CREATE INDEX IF NOT EXISTS idx_trades_account ON trades(account_id);

CREATE TABLE IF NOT EXISTS snapshots (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id INTEGER NOT NULL REFERENCES accounts(id),
    snapshot_date TEXT NOT NULL,
    symbol TEXT NOT NULL,
    quantity REAL NOT NULL,
    price REAL,
    market_value REAL NOT NULL,
    cost_basis REAL,
    source TEXT                 -- 'fidelity_monthly'
);
CREATE INDEX IF NOT EXISTS idx_snapshots_date ON snapshots(snapshot_date DESC);
CREATE INDEX IF NOT EXISTS idx_snapshots_symbol ON snapshots(symbol);

CREATE TABLE IF NOT EXISTS deposits (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    account_id INTEGER NOT NULL REFERENCES accounts(id),
    activity_date TEXT NOT NULL,
    settle_date TEXT,
    amount_usd REAL NOT NULL,
    deposit_type TEXT NOT NULL,  -- 'contribution', 'withdrawal', 'rollover_in', 'rollover_out', 'dividend', 'interest', 'fee', 'conversion'
    description TEXT,
    source TEXT,
    source_id TEXT
);
CREATE INDEX IF NOT EXISTS idx_deposits_date ON deposits(activity_date DESC);
CREATE INDEX IF NOT EXISTS idx_deposits_account ON deposits(account_id);
"""


def get_db():
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA)
    return conn


def ensure_accounts(conn):
    """Create default accounts if they don't exist."""
    accounts = [
        ("Fidelity 401k BrokerageLink", "fidelity", "401k"),
        ("Robinhood", "robinhood", "taxable"),
    ]
    for name, provider, atype in accounts:
        conn.execute(
            "INSERT OR IGNORE INTO accounts (name, provider, account_type) VALUES (?, ?, ?)",
            (name, provider, atype),
        )
    conn.commit()


# ─────────────────────────────── IMPORTS ───────────────────────────────


def import_fidelity(conn, csv_path):
    """Import Fidelity collated monthly holdings CSV. Infers trades from quantity changes."""
    account = conn.execute(
        "SELECT id FROM accounts WHERE provider = 'fidelity'"
    ).fetchone()
    if not account:
        print("ERROR: Fidelity account not found. Run ensure_accounts first.")
        return
    account_id = account["id"]

    rows = []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for r in reader:
            rows.append(r)

    # Group by symbol, sorted by date
    symbol_history = defaultdict(list)
    for r in rows:
        if r["Section"] != "Stocks":
            continue
        symbol_history[r["Symbol"]].append(r)

    trades_inserted = 0
    snapshots_inserted = 0

    # Clear old fidelity data
    conn.execute("DELETE FROM trades WHERE account_id = ? AND source = 'fidelity_csv'", (account_id,))
    conn.execute("DELETE FROM snapshots WHERE account_id = ?", (account_id,))
    conn.execute("DELETE FROM deposits WHERE account_id = ? AND source = 'fidelity_csv'", (account_id,))

    def parse_cost_basis(cb_str):
        """Parse cost basis from string, handling 'not applicable' and 'unknown'."""
        if not cb_str or cb_str in ("not applicable", "unknown"):
            return None
        try:
            return float(cb_str)
        except ValueError:
            return None

    # Insert all snapshots
    for r in rows:
        if r["Section"] != "Stocks":
            continue
        qty = float(r["Quantity"])
        price = float(r["Price"])
        value = float(r["EndingValue"])
        cb = parse_cost_basis(r["CostBasis"])
        conn.execute(
            """INSERT INTO snapshots (account_id, snapshot_date, symbol, quantity, price, market_value, cost_basis, source)
               VALUES (?, ?, ?, ?, ?, ?, ?, 'fidelity_monthly')""",
            (account_id, r["Date"], r["Symbol"], qty, price, value, cb),
        )
        snapshots_inserted += 1

    # Infer trades from quantity changes
    for sym, history in sorted(symbol_history.items()):
        history.sort(key=lambda r: r["Date"])
        prev_qty = None
        prev_cb = None

        for r in history:
            qty = float(r["Quantity"])
            cb = parse_cost_basis(r["CostBasis"])
            price = float(r["Price"])

            if prev_qty is not None:
                qty_diff = round(qty - prev_qty, 6)
                cb_diff = round(cb - prev_cb, 2) if (cb is not None and prev_cb is not None) else 0

                # Quantity increased => buy (or dividend reinvestment)
                if qty_diff > 0.0001 and cb_diff > 0:
                    conn.execute(
                        """INSERT INTO trades (account_id, activity_date, symbol, asset_class, side, quantity, price, amount_usd, cost_basis, description, source)
                           VALUES (?, ?, ?, 'etf', 'buy', ?, NULL, ?, ?, ?, 'fidelity_csv')""",
                        (account_id, r["Date"], sym, abs(qty_diff), -abs(cb_diff), cb,
                         f"Buy {abs(qty_diff):.3f} {sym} @ ~${price:.2f} (cost basis +${cb_diff:,.2f})"),
                    )
                    trades_inserted += 1

                # Quantity decreased => sell
                elif qty_diff < -0.0001 and cb_diff < 0:
                    proceeds = abs(cb_diff)  # approximate — cost basis reduction = proceeds
                    conn.execute(
                        """INSERT INTO trades (account_id, activity_date, symbol, asset_class, side, quantity, price, amount_usd, cost_basis, description, source)
                           VALUES (?, ?, ?, 'etf', 'sell', ?, NULL, ?, ?, ?, 'fidelity_csv')""",
                        (account_id, r["Date"], sym, abs(qty_diff), abs(cb_diff), cb,
                         f"Sell {abs(qty_diff):.3f} {sym} @ ~${price:.2f} (proceeds ~${abs(cb_diff):,.2f})"),
                    )
                    trades_inserted += 1

                # Quantity unchanged but cost basis changed => could be wash sale adj or fee
                # Skip for now — too ambiguous

            prev_qty = qty
            prev_cb = cb

    # Import cash (FDRXX) as deposits — track monthly ending cash balances
    cash_rows = [r for r in rows if r["Section"] == "Core Account"]
    prev_cash = None
    for r in sorted(cash_rows, key=lambda r: r["Date"]):
        cash_val = float(r["EndingValue"])
        if prev_cash is not None:
            diff = round(cash_val - prev_cash, 2)
            if abs(diff) > 100:  # meaningful change
                conn.execute(
                    """INSERT INTO deposits (account_id, activity_date, amount_usd, deposit_type, description, source)
                       VALUES (?, ?, ?, ?, ?, 'fidelity_csv')""",
                    (account_id, r["Date"], diff,
                     "contribution" if diff > 0 else "withdrawal",
                     f"Cash change: ${diff:+,.2f} (${prev_cash:,.2f} → ${cash_val:,.2f})"),
                )
        prev_cash = cash_val

    conn.commit()
    print(f"Fidelity: {snapshots_inserted} snapshots, {trades_inserted} inferred trades imported")


def import_robinhood(conn, csv_path):
    """Import Robinhood combined transactions CSV."""
    account = conn.execute(
        "SELECT id FROM accounts WHERE provider = 'robinhood'"
    ).fetchone()
    if not account:
        print("ERROR: Robinhood account not found.")
        return
    account_id = account["id"]

    with open(csv_path) as f:
        reader = csv.DictReader(f)
        rows = [r for r in reader if r["entry_kind"] == "transaction"]

    # Clear old robinhood data
    conn.execute("DELETE FROM trades WHERE account_id = ? AND source LIKE 'robinhood%'", (account_id,))
    conn.execute("DELETE FROM deposits WHERE account_id = ? AND source LIKE 'robinhood%'", (account_id,))

    trades_inserted = 0
    deposits_inserted = 0

    for r in rows:
        date = r["activity_date"]
        asset_class = r["asset_class"]
        symbol = r["symbol"] or ""
        txn_code = (r["transaction_code"] or "").upper()
        side = r["side"]
        qty = float(r["quantity"]) if r["quantity"] else None
        price = float(r["price_usd"]) if r["price_usd"] else None
        amount = float(r["amount_usd"]) if r["amount_usd"] else 0
        fee = float(r["fee_usd"]) if r["fee_usd"] else 0
        desc = r["description"] or ""
        source_id = r["id"]

        # --- TRADES (equity, etf, crypto, option) ---
        if asset_class in ("equity", "crypto") and txn_code in ("BUY", "SELL"):
            conn.execute(
                """INSERT INTO trades (account_id, activity_date, symbol, asset_class, side, quantity, price, amount_usd, fee_usd, description, source, source_id)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'robinhood_txn', ?)""",
                (account_id, date, symbol, asset_class,
                 "buy" if txn_code == "BUY" else "sell",
                 qty, price, amount, fee, desc, source_id),
            )
            trades_inserted += 1

        elif asset_class == "option" and txn_code in ("BTO", "STC"):
            conn.execute(
                """INSERT INTO trades (account_id, activity_date, symbol, asset_class, side, quantity, price, amount_usd, fee_usd, description, source, source_id)
                   VALUES (?, ?, ?, 'option', ?, ?, ?, ?, ?, ?, 'robinhood_txn', ?)""",
                (account_id, date, symbol,
                 "buy" if txn_code == "BTO" else "sell",
                 qty, price, amount, fee, desc, source_id),
            )
            trades_inserted += 1

        # --- DEPOSITS (cash) ---
        elif asset_class == "cash" and txn_code in ("ACH", "CONV"):
            d_type = "contribution" if (amount > 0 and "deposit" in desc.lower()) else \
                     "withdrawal" if amount < 0 else \
                     "conversion" if txn_code == "CONV" else \
                     "contribution" if amount > 0 else "withdrawal"
            conn.execute(
                """INSERT INTO deposits (account_id, activity_date, amount_usd, deposit_type, description, source, source_id)
                   VALUES (?, ?, ?, ?, ?, 'robinhood_txn', ?)""",
                (account_id, date, amount, d_type, desc, source_id),
            )
            deposits_inserted += 1

        elif asset_class == "cash" and txn_code in ("CDIV", "DIV"):
            conn.execute(
                """INSERT INTO deposits (account_id, activity_date, amount_usd, deposit_type, description, source, source_id)
                   VALUES (?, ?, ?, 'dividend', ?, 'robinhood_txn', ?)""",
                (account_id, date, amount, desc, source_id),
            )
            deposits_inserted += 1

        elif asset_class == "cash" and "interest" in desc.lower():
            conn.execute(
                """INSERT INTO deposits (account_id, activity_date, amount_usd, deposit_type, description, source, source_id)
                   VALUES (?, ?, ?, 'interest', ?, 'robinhood_txn', ?)""",
                (account_id, date, amount, desc, source_id),
            )
            deposits_inserted += 1

        elif asset_class == "cash" and ("fee" in desc.lower() or "gold" in desc.lower()):
            conn.execute(
                """INSERT INTO deposits (account_id, activity_date, amount_usd, deposit_type, description, source, source_id)
                   VALUES (?, ?, ?, 'fee', ?, 'robinhood_txn', ?)""",
                (account_id, date, amount, desc, source_id),
            )
            deposits_inserted += 1

    conn.commit()
    print(f"Robinhood: {trades_inserted} trades, {deposits_inserted} deposits imported")


# ─────────────────────────────── QUERIES ───────────────────────────────


def cmd_trades(conn, symbol=None, account=None, limit=50):
    """List trades, optionally filtered."""
    query = """
        SELECT t.*, a.name as account_name
        FROM trades t JOIN accounts a ON t.account_id = a.id
        WHERE 1=1
    """
    params = []
    if symbol:
        query += " AND upper(t.symbol) = upper(?)"
        params.append(symbol)
    if account:
        query += " AND a.provider = ?"
        params.append(account.lower())
    query += " ORDER BY t.activity_date DESC LIMIT ?"
    params.append(limit)

    rows = conn.execute(query, params).fetchall()
    if not rows:
        print("No trades found.")
        return

    print(f"{'Date':<12} {'Acct':<10} {'Sym':<8} {'Side':<6} {'Qty':>10} {'Price':>10} {'Amount':>12} {'Description'}")
    print("-" * 90)
    for r in rows:
        acct_short = r["account_name"][:10]
        sym = r["symbol"][:7]
        side = r["side"]
        qty = f"{r['quantity']:,.2f}" if r["quantity"] else ""
        price = f"${r['price']:,.2f}" if r["price"] else ""
        amt = f"${r['amount_usd']:+,.2f}" if r["amount_usd"] else ""
        desc = (r["description"] or "")[:40]
        print(f"{r['activity_date']:<12} {acct_short:<10} {sym:<8} {side:<6} {qty:>10} {price:>10} {amt:>12} {desc}")


def cmd_summary(conn, account=None):
    """Show current portfolio summary across accounts."""
    where = ""
    params = []
    if account:
        where = "AND a.provider = ?"
        params.append(account.lower())

    # Get latest snapshot date per account for the header, and all positions from that date
    latest_dates = conn.execute(f"""
        SELECT account_id, MAX(snapshot_date) as latest_date
        FROM snapshots s
        JOIN accounts a ON s.account_id = a.id
        WHERE 1=1 {where}
        GROUP BY account_id
    """, params).fetchall()

    if not latest_dates:
        print("No snapshots found. Run import first.")
        return

    # Get snapshots from each account's latest date
    all_positions = []
    latest_date_map = {}
    for ld in latest_dates:
        latest_date_map[ld["account_id"]] = ld["latest_date"]
        rows = conn.execute("""
            SELECT a.name as account_name, a.provider, s.symbol, s.quantity, s.market_value, s.cost_basis,
                   s.snapshot_date
            FROM snapshots s
            JOIN accounts a ON s.account_id = a.id
            WHERE s.account_id = ? AND s.snapshot_date = ?
            ORDER BY s.market_value DESC
        """, (ld["account_id"], ld["latest_date"])).fetchall()
        all_positions.extend(rows)

    by_account = defaultdict(list)
    for r in all_positions:
        by_account[r["provider"]].append(r)

    grand_total = 0
    for provider, positions in by_account.items():
        total = sum(p["market_value"] for p in positions)
        grand_total += total
        print(f"\n{'='*70}")
        print(f"  {provider.upper()}  (as of {positions[0]['snapshot_date']})")
        print(f"{'='*70}")
        print(f"{'Symbol':<10} {'Qty':>12} {'Value':>14} {'Cost Basis':>14} {'Gain':>12}")
        print("-" * 70)
        for p in positions:
            gain = p["market_value"] - (p["cost_basis"] or 0)
            print(f"{p['symbol']:<10} {p['quantity']:>12,.3f} ${p['market_value']:>13,.2f} ${(p['cost_basis'] or 0):>13,.2f} ${gain:>+11,.2f}")
        print("-" * 70)
        print(f"{'TOTAL':<10} {'':>12} ${total:>13,.2f}")
    if len(by_account) > 1:
        print(f"\n  COMBINED TOTAL: ${grand_total:,.2f}")


def cmd_pnl(conn, year=None, account=None):
    """Show realized P&L by symbol, optionally filtered by year."""
    where = "WHERE t.side IN ('buy', 'sell')"
    params = []
    if year:
        where += " AND strftime('%Y', t.activity_date) = ?"
        params.append(str(year))
    if account:
        where += " AND a.provider = ?"
        params.append(account.lower())

    rows = conn.execute(f"""
        SELECT t.symbol, t.side, t.activity_date, t.quantity, t.amount_usd, t.description, a.name as account_name
        FROM trades t JOIN accounts a ON t.account_id = a.id
        {where}
        ORDER BY t.symbol, t.activity_date
    """, params).fetchall()

    if not rows:
        print("No trades found.")
        return

    # Group by symbol, compute realized P&L
    by_symbol = defaultdict(list)
    for r in rows:
        by_symbol[r["symbol"]].append(r)

    print(f"\n{'Symbol':<10} {'Buys':>10} {'Sells':>10} {'Net Proceeds':>14} {'Trades':>8}")
    print("-" * 60)

    total_net = 0
    for sym, trades in sorted(by_symbol.items()):
        buy_total = sum(t["amount_usd"] for t in trades if t["side"] == "buy")
        sell_total = sum(t["amount_usd"] for t in trades if t["side"] == "sell")
        # For buys, amount is negative (cash out). For sells, positive (cash in).
        # Net proceeds = sells + buys (buys are negative)
        net = sell_total + buy_total
        total_net += net
        print(f"{sym:<10} ${buy_total:>9,.2f} ${sell_total:>9,.2f} ${net:>13,.2f} {len(trades):>8}")

    print("-" * 60)
    print(f"{'TOTAL':<10} {'':>10} {'':>10} ${total_net:>13,.2f}")


def cmd_symbols(conn):
    """List all traded symbols with trade counts."""
    rows = conn.execute("""
        SELECT symbol, asset_class, COUNT(*) as cnt,
               SUM(CASE WHEN side='buy' THEN 1 ELSE 0 END) as buys,
               SUM(CASE WHEN side='sell' THEN 1 ELSE 0 END) as sells
        FROM trades
        GROUP BY symbol, asset_class
        ORDER BY cnt DESC
    """).fetchall()

    print(f"\n{'Symbol':<12} {'Class':<10} {'Trades':>8} {'Buys':>8} {'Sells':>8}")
    print("-" * 50)
    for r in rows:
        print(f"{r['symbol']:<12} {r['asset_class']:<10} {r['cnt']:>8} {r['buys']:>8} {r['sells']:>8}")


def cmd_history(conn, symbol):
    """Show holdings over time for a symbol from monthly snapshots."""
    rows = conn.execute("""
        SELECT s.snapshot_date, s.quantity, s.price, s.market_value, a.name as account_name
        FROM snapshots s JOIN accounts a ON s.account_id = a.id
        WHERE upper(s.symbol) = upper(?)
        ORDER BY s.snapshot_date
    """, (symbol,)).fetchall()

    if not rows:
        print(f"No snapshot history for {symbol}.")
        return

    print(f"\n{symbol.upper()} — Holdings Over Time")
    print(f"{'Date':<12} {'Account':<12} {'Qty':>12} {'Price':>10} {'Value':>12}")
    print("-" * 60)
    for r in rows:
        print(f"{r['snapshot_date']:<12} {r['account_name']:<12} {r['quantity']:>12,.3f} ${r['price']:>9,.2f} ${r['market_value']:>11,.2f}")


def cmd_deposits(conn, account=None, year=None):
    """Show deposit/withdrawal history."""
    where = "WHERE 1=1"
    params = []
    if account:
        where += " AND a.provider = ?"
        params.append(account.lower())
    if year:
        where += " AND strftime('%Y', d.activity_date) = ?"
        params.append(str(year))

    rows = conn.execute(f"""
        SELECT d.*, a.name as account_name
        FROM deposits d JOIN accounts a ON d.account_id = a.id
        {where}
        ORDER BY d.activity_date DESC
        LIMIT 100
    """, params).fetchall()

    print(f"\n{'Date':<12} {'Account':<12} {'Type':<14} {'Amount':>12} {'Description'}")
    print("-" * 80)
    for r in rows:
        desc = (r["description"] or "")[:40]
        print(f"{r['activity_date']:<12} {r['account_name']:<12} {r['deposit_type']:<14} ${r['amount_usd']:>11,.2f} {desc}")


# ─────────────────────────────── MAIN ───────────────────────────────


def main():
    parser = argparse.ArgumentParser(description="Investment Tracker CLI")
    sub = parser.add_subparsers(dest="command")

    # import
    p_import = sub.add_parser("import", help="Import data from CSV")
    p_import.add_argument("source", choices=["fidelity", "robinhood", "all"])
    p_import.add_argument("--csv", help="Path to CSV file")

    # trades
    p_trades = sub.add_parser("trades", help="List trades")
    p_trades.add_argument("symbol", nargs="?", help="Filter by symbol")
    p_trades.add_argument("--account", "-a", help="Filter by account (fidelity/robinhood)")
    p_trades.add_argument("--limit", "-n", type=int, default=50, help="Max results")

    # summary
    p_summary = sub.add_parser("summary", help="Current portfolio summary")
    p_summary.add_argument("--account", "-a", help="Filter by account")

    # pnl
    p_pnl = sub.add_parser("pnl", help="Realized P&L")
    p_pnl.add_argument("year", nargs="?", help="Filter by year (e.g. 2025)")
    p_pnl.add_argument("--account", "-a", help="Filter by account")

    # history
    p_hist = sub.add_parser("history", help="Holdings over time for a symbol")
    p_hist.add_argument("symbol", help="Symbol (e.g. SLV)")

    # symbols
    sub.add_parser("symbols", help="List all traded symbols")

    # deposits
    p_dep = sub.add_parser("deposits", help="Deposit/withdrawal history")
    p_dep.add_argument("--account", "-a", help="Filter by account")
    p_dep.add_argument("--year", "-y", help="Filter by year")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return

    conn = get_db()
    ensure_accounts(conn)

    if args.command == "import":
        if args.source in ("fidelity", "all"):
            csv_path = args.csv or FIDELITY_DEFAULT_CSV
            if not Path(csv_path).exists():
                print(f"ERROR: CSV not found: {csv_path}")
                sys.exit(1)
            import_fidelity(conn, csv_path)
        if args.source in ("robinhood", "all"):
            csv_path = args.csv or ROBINHOOD_DEFAULT_CSV
            if not Path(csv_path).exists():
                print(f"ERROR: CSV not found: {csv_path}")
                sys.exit(1)
            import_robinhood(conn, csv_path)

    elif args.command == "trades":
        cmd_trades(conn, args.symbol, args.account, args.limit)

    elif args.command == "summary":
        cmd_summary(conn, args.account)

    elif args.command == "pnl":
        cmd_pnl(conn, args.year, args.account)

    elif args.command == "history":
        cmd_history(conn, args.symbol)

    elif args.command == "symbols":
        cmd_symbols(conn)

    elif args.command == "deposits":
        cmd_deposits(conn, args.account, args.year)

    conn.close()


if __name__ == "__main__":
    main()
