"""
Backtesting engine: executes user-written Python strategies against historical kline data.
Computes performance metrics, equity curve, drawdown, trade log, etc.
"""

import logging
import traceback
import uuid
from datetime import datetime
from typing import Any

import numpy as np

from app.models.schemas import KlineData

logger = logging.getLogger(__name__)


class BacktestContext:
    """Injected into user strategy code as `ctx`. Provides data + order API."""

    def __init__(self, klines: list[dict], initial_capital: float, commission: float, slippage: float):
        self.klines = klines
        self.initial_capital = initial_capital
        self.commission = commission
        self.slippage = slippage

        self.capital = initial_capital
        self.position = 0
        self.position_side = ""
        self.entry_price = 0.0
        self.current_bar = 0

        self.trades: list[dict] = []
        self.equity_curve: list[float] = []
        self.daily_returns: list[float] = []
        self._signals: list[dict] = []

    @property
    def bar(self) -> dict:
        return self.klines[self.current_bar]

    @property
    def close(self) -> float:
        return self.bar["close"]

    @property
    def open(self) -> float:
        return self.bar["open"]

    @property
    def high(self) -> float:
        return self.bar["high"]

    @property
    def low(self) -> float:
        return self.bar["low"]

    @property
    def volume(self) -> int:
        return self.bar["volume"]

    def closes(self, n: int = 0) -> np.ndarray:
        end = self.current_bar + 1
        start = max(0, end - n) if n > 0 else 0
        return np.array([k["close"] for k in self.klines[start:end]])

    def highs(self, n: int = 0) -> np.ndarray:
        end = self.current_bar + 1
        start = max(0, end - n) if n > 0 else 0
        return np.array([k["high"] for k in self.klines[start:end]])

    def lows(self, n: int = 0) -> np.ndarray:
        end = self.current_bar + 1
        start = max(0, end - n) if n > 0 else 0
        return np.array([k["low"] for k in self.klines[start:end]])

    def volumes(self, n: int = 0) -> np.ndarray:
        end = self.current_bar + 1
        start = max(0, end - n) if n > 0 else 0
        return np.array([k["volume"] for k in self.klines[start:end]])

    def sma(self, period: int) -> float:
        c = self.closes(period)
        return float(np.mean(c)) if len(c) >= period else 0.0

    def ema(self, period: int) -> float:
        c = self.closes(period * 2)
        if len(c) < period:
            return 0.0
        weights = np.exp(np.linspace(-1., 0., period))
        weights /= weights.sum()
        return float(np.convolve(c, weights, mode='valid')[-1])

    def std(self, period: int) -> float:
        c = self.closes(period)
        return float(np.std(c)) if len(c) >= period else 0.0

    def highest(self, n: int) -> float:
        return float(np.max(self.highs(n)))

    def lowest(self, n: int) -> float:
        return float(np.min(self.lows(n)))

    def buy(self, quantity: int = 1, price: float | None = None):
        exec_price = price or self.close
        exec_price *= (1 + self.slippage)
        cost = exec_price * quantity * (1 + self.commission)

        if self.position < 0:
            pnl = (self.entry_price - exec_price) * abs(self.position)
            self.capital += pnl - abs(pnl) * self.commission
            self._record_trade("CLOSE_SHORT", exec_price, abs(self.position), pnl)
            self.position = 0

        self.position += quantity
        self.position_side = "LONG"
        self.entry_price = exec_price
        self.capital -= cost
        self._signals.append({"bar": self.current_bar, "type": "BUY", "price": exec_price, "qty": quantity})

    def sell(self, quantity: int = 1, price: float | None = None):
        exec_price = price or self.close
        exec_price *= (1 - self.slippage)

        if self.position > 0:
            pnl = (exec_price - self.entry_price) * self.position
            self.capital += pnl - abs(pnl) * self.commission
            self._record_trade("CLOSE_LONG", exec_price, self.position, pnl)
            self.position = 0

        self.position -= quantity
        self.position_side = "SHORT"
        self.entry_price = exec_price
        self.capital += exec_price * quantity * (1 - self.commission)
        self._signals.append({"bar": self.current_bar, "type": "SELL", "price": exec_price, "qty": quantity})

    def close_position(self):
        if self.position > 0:
            self.sell(abs(self.position))
        elif self.position < 0:
            self.buy(abs(self.position))

    def _record_trade(self, action: str, price: float, qty: int, pnl: float):
        self.trades.append({
            "bar": self.current_bar,
            "timestamp": self.bar.get("timestamp", ""),
            "action": action,
            "price": round(price, 2),
            "quantity": qty,
            "pnl": round(pnl, 2),
            "capital": round(self.capital, 2),
        })

    def _update_equity(self):
        unrealized = 0.0
        if self.position > 0:
            unrealized = (self.close - self.entry_price) * self.position
        elif self.position < 0:
            unrealized = (self.entry_price - self.close) * abs(self.position)
        equity = self.capital + unrealized
        self.equity_curve.append(round(equity, 2))
        if len(self.equity_curve) > 1:
            prev = self.equity_curve[-2]
            ret = (equity - prev) / prev if prev != 0 else 0
            self.daily_returns.append(ret)
        else:
            self.daily_returns.append(0.0)


def _compute_metrics(ctx: BacktestContext) -> dict:
    eq = np.array(ctx.equity_curve)
    if len(eq) < 2:
        return {}

    total_return = (eq[-1] - ctx.initial_capital) / ctx.initial_capital
    returns = np.array(ctx.daily_returns)

    peak = np.maximum.accumulate(eq)
    drawdown = (eq - peak) / peak
    max_dd = float(np.min(drawdown))

    sharpe = float(np.mean(returns) / np.std(returns) * np.sqrt(252)) if np.std(returns) > 0 else 0
    sortino_denom = np.std(returns[returns < 0]) if len(returns[returns < 0]) > 0 else 1e-10
    sortino = float(np.mean(returns) / sortino_denom * np.sqrt(252))

    wins = [t for t in ctx.trades if t["pnl"] > 0]
    losses = [t for t in ctx.trades if t["pnl"] < 0]
    win_rate = len(wins) / len(ctx.trades) * 100 if ctx.trades else 0
    avg_win = np.mean([t["pnl"] for t in wins]) if wins else 0
    avg_loss = np.mean([abs(t["pnl"]) for t in losses]) if losses else 1e-10
    profit_factor = float(sum(t["pnl"] for t in wins) / abs(sum(t["pnl"] for t in losses))) if losses else float('inf')

    commission_paid = sum(abs(t.get("pnl", 0)) * ctx.commission for t in ctx.trades)
    total_pnl = eq[-1] - ctx.initial_capital
    turnover_rate = len(ctx._signals) / max(len(ctx.equity_curve), 1)

    return {
        "initial_capital": ctx.initial_capital,
        "final_capital": round(float(eq[-1]), 2),
        "total_return": round(total_return * 100, 2),
        "total_pnl": round(float(total_pnl), 2),
        "max_drawdown": round(max_dd * 100, 2),
        "sharpe_ratio": round(sharpe, 3),
        "sortino_ratio": round(sortino, 3),
        "total_trades": len(ctx.trades),
        "winning_trades": len(wins),
        "losing_trades": len(losses),
        "win_rate": round(win_rate, 2),
        "avg_win": round(float(avg_win), 2),
        "avg_loss": round(float(avg_loss), 2),
        "profit_factor": round(profit_factor, 3) if profit_factor != float('inf') else 999.0,
        "commission_paid": round(commission_paid, 2),
        "total_bars": len(ctx.equity_curve),
        "turnover_rate": round(turnover_rate * 100, 2),
    }


def run_backtest(
    code: str,
    klines: list[KlineData],
    initial_capital: float = 1_000_000,
    commission: float = 0.0003,
    slippage: float = 0.0001,
) -> dict:
    kline_dicts = [
        {"open": k.open, "high": k.high, "low": k.low, "close": k.close,
         "volume": k.volume, "timestamp": k.timestamp.isoformat() if isinstance(k.timestamp, datetime) else str(k.timestamp)}
        for k in klines
    ]

    ctx = BacktestContext(kline_dicts, initial_capital, commission, slippage)
    backtest_id = f"BT-{uuid.uuid4().hex[:8].upper()}"

    user_ns: dict[str, Any] = {"ctx": ctx, "np": np}

    try:
        exec(compile(code, "<strategy>", "exec"), user_ns)
    except Exception as e:
        return {
            "backtest_id": backtest_id,
            "status": "error",
            "error": f"Strategy compilation error: {e}\n{traceback.format_exc()}",
        }

    on_bar = user_ns.get("on_bar")
    on_init = user_ns.get("on_init")

    if on_bar is None:
        return {
            "backtest_id": backtest_id,
            "status": "error",
            "error": "Strategy must define an `on_bar(ctx)` function.",
        }

    try:
        if on_init:
            on_init(ctx)

        for i in range(len(kline_dicts)):
            ctx.current_bar = i
            on_bar(ctx)
            ctx._update_equity()

        if ctx.position != 0:
            ctx.close_position()
            ctx._update_equity()

    except Exception as e:
        return {
            "backtest_id": backtest_id,
            "status": "error",
            "error": f"Runtime error at bar {ctx.current_bar}: {e}\n{traceback.format_exc()}",
        }

    metrics = _compute_metrics(ctx)
    eq = ctx.equity_curve
    peak = np.maximum.accumulate(np.array(eq))
    dd_curve = ((np.array(eq) - peak) / peak * 100).tolist()

    timestamps = [k["timestamp"] for k in kline_dicts[:len(eq)]]
    returns_hist = np.array(ctx.daily_returns)
    # If final forced-close adds one extra equity point, pad close series
    close_vals = [k["close"] for k in kline_dicts]
    if len(close_vals) < len(eq):
        close_vals.extend([close_vals[-1]] * (len(eq) - len(close_vals)))
    close_series = np.array(close_vals[:len(eq)], dtype=float)
    bench_returns = np.zeros(len(close_series), dtype=float)
    if len(close_series) > 1:
        bench_returns[1:] = np.diff(close_series) / close_series[:-1]
    benchmark_curve = (1 + bench_returns).cumprod() * ctx.initial_capital
    excess_curve = np.array(eq) - benchmark_curve

    rolling_window = min(20, max(5, len(returns_hist) // 8))
    rolling_sharpe = []
    rolling_vol = []
    for i in range(len(returns_hist)):
        if i < rolling_window:
            rolling_sharpe.append(0.0)
            rolling_vol.append(0.0)
            continue
        seg = returns_hist[i - rolling_window + 1:i + 1]
        vol = float(np.std(seg))
        sharpe = float(np.mean(seg) / vol * np.sqrt(252)) if vol > 0 else 0.0
        rolling_sharpe.append(sharpe)
        rolling_vol.append(vol * np.sqrt(252))
    hist_counts, hist_edges = np.histogram(returns_hist[~np.isnan(returns_hist)], bins=50)

    return {
        "backtest_id": backtest_id,
        "status": "success",
        "metrics": metrics,
        "equity_curve": {"timestamps": timestamps, "values": [round(v, 2) for v in eq]},
        "benchmark_curve": {"timestamps": timestamps, "values": [round(float(v), 2) for v in benchmark_curve.tolist()]},
        "excess_curve": {"timestamps": timestamps, "values": [round(float(v), 2) for v in excess_curve.tolist()]},
        "drawdown_curve": {"timestamps": timestamps, "values": [round(v, 4) for v in dd_curve]},
        "rolling_stats": {
            "timestamps": timestamps,
            "rolling_sharpe": [round(float(v), 4) for v in rolling_sharpe],
            "rolling_volatility": [round(float(v), 4) for v in rolling_vol],
            "window": rolling_window,
        },
        "turnover_curve": {
            "timestamps": timestamps,
            "values": [1.0 if any(s["bar"] == i for s in ctx._signals) else 0.0 for i in range(len(timestamps))],
        },
        "returns_distribution": {
            "edges": [round(float(e) * 100, 4) for e in hist_edges.tolist()],
            "counts": hist_counts.tolist(),
        },
        "trades": ctx.trades[-200:],
        "signals": ctx._signals[-500:],
    }


STRATEGY_TEMPLATES = {
    "ma_crossover": {
        "name": "均线交叉策略",
        "description": "快慢均线金叉买入，死叉卖出",
        "code": '''# 均线交叉策略 (MA Crossover)
# ctx: 回测上下文，提供数据访问和下单接口
# ctx.sma(n): n周期简单移动平均
# ctx.buy(qty): 买入开多
# ctx.sell(qty): 卖出开空
# ctx.close_position(): 平仓

FAST = 5
SLOW = 20

def on_bar(ctx):
    if ctx.current_bar < SLOW + 1:
        return
    fast_ma = ctx.sma(FAST)
    slow_ma = ctx.sma(SLOW)
    prev_closes = ctx.closes(SLOW + 1)
    prev_fast = float(np.mean(prev_closes[-FAST-1:-1]))
    prev_slow = float(np.mean(prev_closes[-SLOW-1:-1]))

    if prev_fast <= prev_slow and fast_ma > slow_ma:
        if ctx.position <= 0:
            ctx.close_position()
            ctx.buy(1)
    elif prev_fast >= prev_slow and fast_ma < slow_ma:
        if ctx.position >= 0:
            ctx.close_position()
            ctx.sell(1)
''',
    },
    "bollinger_breakout": {
        "name": "布林带突破策略",
        "description": "价格突破上轨做多，突破下轨做空，回归中轨平仓",
        "code": '''# 布林带突破策略 (Bollinger Bands Breakout)
PERIOD = 20
STD_DEV = 2.0

def on_bar(ctx):
    if ctx.current_bar < PERIOD + 1:
        return
    ma = ctx.sma(PERIOD)
    std = ctx.std(PERIOD)
    upper = ma + STD_DEV * std
    lower = ma - STD_DEV * std
    price = ctx.close

    if price > upper and ctx.position <= 0:
        ctx.close_position()
        ctx.buy(1)
    elif price < lower and ctx.position >= 0:
        ctx.close_position()
        ctx.sell(1)
    elif ctx.position != 0 and abs(price - ma) < std * 0.3:
        ctx.close_position()
''',
    },
    "dual_thrust": {
        "name": "Dual Thrust 突破策略",
        "description": "经典日内突破策略，基于N日range计算上下轨",
        "code": '''# Dual Thrust 突破策略
LOOKBACK = 5
K1 = 0.5
K2 = 0.5

def on_bar(ctx):
    if ctx.current_bar < LOOKBACK + 2:
        return
    highs = ctx.highs(LOOKBACK + 1)[:-1]
    lows = ctx.lows(LOOKBACK + 1)[:-1]
    closes = ctx.closes(LOOKBACK + 1)[:-1]

    hh = float(np.max(highs))
    hc = float(np.max(closes))
    ll = float(np.min(lows))
    lc = float(np.min(closes))
    range_val = max(hh - lc, hc - ll)

    open_price = ctx.open
    upper = open_price + K1 * range_val
    lower = open_price - K2 * range_val

    if ctx.close > upper and ctx.position <= 0:
        ctx.close_position()
        ctx.buy(1)
    elif ctx.close < lower and ctx.position >= 0:
        ctx.close_position()
        ctx.sell(1)
''',
    },
    "rsi_mean_reversion": {
        "name": "RSI均值回归策略",
        "description": "RSI超卖买入，超买卖出",
        "code": '''# RSI 均值回归策略
PERIOD = 14
OVERSOLD = 30
OVERBOUGHT = 70

def on_bar(ctx):
    if ctx.current_bar < PERIOD + 2:
        return
    closes = ctx.closes(PERIOD + 1)
    deltas = np.diff(closes)
    gains = np.where(deltas > 0, deltas, 0)
    losses = np.where(deltas < 0, -deltas, 0)
    avg_gain = np.mean(gains[-PERIOD:])
    avg_loss = np.mean(losses[-PERIOD:])
    rs = avg_gain / avg_loss if avg_loss > 0 else 100
    rsi = 100 - (100 / (1 + rs))

    if rsi < OVERSOLD and ctx.position <= 0:
        ctx.close_position()
        ctx.buy(1)
    elif rsi > OVERBOUGHT and ctx.position >= 0:
        ctx.close_position()
        ctx.sell(1)
''',
    },
    "channel_breakout": {
        "name": "通道突破策略",
        "description": "突破N周期最高价做多，突破最低价做空",
        "code": '''# 通道突破策略 (Donchian Channel)
ENTRY_PERIOD = 20
EXIT_PERIOD = 10

def on_bar(ctx):
    if ctx.current_bar < ENTRY_PERIOD + 1:
        return
    entry_high = ctx.highest(ENTRY_PERIOD)
    entry_low = ctx.lowest(ENTRY_PERIOD)
    exit_high = ctx.highest(EXIT_PERIOD)
    exit_low = ctx.lowest(EXIT_PERIOD)

    if ctx.close > entry_high and ctx.position <= 0:
        ctx.close_position()
        ctx.buy(1)
    elif ctx.close < entry_low and ctx.position >= 0:
        ctx.close_position()
        ctx.sell(1)
    elif ctx.position > 0 and ctx.close < exit_low:
        ctx.close_position()
    elif ctx.position < 0 and ctx.close > exit_high:
        ctx.close_position()
''',
    },
}