"""
FTDT Quant Lab — PostgreSQL Persistence Layer.

Tables:
  strategies_snap  — per-tick strategy state (PnL, position, trades)
  trade_log        — every fill with PnL attribution
  equity_history   — per-strategy equity curve
  fill_tracker     — seen_fills persistence (prevents cross-restart blocking)
"""

import os, json, time
import psycopg2
from datetime import datetime

DB = os.getenv("FTDT_DB", "dbname=ftdt_quant user=ftdt password=ftdt_quant_2024 host=localhost")

def get_conn():
    return psycopg2.connect(DB)

def init_db():
    """Create tables if they don't exist."""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        CREATE TABLE IF NOT EXISTS strategies_snap (
            id SERIAL PRIMARY KEY,
            ts TIMESTAMPTZ DEFAULT NOW(),
            name TEXT NOT NULL,
            pnl DOUBLE PRECISION DEFAULT 0,
            pnl_pct DOUBLE PRECISION DEFAULT 0,
            position DOUBLE PRECISION DEFAULT 0,
            trades_today INTEGER DEFAULT 0,
            wins INTEGER DEFAULT 0,
            win_rate DOUBLE PRECISION DEFAULT 0,
            equity DOUBLE PRECISION DEFAULT 100,
            status TEXT DEFAULT 'idle',
            instrument TEXT DEFAULT ''
        );
        CREATE INDEX IF NOT EXISTS idx_strat_name_ts ON strategies_snap(name, ts);
        
        CREATE TABLE IF NOT EXISTS trade_log (
            id SERIAL PRIMARY KEY,
            ts TIMESTAMPTZ DEFAULT NOW(),
            strategy TEXT NOT NULL,
            side TEXT,
            size DOUBLE PRECISION,
            price DOUBLE PRECISION,
            pnl DOUBLE PRECISION DEFAULT 0,
            fee DOUBLE PRECISION DEFAULT 0,
            fill_tid BIGINT,
            reason TEXT DEFAULT ''
        );
        CREATE INDEX IF NOT EXISTS idx_trade_strat_ts ON trade_log(strategy, ts);
        
        CREATE TABLE IF NOT EXISTS equity_history (
            id SERIAL PRIMARY KEY,
            ts TIMESTAMPTZ DEFAULT NOW(),
            strategy TEXT NOT NULL,
            equity DOUBLE PRECISION
        );
        CREATE INDEX IF NOT EXISTS idx_equity_strat_ts ON equity_history(strategy, ts);
        
        CREATE TABLE IF NOT EXISTS fill_tracker (
            tid BIGINT PRIMARY KEY,
            seen_at TIMESTAMPTZ DEFAULT NOW()
        );
    """)
    conn.commit()
    cur.close()
    conn.close()
    return True

def save_strategies(strategies: dict):
    """Save current strategy states to PG."""
    conn = get_conn()
    cur = conn.cursor()
    now = datetime.utcnow()
    for name, s in strategies.items():
        cur.execute(
            "INSERT INTO strategies_snap (ts, name, pnl, pnl_pct, position, trades_today, wins, win_rate, equity, status, instrument) "
            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
            (now, name,
             s.get("pnl", 0), s.get("pnl_pct", 0), s.get("position", 0),
             s.get("trades_today", 0), s.get("wins", 0), s.get("win_rate", 0),
             s.get("allocation", 100) + s.get("pnl", 0),
             s.get("status", "idle"), s.get("instrument", ""))
        )
    conn.commit()
    cur.close()
    conn.close()

def save_trade(strategy: str, side: str, size: float, price: float, pnl: float, fee: float, tid: int, reason: str = ""):
    """Save a single trade fill to PG."""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO trade_log (ts, strategy, side, size, price, pnl, fee, fill_tid, reason) "
        "VALUES (NOW(), %s, %s, %s, %s, %s, %s, %s, %s)",
        (strategy, side, size, price, pnl, fee, tid, reason)
    )
    conn.commit()
    cur.close()
    conn.close()

def save_equity(strategy: str, equity: float):
    """Save equity point for a strategy."""
    conn = get_conn()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO equity_history (ts, strategy, equity) VALUES (NOW(), %s, %s)",
        (strategy, equity)
    )
    conn.commit()
    cur.close()
    conn.close()

# ═══════════ Fill Tracker (seen_fills) ═══════════

def load_fill_tracker() -> set:
    """Load seen_fills from PG — avoids reloading ALL history from API on restart."""
    seen = set()
    try:
        conn = get_conn()
        cur = conn.cursor()
        cur.execute("SELECT tid FROM fill_tracker")
        for row in cur.fetchall():
            seen.add(row[0])
        cur.close()
        conn.close()
    except Exception:
        pass
    return seen

def save_fill_tids(tids: set):
    """Batch save new fill TIDs to PG."""
    if not tids:
        return
    conn = get_conn()
    cur = conn.cursor()
    for tid in tids:
        try:
            cur.execute(
                "INSERT INTO fill_tracker (tid) VALUES (%s) ON CONFLICT (tid) DO NOTHING",
                (tid,)
            )
        except Exception:
            pass
    conn.commit()
    cur.close()
    conn.close()

# ═══════════ Query Helpers ═══════════

def get_trades(strategy: str = None, limit: int = 200):
    conn = get_conn()
    cur = conn.cursor()
    if strategy:
        cur.execute("SELECT * FROM trade_log WHERE strategy=%s ORDER BY ts DESC LIMIT %s", (strategy, limit))
    else:
        cur.execute("SELECT * FROM trade_log ORDER BY ts DESC LIMIT %s", (limit,))
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return rows

def get_equity(strategy: str, limit: int = 500):
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT ts, equity FROM equity_history WHERE strategy=%s ORDER BY ts ASC LIMIT %s", (strategy, limit))
    rows = cur.fetchall()
    cur.close()
    conn.close()
    return [(str(r[0]), r[1]) for r in rows]