"""QuantTrade research workstation."""

from __future__ import annotations

import os

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st
from plotly.subplots import make_subplots

from auth import get_current_user, logout
from engine import (
    get_strategy_source,
    load_ohlcv,
    run_backtest,
    run_optimization,
)
from strategies.executor import CUSTOM_TEMPLATE
from strategies.registry import list_builtins
from strategy_db import delete_strategy, init_db, list_strategies, load_strategy, save_strategy
from telemetry import capture_exception, init_telemetry

init_telemetry("quant-streamlit")
init_db()

st.set_page_config(page_title="QuantTrade", page_icon="📈", layout="wide")

DEFAULT_TICKERS = [
    t.strip().upper()
    for t in os.environ.get(
        "CORE_TICKERS",
        "SPY,QQQ,AAPL,MSFT,GOOGL,AMZN,NVDA,META,IWM,TLT",
    ).split(",")
    if t.strip()
]

METRICS = {
    "sharpe_ratio": "Sharpe Ratio",
    "sortino_ratio": "Sortino Ratio",
    "total_return": "Total Return",
    "max_drawdown": "Max Drawdown (minimize)",
}


def sidebar_account(user: str) -> None:
    st.subheader("Account")
    st.write(f"**{user}**")
    if st.button("Logout", use_container_width=True):
        logout()
        st.rerun()


def sidebar_market() -> tuple[str, float, float]:
    st.subheader("Market")
    ticker = st.selectbox("Ticker", options=DEFAULT_TICKERS)
    init_cash = st.number_input("Initial capital ($)", min_value=1000.0, value=10_000.0, step=1000.0)
    fees = st.number_input("Fees (per trade, fraction)", min_value=0.0, max_value=0.05, value=0.001, step=0.0005)
    return ticker, init_cash, fees


def sidebar_strategy_picker() -> str:
    st.subheader("Strategy")
    builtin_options = {b.key: b.display_name for b in list_builtins()}
    kind = st.radio("Type", options=["Built-in", "Custom Python"], horizontal=True)
    if kind == "Built-in":
        return st.selectbox(
            "Model",
            options=list(builtin_options.keys()),
            format_func=lambda k: builtin_options[k],
        )
    if "custom_code" not in st.session_state:
        st.session_state.custom_code = CUSTOM_TEMPLATE
    return "custom"


def render_metrics_row(result) -> None:
    c1, c2, c3, c4, c5, c6 = st.columns(6)
    c1.metric("Sharpe", f"{result.sharpe_ratio:.2f}")
    c2.metric("Sortino", f"{result.sortino_ratio:.2f}")
    c3.metric("Return", f"{result.total_return:.1%}")
    c4.metric("Max DD", f"{result.max_drawdown:.1%}")
    c5.metric("Win rate", f"{result.win_rate:.1%}")
    c6.metric("Trades", f"{result.total_trades:,}")


def render_backtest_chart(result) -> None:
    fig = make_subplots(
        rows=3,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.05,
        row_heights=[0.5, 0.25, 0.25],
        subplot_titles=(f"{result.ticker} — price & signals", "Equity curve", "Position"),
    )
    fig.add_trace(
        go.Scatter(x=result.price.index, y=result.price, name="Close", line=dict(color="#60a5fa")),
        row=1,
        col=1,
    )
    buys = result.entries & ~result.entries.shift(1, fill_value=False)
    sells = result.exits & ~result.exits.shift(1, fill_value=False)
    fig.add_trace(
        go.Scatter(
            x=result.price.index[buys],
            y=result.price[buys],
            mode="markers",
            name="Entry",
            marker=dict(color="#34d399", size=8, symbol="triangle-up"),
        ),
        row=1,
        col=1,
    )
    fig.add_trace(
        go.Scatter(
            x=result.price.index[sells],
            y=result.price[sells],
            mode="markers",
            name="Exit",
            marker=dict(color="#f87171", size=8, symbol="triangle-down"),
        ),
        row=1,
        col=1,
    )
    fig.add_trace(
        go.Scatter(x=result.equity_curve.index, y=result.equity_curve, name="Equity", line=dict(color="#a78bfa")),
        row=2,
        col=1,
    )
    position = result.entries.astype(int).replace(0, -1).cumsum().clip(lower=0)
    fig.add_trace(
        go.Scatter(x=position.index, y=position, name="In market", fill="tozeroy", line=dict(color="#22d3ee")),
        row=3,
        col=1,
    )
    fig.update_layout(height=760, template="plotly_dark", margin=dict(l=12, r=12, t=40, b=12), showlegend=False)
    st.plotly_chart(fig, use_container_width=True)


def render_heatmap(results: pd.DataFrame, x: str, y: str, metric: str) -> None:
    if x not in results.columns or y not in results.columns:
        return
    pivot = results.pivot_table(index=y, columns=x, values="score", aggfunc="mean")
    fig = px.imshow(
        pivot,
        labels=dict(x=x, y=y, color="Score"),
        color_continuous_scale="Viridis",
        aspect="auto",
        title=f"Parameter surface — {METRICS.get(metric, metric)}",
    )
    fig.update_layout(template="plotly_dark", height=420)
    st.plotly_chart(fig, use_container_width=True)


def tab_backtest(user: str, ticker: str, init_cash: float, fees: float, strategy_key: str, source_code: str) -> None:
    st.markdown("### Single run")
    st.caption("Validate one parameter set before running a full scan.")

    params: dict = {"init_cash": init_cash, "fees": fees}
    if strategy_key != "custom":
        builtin = next(b for b in list_builtins() if b.key == strategy_key)
        st.info(builtin.description)
        for key, default in builtin.default_params.items():
            if isinstance(default, int):
                params[key] = st.number_input(key, value=int(default), step=1)
            else:
                params[key] = st.number_input(key, value=float(default))
    else:
        params.update(st.session_state.get("custom_defaults", {}))

    if st.button("Run backtest", type="primary"):
        try:
            load_ohlcv(ticker)
            result = run_backtest(
                ticker=ticker,
                strategy_key=strategy_key,
                params=params,
                source_code=source_code if strategy_key == "custom" else None,
            )
            st.session_state["last_backtest"] = result
        except Exception as exc:
            capture_exception(exc)
            st.error(str(exc))

    result = st.session_state.get("last_backtest")
    if result and result.ticker == ticker.upper():
        render_metrics_row(result)
        st.json(result.params)
        render_backtest_chart(result)


def tab_optimize(user: str, ticker: str, init_cash: float, fees: float, strategy_key: str, source_code: str) -> None:
    st.markdown("### Parameter scan")
    st.caption("Exhaustively test parameter combinations and rank by objective.")

    metric = st.selectbox("Objective", options=list(METRICS.keys()), format_func=lambda m: METRICS[m])

    if strategy_key != "custom":
        builtin = next(b for b in list_builtins() if b.key == strategy_key)
        st.write("**Search space**")
        st.json(builtin.param_grid)
        combo_hint = len(builtin.param_grid.get("window_pool", []))
        if strategy_key == "ma_crossover":
            n = len(builtin.param_grid["window_pool"])
            combo_hint = n * (n - 1) // 2
        elif strategy_key == "rsi_reversion":
            import itertools

            combo_hint = sum(
                1
                for p, os, ob in itertools.product(
                    builtin.param_grid["rsi_period"],
                    builtin.param_grid["oversold"],
                    builtin.param_grid["overbought"],
                )
                if os < ob
            )
        else:
            import itertools

            combo_hint = sum(1 for _ in itertools.product(*builtin.param_grid.values()))
        st.write(f"~**{combo_hint:,}** combinations")
    else:
        st.write("Uses `PARAM_GRID` defined in your Python strategy.")

    if st.button("Run optimization", type="primary"):
        with st.spinner("Scanning parameter space…"):
            try:
                load_ohlcv(ticker)
                opt = run_optimization(
                    ticker=ticker,
                    strategy_key=strategy_key,
                    metric=metric,
                    init_cash=init_cash,
                    fees=fees,
                    source_code=source_code if strategy_key == "custom" else None,
                )
                st.session_state["last_optimization"] = opt
            except Exception as exc:
                capture_exception(exc)
                st.error(str(exc))

    opt = st.session_state.get("last_optimization")
    if opt and opt.ticker == ticker.upper() and opt.strategy_key == strategy_key:
        st.success(
            f"Tested **{opt.combinations_tested:,}** combinations · "
            f"Best {METRICS[opt.metric]}: **{opt.best_score:.3f}**"
        )
        st.write("**Optimal parameters**")
        st.json(opt.best_params)

        top = opt.results.head(25)
        st.dataframe(
            top.style.format(
                {
                    "score": "{:.3f}",
                    "sharpe_ratio": "{:.2f}",
                    "sortino_ratio": "{:.2f}",
                    "total_return": "{:.1%}",
                    "max_drawdown": "{:.1%}",
                    "win_rate": "{:.1%}",
                },
                na_rep="—",
            ),
            use_container_width=True,
            height=360,
        )

        if strategy_key == "ma_crossover":
            render_heatmap(opt.results, "fast_window", "slow_window", opt.metric)
        elif strategy_key == "rsi_reversion":
            render_heatmap(opt.results, "rsi_period", "oversold", opt.metric)

        if st.button("Apply best params to backtest"):
            st.session_state["apply_best_params"] = opt.best_params
            st.toast("Best parameters saved — switch to Backtest tab.")


def tab_editor(strategy_key: str, source_code: str) -> str:
    st.markdown("### Strategy code")
    st.caption(
        "Write Python that defines `generate_signals(close, high, low, volume, **params)` "
        "returning `(entries, exits)` booleans. Optional: `PARAM_GRID` and `DEFAULT_PARAMS`."
    )

    if strategy_key == "custom":
        code = st.text_area("Python strategy", value=source_code, height=420, label_visibility="collapsed")
        st.session_state["custom_code"] = code
        return code

    st.code(get_strategy_source(strategy_key), language="python")
    return ""


def tab_library(user: str, ticker: str, strategy_key: str, source_code: str, params: dict) -> None:
    st.markdown("### Saved strategies")
    saved = list_strategies(user)
    names = [s.name for s in saved]
    pick = st.selectbox("Load saved", ["—"] + names)

    name = st.text_input("Save as", placeholder="SPY MA sweep v1")
    c1, c2 = st.columns(2)
    with c1:
        if st.button("Save", use_container_width=True):
            if not name.strip():
                st.error("Name required.")
            else:
                save_strategy(
                    user,
                    name.strip(),
                    ticker,
                    strategy_key,
                    params,
                    source_code if strategy_key == "custom" else None,
                )
                st.success(f"Saved **{name.strip()}**")
                st.rerun()
    with c2:
        if st.button("Delete", use_container_width=True) and pick != "—":
            delete_strategy(user, pick)
            st.success(f"Deleted **{pick}**")
            st.rerun()

    if pick != "—":
        loaded = load_strategy(user, pick)
        if loaded and st.button("Apply loaded strategy", type="primary"):
            st.session_state.active_strategy_key = loaded.strategy_key
            st.session_state.active_ticker = loaded.ticker
            st.session_state.active_params = loaded.params
            if loaded.source_code:
                st.session_state.custom_code = loaded.source_code
            st.session_state.apply_best_params = {
                k: v for k, v in loaded.params.items() if k not in ("init_cash", "fees")
            }
            st.rerun()

    if saved:
        st.dataframe(
            pd.DataFrame(
                [
                    {
                        "name": s.name,
                        "ticker": s.ticker,
                        "strategy": s.strategy_key,
                        "updated": s.created_at[:19],
                    }
                    for s in saved
                ]
            ),
            use_container_width=True,
            hide_index=True,
        )


def main() -> None:
    user = get_current_user()
    if not user:
        return

    st.title("QuantTrade")
    st.caption("Research desk · Python strategies · VectorBT parameter scans · Parquet data")

    with st.sidebar:
        sidebar_account(user)
        st.divider()
        ticker, init_cash, fees = sidebar_market()
        if "active_ticker" in st.session_state:
            ticker = st.session_state.active_ticker
        st.divider()
        strategy_key = sidebar_strategy_picker()
        if "active_strategy_key" in st.session_state:
            strategy_key = st.session_state.active_strategy_key

    source_code = st.session_state.get("custom_code", CUSTOM_TEMPLATE)
    params: dict = {"init_cash": init_cash, "fees": fees}
    if strategy_key != "custom":
        builtin = next(b for b in list_builtins() if b.key == strategy_key)
        params.update(builtin.default_params)

    if "active_params" in st.session_state:
        params.update(st.session_state.active_params)

    best = st.session_state.pop("apply_best_params", None)
    if best:
        params.update(best)

    tab_bt, tab_opt, tab_code, tab_save = st.tabs(["Backtest", "Optimize", "Python", "Library"])

    with tab_code:
        source_code = tab_editor(strategy_key, source_code)
        st.session_state.custom_code = source_code

    with tab_save:
        tab_library(user, ticker, strategy_key, source_code, params)

    with tab_bt:
        tab_backtest(user, ticker, init_cash, fees, strategy_key, source_code)

    with tab_opt:
        tab_optimize(user, ticker, init_cash, fees, strategy_key, source_code)


if __name__ == "__main__":
    main()