"""Local strategy workbench. Run with a licensed VectorBT PRO installation."""

from __future__ import annotations

import argparse
from functools import lru_cache
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
from pathlib import Path
import platform
import threading
from time import perf_counter
from urllib.parse import urlsplit

import numpy as np
import pandas as pd
import vectorbtpro as vbt

from data_sources import EXCHANGES, fetch_history, provider_catalog


ROOT = Path(__file__).resolve().parent
INITIAL_CASH = 10_000.0
MAX_BARS = 50_000
LOOKBACK_GRID = list(range(10, 81, 5))
KC_GRID = [float(x) / 4 for x in range(4, 15)]
GRID_CHUNK = 48


def validate_prices(frame):
    """Reject broken data rather than repairing it invisibly."""
    if not 250 <= len(frame) <= MAX_BARS:
        raise ValueError(f"Supply 250 to {MAX_BARS:,} bars.")
    if not isinstance(frame.index, pd.DatetimeIndex) or frame.index.tz is None:
        raise ValueError("Timestamps must have a timezone.")
    if frame.index.hasnans or not frame.index.is_monotonic_increasing:
        raise ValueError("Timestamps must be valid and sorted oldest first.")
    if frame.index.has_duplicates:
        raise ValueError("Duplicate timestamps are not allowed.")
    columns = ["Open", "High", "Low", "Close"]
    values = frame[columns].to_numpy(dtype=float)
    if not np.isfinite(values).all() or (values <= 0).any():
        raise ValueError("OHLC values must be finite and positive.")
    if (frame.High < frame[["Open", "Close", "Low"]].max(axis=1)).any():
        raise ValueError("High must be at least open, close, and low.")
    if (frame.Low > frame[["Open", "Close", "High"]].min(axis=1)).any():
        raise ValueError("Low must be at most open, close, and high.")
    if "Volume" in frame and (
        not np.isfinite(frame.Volume).all() or (frame.Volume < 0).any()
    ):
        raise ValueError("Volume must be finite and nonnegative.")
    return frame


def load_csv(path):
    frame = pd.read_csv(path, float_precision="round_trip")
    if "timestamp" not in frame:
        raise ValueError("CSV needs timestamp, Open, High, Low, Close columns.")
    # Naive input timestamps are explicitly interpreted as UTC.
    frame.index = pd.to_datetime(frame.pop("timestamp"), utc=True, errors="raise")
    required = ["Open", "High", "Low", "Close"]
    if not set(required).issubset(frame.columns):
        raise ValueError("CSV needs timestamp, Open, High, Low, Close columns.")
    for name in required + (["Volume"] if "Volume" in frame else []):
        frame[name] = pd.to_numeric(frame[name], errors="raise")
    return validate_prices(frame)


def demo_prices(regime, periods=10_000):
    """Deterministic synthetic OHLC, explicitly labeled in the interface."""
    rng = np.random.default_rng(42 if regime == "trend" else 19)
    t = np.arange(periods)
    if regime == "trend":
        drift = 0.00003 + 0.00025 * np.sin(t / 330)
        close = 100 * np.exp(np.cumsum(drift + rng.normal(0, 0.004, periods)))
    else:
        close = 100 * np.exp(0.07 * np.sin(t / 47) + rng.normal(0, 0.008, periods))
    opening = np.r_[close[0], close[:-1]] * np.exp(rng.normal(0, 0.001, periods))
    span = rng.uniform(0.001, 0.008, periods)
    return validate_prices(pd.DataFrame({
        "Open": opening,
        "High": np.maximum(opening, close) * (1 + span),
        "Low": np.minimum(opening, close) * (1 - span),
        "Close": close,
        "Volume": rng.integers(100, 10_000, periods),
    }, index=pd.date_range("2024-01-01", periods=periods, freq="h", tz="UTC")))


def prepare_features(frame, window):
    """Reuse installed VBT indicators. All windows contain past/completed bars only."""
    close = frame.Close
    bb = vbt.BBANDS.run(close, window=window, wtype="simple", alpha=1,
                        ddof=0, hide_params=True)
    center = vbt.MA.run(close, window=window, wtype="exp", adjust=False,
                        hide_params=True).ma
    atr = vbt.ATR.run(frame.High, frame.Low, close, window=window,
                      wtype="wilder", adjust=False, hide_params=True).atr
    midpoint = (frame.High.vbt.rolling_max(window) + frame.Low.vbt.rolling_min(window)) / 2
    detrended = close - (midpoint + bb.middle) / 2
    clock = pd.Series(np.arange(len(frame), dtype=float), index=frame.index)
    fit = vbt.OLS.run(clock, detrended, window=window, with_zscore=False,
                      hide_params=True)
    momentum = fit.slope * clock + fit.intercept
    return {"basis": bb.middle.to_numpy(), "std": (bb.upper - bb.middle).to_numpy(),
            "center": center.to_numpy(), "atr": atr.to_numpy(),
            "momentum": momentum.to_numpy()}


def squeeze_signals(features, bb_mult, kc_mult):
    """Compression release with positive momentum, exit on a zero-line cross down."""
    upper = features["basis"] + bb_mult * features["std"]
    lower = features["basis"] - bb_mult * features["std"]
    kc_upper = features["center"] + kc_mult * features["atr"]
    kc_lower = features["center"] - kc_mult * features["atr"]
    momentum = features["momentum"]
    valid = np.isfinite(momentum) & np.isfinite(upper) & np.isfinite(kc_upper)
    compressed = (lower > kc_lower) & (upper < kc_upper)
    released = np.zeros(len(momentum), dtype=bool)
    released[1:] = compressed[:-1] & ~compressed[1:] & valid[:-1] & valid[1:]
    entries = released & (momentum > 0)
    exits = np.zeros(len(momentum), dtype=bool)
    exits[1:] = (momentum[1:] <= 0) & (momentum[:-1] > 0)
    return entries, exits, {"bb_upper": upper, "bb_lower": lower,
        "basis": features["basis"], "kc_upper": kc_upper, "kc_lower": kc_lower,
        "momentum": momentum, "compressed": compressed, "released": released}


def simulate(frame, entries, exits, settings):
    return vbt.PF.from_signals(
        frame.Close, entries=entries, exits=exits,
        open=frame.Open, high=frame.High, low=frame.Low,
        price="nextopen", direction="longonly",
        size=settings["allocation"] / 100, size_type="percent",
        init_cash=INITIAL_CASH, fees=settings["fees"] / 10_000,
        slippage=settings["slippage"] / 10_000,
        accumulate=False, use_stops=False,
    )


class Workbench:
    def __init__(self, datasets):
        self.datasets = datasets

    @lru_cache(maxsize=32)
    def features(self, key, window):
        # Bound feature memory. Never cache a portfolio or a backtest result.
        return prepare_features(self.datasets[key]["frame"], window)

    @lru_cache(maxsize=6)
    def rsi(self, key):
        return vbt.RSI.run(self.datasets[key]["frame"].Close, window=14,
                           wtype="wilder", hide_params=True).rsi.to_numpy()

    def signals(self, settings, window=None, kc=None):
        features = self.features(settings["dataset"], int(settings["lookback"] if window is None else window))
        entries, exits, indicator = squeeze_signals(features, settings["bb_mult"], settings["kc_mult"] if kc is None else kc)
        rsi = self.rsi(settings["dataset"])
        eligible = np.isfinite(rsi) & (rsi > settings["rsi_threshold"])
        blocked = entries & ~eligible if settings["rsi_filter"] else np.zeros(len(entries), dtype=bool)
        return entries & ~blocked, exits, {**indicator, "rsi": rsi, "blocked": blocked}

    def load_history(self, raw):
        item = fetch_history(raw, MAX_BARS)
        validate_prices(item["frame"])
        key = item.pop("id")
        self.datasets[key] = item
        remote = [k for k in self.datasets if k.startswith("remote-")]
        for stale in remote[:-4]:
            del self.datasets[stale]
        return {"dataset": key, "metadata": self.metadata()}

    def settings(self, raw):
        result = {}
        limits = {"lookback": (5, 100), "bb_mult": (1, 3.5), "kc_mult": (0.5, 4), "fees": (0, 50),
                  "slippage": (0, 30), "allocation": (5, 100), "rsi_threshold": (1, 99)}
        for name, (low, high) in limits.items():
            value = raw.get(name)
            if isinstance(value, bool) or not isinstance(value, (int, float)):
                raise ValueError(f"{name} must be numeric.")
            if not np.isfinite(value) or not low <= value <= high:
                raise ValueError(f"{name} must be between {low} and {high}.")
            if name == "lookback" and value != int(value):
                raise ValueError("Lookback must be a whole number of bars.")
            result[name] = value
        if raw.get("dataset") not in self.datasets:
            raise ValueError("Unknown dataset.")
        if raw.get("period") not in ("all", "2160", "720", "development", "holdout"):
            raise ValueError("Unknown evaluation period.")
        if not isinstance(raw.get("rsi_filter"), bool):
            raise ValueError("RSI filter must be on or off.")
        try:
            split = pd.Timestamp(raw["split"], tz="UTC")
        except (ValueError, TypeError, KeyError):
            raise ValueError("Supply a split date in UTC.") from None
        if pd.isna(split):
            raise ValueError("Supply a valid split date.")
        result.update(dataset=raw["dataset"], period=raw["period"],
                      rsi_filter=raw["rsi_filter"], split=split.strftime("%Y-%m-%d"))
        return result

    def inputs(self, settings):
        key = settings["dataset"]
        full = self.datasets[key]["frame"]
        period = settings["period"]
        if period in ("development", "holdout"):
            split = pd.Timestamp(settings["split"], tz="UTC")
            boundary = int(full.index.searchsorted(split))
            start, stop = (0, boundary) if period == "development" else (boundary, len(full))
            if min(boundary, len(full)-boundary) < 100:
                raise ValueError("Split must leave at least 100 loaded bars on each side.")
        else:
            count = len(full) if period == "all" else int(period)
            start, stop = max(0, len(full)-count), len(full)
        return full.iloc[start:stop], start

    def metadata(self):
        return {"version": vbt.__version__, "python": platform.python_version(), "cash": INITIAL_CASH, "providers": provider_catalog(), "exchanges": EXCHANGES,
                "datasets": [{"id": key, "name": item["name"],
                    "source": item["source"], "bars": len(item["frame"]), "provenance": item.get("provenance")}
                    for key, item in self.datasets.items()]}

    @lru_cache(maxsize=6)
    def market(self, key, period, split):
        frame, _ = self.inputs({"dataset": key, "period": period, "split": split})
        result = {"time": [stamp.isoformat() for stamp in frame.index]}
        for name in ["Open", "High", "Low", "Close"]:
            result[name.lower()] = frame[name].round(6).tolist()
        return result

    def backtest(self, raw):
        started = perf_counter()
        settings = self.settings(raw)
        frame, start = self.inputs(settings)
        entries, exits, indicator = self.signals(settings)
        stop = start + len(frame)
        engine_started = perf_counter()
        pf = simulate(frame, entries[start:stop], exits[start:stop], settings)
        # Access lazy account arrays before stopping the engine timer.
        equity = pf.value.to_numpy()
        orders = pf.orders.values
        trades = pf.trades.values
        engine_ms = (perf_counter() - engine_started) * 1000
        peak = np.maximum.accumulate(np.r_[INITIAL_CASH, equity])[1:]
        drawdown = equity / peak - 1
        closed = trades[trades["status"] == 1]
        wins = closed["pnl"][closed["pnl"] > 0].sum()
        losses = -closed["pnl"][closed["pnl"] < 0].sum()
        # A passive benchmark has the same starting allocation and per-side costs.
        buy = frame.Open.iloc[0] * (1 + settings["slippage"] / 10_000)
        budget = INITIAL_CASH * settings["allocation"] / 100
        units = budget / (buy * (1 + settings["fees"] / 10_000))
        benchmark = INITIAL_CASH - budget + units * frame.Close.to_numpy()
        fill_rows = [{"row": int(r["idx"]), "signal": int(r["signal_idx"]),
                      "side": "Buy" if r["side"] == 0 else "Sell",
                      "price": float(r["price"]), "size": float(r["size"]),
                      "fee": float(r["fees"])} for r in orders]
        trade_rows = [{"entry": int(r["entry_idx"]), "exit": int(r["exit_idx"]),
                       "entry_price": float(r["entry_price"]),
                       "exit_price": float(r["exit_price"]),
                       "pnl": float(r["pnl"]), "return": float(r["return"])}
                      for r in closed]
        # Chart lines are sampled only for display. Statistics use every observation.
        points = np.unique(np.r_[np.linspace(0, len(frame)-1, min(1400, len(frame)),
                                            dtype=int), np.argmin(drawdown)])
        result = {"settings": settings, "bars": len(frame),
            "first_bar": frame.index[0].isoformat(), "last_bar": frame.index[-1].isoformat(),
            "blocked_entries": int(indicator["blocked"][start:stop].sum()),
            "metrics": {"return": equity[-1] / INITIAL_CASH - 1,
                        "profit": equity[-1] - INITIAL_CASH,
                        "benchmark_return": benchmark[-1] / INITIAL_CASH - 1,
                        "drawdown": -drawdown.min(), "trades": len(closed),
                        "win_rate": float((closed["pnl"] > 0).mean()) if len(closed) else None,
                        "profit_factor": float(wins / losses) if losses else None,
                        "fees": float(orders["fees"].sum()),
                        "open_position": bool(len(trades) and trades[-1]["status"] == 0)},
            "indicator": {key: values[start:stop].tolist() if values.dtype == bool
                          else nullable(values[start:stop]) for key, values in indicator.items()},
            "curve": {"row": points.tolist(), "equity": equity[points].round(4).tolist(),
                      "benchmark": benchmark[points].round(4).tolist(),
                      "drawdown": (drawdown[points] * 100).round(4).tolist()},
            "fills": fill_rows, "trades": trade_rows, "engine_ms": engine_ms}
        result["compute_ms"] = (perf_counter() - started) * 1000
        return result

    def grid(self, raw):
        started = perf_counter()
        settings = self.settings(raw)
        frame, start = self.inputs(settings)
        stop = start + len(frame)
        pairs = [(window, kc) for kc in KC_GRID for window in LOOKBACK_GRID]
        values = {}
        # Limit peak simulation memory even with the largest accepted CSV.
        for offset in range(0, len(pairs), GRID_CHUNK):
            chunk = pairs[offset:offset + GRID_CHUNK]
            masks = [self.signals(settings, window=window, kc=kc)[:2] for window, kc in chunk]
            entries = pd.DataFrame(np.column_stack([x[0][start:stop] for x in masks]), index=frame.index)
            exits = pd.DataFrame(np.column_stack([x[1][start:stop] for x in masks]), index=frame.index)
            pf = simulate(frame, entries, exits, settings)
            values.update({pair: float(value * 100) for pair, value in zip(chunk, np.asarray(pf.total_return))})
        return {"lookback": LOOKBACK_GRID, "kc_mult": KC_GRID, "bb_mult": settings["bb_mult"],
                "count": len(pairs), "chunk_size": GRID_CHUNK,
                "returns": [[values[(window, kc)] for window in LOOKBACK_GRID] for kc in KC_GRID],
                "compute_ms": (perf_counter() - started) * 1000}

    def compare(self, raw):
        settings = self.settings(raw)
        rows = []
        for period in ("development", "holdout"):
            for enabled in (False, True):
                result = self.backtest(dict(settings, period=period, rsi_filter=enabled))
                rows.append({key: result[key] for key in ("settings", "bars", "first_bar", "last_bar", "metrics", "blocked_entries")})
        return {"rows": rows, "threshold": settings["rsi_threshold"], "split": settings["split"]}


def nullable(values):
    return [round(float(x), 6) if np.isfinite(x) else None for x in values]


def make_handler(workbench):
    # One simulation at a time avoids competing requests and first-use compilation races.
    lock = threading.Lock()
    charts = (ROOT / "vendor" / "lightweight-charts.js").read_bytes()

    class Handler(BaseHTTPRequestHandler):
        def send(self, data, content_type, status=200):
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(len(data)))
            self.send_header("X-Content-Type-Options", "nosniff")
            self.end_headers()
            try:
                self.wfile.write(data)
            except (BrokenPipeError, ConnectionResetError):
                pass

        def json(self, data, status=200):
            self.send(json.dumps(data, allow_nan=False).encode(), "application/json", status)

        def do_GET(self):
            route = urlsplit(self.path).path
            if route == "/":
                self.send((ROOT / "dashboard.html.txt").read_bytes(), "text/html; charset=utf-8")
            elif route == "/charts.js":
                self.send(charts, "text/javascript; charset=utf-8")
            elif route == "/api/meta":
                with lock:
                    self.json(workbench.metadata())
            else:
                self.json({"error": "Not found"}, 404)

        def do_POST(self):
            # Browser tools on another origin should not invoke this local calculator.
            origin = self.headers.get("Origin")
            if origin and origin != f"http://{self.headers.get('Host')}":
                return self.json({"error": "Origin not allowed"}, 403)
            try:
                length = int(self.headers.get("Content-Length", "0"))
                if not 0 < length <= 4096:
                    raise ValueError("Invalid request size.")
                raw = json.loads(self.rfile.read(length))
                if not isinstance(raw, dict):
                    raise ValueError("Expected a settings object.")
                route = urlsplit(self.path).path
                with lock:
                    if route == "/api/backtest":
                        self.json(workbench.backtest(raw))
                    elif route == "/api/market":
                        settings = workbench.settings(raw)
                        self.json(workbench.market(settings["dataset"], settings["period"], settings["split"]))
                    elif route == "/api/history":
                        self.json(workbench.load_history(raw))
                    elif route == "/api/compare":
                        self.json(workbench.compare(raw))
                    elif route == "/api/grid":
                        self.json(workbench.grid(raw))
                    else:
                        self.json({"error": "Not found"}, 404)
            except (ValueError, KeyError, TypeError) as exc:
                self.json({"error": str(exc)}, 400)
            except Exception:
                import traceback
                traceback.print_exc()
                self.json({"error": "Calculation failed. See the server terminal."}, 500)

        def log_message(self, *_):
            pass

    return Handler


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", type=Path, help="Optional local timestamp/OHLC CSV")
    parser.add_argument("--name", default="My market data")
    parser.add_argument("--port", type=int, default=8765)
    args = parser.parse_args()
    datasets = {
        "trend": {"name": "Trend / synthetic", "source": "Synthetic hourly OHLC. Seed 42. Not market data.", "frame": demo_prices("trend")},
        "range": {"name": "Range / synthetic", "source": "Synthetic hourly OHLC. Seed 19. Not market data.", "frame": demo_prices("range")},
    }
    if args.csv:
        datasets["csv"] = {"name": args.name, "source": f"Local CSV: {args.csv.name}. Naive timestamps interpreted as UTC.", "frame": load_csv(args.csv)}
    started = perf_counter()
    print("Preparing squeeze indicators and warming the backtest engine...", flush=True)
    workbench = Workbench(datasets)
    workbench.backtest(dict(dataset="trend", period="all", lookback=20, bb_mult=2, kc_mult=1.5,
                           fees=5, slippage=2, allocation=95, rsi_filter=False, rsi_threshold=55, split="2024-01-01"))
    print(f"Ready in {perf_counter()-started:.2f}s. Open http://127.0.0.1:{args.port}", flush=True)
    server = ThreadingHTTPServer(("127.0.0.1", args.port), make_handler(workbench))
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()