import os
import math
from datetime import datetime
from dateutil import parser as dateparser
import pandas as pd
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt

from supabase import create_client, Client

# --- Config ---
SUPABASE_URL = os.environ.get("SUPABASE_URL", "")
SUPABASE_ANON_KEY = os.environ.get("SUPABASE_ANON_KEY", "")
DEFAULT_MAX_ROWS = int(os.environ.get("DEFAULT_MAX_ROWS", "10000"))

if not SUPABASE_URL or not SUPABASE_ANON_KEY:
    print("WARNING: SUPABASE_URL / SUPABASE_ANON_KEY not set. App will show a banner.")

def sb_client() -> Client:
    return create_client(SUPABASE_URL, SUPABASE_ANON_KEY)

# --- Data fetch ---
def fetch_decisions(limit=DEFAULT_MAX_ROWS, filters=None):
    """
    Fetch trading decisions from Supabase.
    Filters is a dict with optional keys: agent_name, asset, model, start_date, end_date.
    """
    filters = filters or {}
    supabase = sb_client()
    q = supabase.table("trading_decisions").select("*")

    if filters.get("agent_name"):
        q = q.eq("agent_name", filters["agent_name"])
    if filters.get("asset"):
        q = q.eq("asset", filters["asset"])
    if filters.get("model"):
        q = q.eq("model", filters["model"])
    if filters.get("start_date"):
        q = q.gte("date", filters["start_date"])
    if filters.get("end_date"):
        q = q.lte("date", filters["end_date"])

    # Order by date ascending for time-series correctness
    q = q.order("date", desc=False)
    # Pull up to limit
    data = q.limit(limit).execute().data
    df = pd.DataFrame(data or [])
    if not df.empty:
        # Normalize types
        df["date"] = pd.to_datetime(df["date"], errors="coerce", utc=True).dt.tz_convert(None)
        df = df.sort_values("date")
    return df

# --- Strategy logic ---
class StrategyConfig:
    def __init__(self, long_only=True, aggressive=False, fee=0.0005):
        self.long_only = long_only         # if True, no SHORT positions
        self.aggressive = aggressive       # if True, HOLD = flatten; BUY/SELL switch directly
        self.fee = float(fee)

def simulate_equity(rows: pd.DataFrame, cfg: StrategyConfig):
    """
    Simulate equity curve given rows with columns: date, price, recommended_action
    Returns: equity DataFrame with columns [date, equity], plus stats dict.
    """
    if rows.empty:
        return pd.DataFrame(columns=["date", "equity"]), {"trades": 0, "win_rate": 0.0, "ret_total": 0.0, "ret_bh": 0.0, "ret_vs_bh": 0.0, "sharpe_daily": 0.0}

    dates = rows["date"].tolist()
    prices = rows["price"].astype(float).tolist()
    actions = rows["recommended_action"].fillna("HOLD").str.upper().tolist()

    equity = []
    capital = 1.0
    fee = cfg.fee
    position = "FLAT"
    entry_price = None

    trades = []
    last_equity = capital

    # Buy & Hold baseline
    p0 = prices[0]
    bh_equity = [1.0 * (p / p0) for p in prices]
    # Returns series for sharpe (daily-ish)
    eq_series = []

    for i, (dt, price, act) in enumerate(zip(dates, prices, actions)):
        # Normalize action
        if act not in ("BUY","SELL","HOLD"):
            act = "HOLD"

        # Aggressive logic: HOLD => flatten
        if cfg.aggressive and act == "HOLD" and position != "FLAT":
            # close pos
            if position == "LONG":
                capital *= (price / entry_price) * (1 - fee)
            elif position == "SHORT":
                capital *= (entry_price / price) * (1 - fee)
            trades.append({"entry": entry_price, "exit": price, "dir": position})
            position, entry_price = "FLAT", None

        if act == "BUY":
            if position == "FLAT":
                position = "LONG"
                entry_price = price
                capital *= (1 - fee)
            elif position == "SHORT":
                # close SHORT, open LONG (if aggressive) or ignore (if baseline)
                # In both modes, we interpret a BUY while short as closing short then long
                capital *= (entry_price / price) * (1 - fee)
                trades.append({"entry": entry_price, "exit": price, "dir": "SHORT"})
                if cfg.long_only:
                    position, entry_price = "FLAT", None
                else:
                    position, entry_price = "LONG", price
                    capital *= (1 - fee)
            else:
                # already LONG: no change
                pass

        elif act == "SELL":
            if cfg.long_only:
                # In long-only, SELL means close long if any
                if position == "LONG":
                    capital *= (price / entry_price) * (1 - fee)
                    trades.append({"entry": entry_price, "exit": price, "dir": "LONG"})
                    position, entry_price = "FLAT", None
            else:
                if position == "FLAT":
                    position = "SHORT"
                    entry_price = price
                    capital *= (1 - fee)
                elif position == "LONG":
                    # close LONG, open SHORT
                    capital *= (price / entry_price) * (1 - fee)
                    trades.append({"entry": entry_price, "exit": price, "dir": "LONG"})
                    position, entry_price = "SHORT", price
                    capital *= (1 - fee)
                else:
                    # already SHORT
                    pass

        # HOLD in non-aggressive does nothing

        equity.append(capital)
        eq_series.append(capital)

    # At end, mark-to-market (position still open)
    if position != "FLAT" and entry_price is not None:
        last_price = prices[-1]
        if position == "LONG":
            mtm = capital * (last_price / entry_price)
        else:
            mtm = capital * (entry_price / last_price)
        equity[-1] = mtm

    eq_df = pd.DataFrame({"date": dates, "equity": equity})
    ret_total = (eq_df["equity"].iloc[-1] / eq_df["equity"].iloc[0]) - 1.0 if len(eq_df) > 1 else 0.0
    ret_bh = (bh_equity[-1] / bh_equity[0]) - 1.0 if len(bh_equity) > 1 else 0.0
    ret_vs_bh = ret_total - ret_bh

    # Sharpe (daily-ish): simple approximation using equity pct change
    eq_series = np.array(eq_series)
    if len(eq_series) > 1:
        rets = np.diff(eq_series) / eq_series[:-1]
        sharpe_daily = (np.mean(rets) / (np.std(rets) + 1e-9)) * np.sqrt(252)
    else:
        sharpe_daily = 0.0

    # Win rate
    wins = 0
    for t in trades:
        if t["dir"] == "LONG":
            wins += 1 if t["exit"] > t["entry"] else 0
        else:
            wins += 1 if t["exit"] < t["entry"] else 0
    win_rate = (wins / len(trades)) if trades else 0.0

    stats = {
        "trades": len(trades),
        "win_rate": win_rate,
        "ret_total": ret_total,
        "ret_bh": ret_bh,
        "ret_vs_bh": ret_vs_bh,
        "sharpe_daily": float(sharpe_daily),
    }
    return eq_df, stats

# --- Plotting ---
def plot_equities(main_eq: pd.DataFrame, bh_eq: pd.DataFrame, title: str = "Equity Curve"):
    fig = plt.figure(figsize=(8,4.5))
    ax = fig.gca()
    ax.plot(main_eq["date"], main_eq["equity"], label="Strategy")
    ax.plot(bh_eq["date"], bh_eq["equity"], label="Buy & Hold", linestyle="--")
    ax.set_title(title)
    ax.set_xlabel("Date")
    ax.set_ylabel("Equity (normalized)")
    ax.legend()
    ax.grid(True, alpha=0.3)
    return fig

def build_bh(df: pd.DataFrame):
    if df.empty: 
        return pd.DataFrame(columns=["date","equity"])
    prices = df["price"].astype(float).values
    base = prices[0]
    eq = prices / base
    return pd.DataFrame({"date": df["date"].values, "equity": eq})

# --- UI Handlers ---
def list_unique(column):
    try:
        df = fetch_decisions(limit=1000)
        if df.empty or column not in df:
            return []
        vals = df[column].dropna().unique().tolist()
        return sorted([v for v in vals if isinstance(v, str)])
    except Exception:
        return []

def run_query(agent, asset, model, start, end, long_only, aggressive, fee, limit_rows):
    if not SUPABASE_URL or not SUPABASE_ANON_KEY:
        banner = "⚠️ Missing SUPABASE_URL or SUPABASE_ANON_KEY. Set them in Space Secrets."
    else:
        banner = ""

    filters = {}
    if agent: filters["agent_name"] = agent
    if asset: filters["asset"] = asset
    if model: filters["model"] = model
    if start: filters["start_date"] = start
    if end: filters["end_date"] = end

    df = fetch_decisions(limit=limit_rows, filters=filters)
    if df.empty:
        return banner or "No data found for the selected filters.", None, pd.DataFrame(), pd.DataFrame()

    # Simulate
    cfg = StrategyConfig(long_only=long_only, aggressive=aggressive, fee=fee)
    eq_df, stats = simulate_equity(df[["date","price","recommended_action"]], cfg)
    bh_df = build_bh(df)

    # Plot
    fig = plot_equities(eq_df, bh_df, title="Equity Curve")

    # Metrics table
    metrics = pd.DataFrame([{
        "Trades": stats["trades"],
        "Win Rate": f"{stats['win_rate']*100:.1f}%",
        "Total Return": f"{stats['ret_total']*100:.1f}%",
        "Buy&Hold Return": f"{stats['ret_bh']*100:.1f}%",
        "Excess vs B&H": f"{stats['ret_vs_bh']*100:.1f}%",
        "Sharpe (daily)": f"{stats['sharpe_daily']:.2f}",
        "Rows Used": len(df)
    }])

    return banner, fig, df[["date","agent_name","asset","model","price","recommended_action"]].tail(25), metrics

with gr.Blocks(title="Paper Trading Agents") as demo:
    gr.Markdown("# 📈 Paper Trading Agents\nVisualize agent decisions from Supabase and compute strategy equity vs buy&hold.\n")
    with gr.Row():
        agent = gr.Dropdown(choices=[], label="Agent Name (optional)", interactive=True)
        asset = gr.Dropdown(choices=[], label="Asset (optional)", interactive=True)
        model = gr.Dropdown(choices=[], label="Model (optional)", interactive=True)
    with gr.Row():
        start = gr.Textbox(label="Start Date (YYYY-MM-DD, optional)")
        end = gr.Textbox(label="End Date (YYYY-MM-DD, optional)")
        limit_rows = gr.Slider(1000, 50000, value=DEFAULT_MAX_ROWS, step=500, label="Max rows")
    with gr.Row():
        long_only = gr.Checkbox(value=True, label="Long Only")
        aggressive = gr.Checkbox(value=False, label="Aggressive Mode (HOLD = flatten; BUY/SELL switch)")
        fee = gr.Number(value=0.0005, label="Fee (per open/close)")

    go = gr.Button("Run")
    banner = gr.Markdown()
    plot = gr.Plot()
    tail = gr.Dataframe(headers=["date","agent_name","asset","model","price","recommended_action"], label="Sample of latest rows", wrap=True)
    metrics = gr.Dataframe(label="Metrics", wrap=True)

    def _init_choices():
        return gr.update(choices=list_unique("agent_name")), gr.update(choices=list_unique("asset")), gr.update(choices=list_unique("model"))

    demo.load(_init_choices, None, [agent, asset, model])
    go.click(run_query, inputs=[agent, asset, model, start, end, long_only, aggressive, fee, limit_rows], outputs=[banner, plot, tail, metrics])

if __name__ == "__main__":
    demo.launch()