#!/usr/bin/env python3
"""
Strategy subprocess runner.

Loads a PandasStrategy from a git worktree path, subscribes to live bar data
(polling DataAPI), runs the paper trading loop, and pushes events to the main
MCP process via ZMQ PUSH.

Usage:
    python -m dexorder.strategy.runner --config <json_config_path>

Config JSON:
    {
        "strategy_name": "My RSI Strategy",
        "impl_path": "/app/data/worktrees/my_rsi_strategy_abc1234/strategy/my_rsi_strategy/implementation.py",
        "feed_configs": [["BTC/USDT.BINANCE", 3600]],
        "allocation": 5000.0,
        "ipc_endpoint": "ipc:///tmp/dexorder-strategy-events.sock",
        "data_dir": "/app/data",
        "poll_interval": 60
    }
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sys
import threading
import time
from pathlib import Path

# Ensure the worktree's parent (which contains dexorder package) is on the path.
# Also ensure the original dexorder package is importable.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

log = logging.getLogger(__name__)


class StrategyRunner:
    """Runs a PandasStrategy in paper trading mode using DataAPI polling."""

    def __init__(self, config: dict, stop_event: threading.Event | None = None):
        self.strategy_name = config["strategy_name"]
        self.impl_path = Path(config["impl_path"])
        self.feed_configs: list[tuple[str, int]] = [
            (f[0], int(f[1])) for f in config["feed_configs"]
        ]
        self.allocation = float(config.get("allocation", 10_000.0))
        self.ipc_endpoint = config.get("ipc_endpoint", "ipc:///tmp/dexorder-strategy-events.sock")
        self.data_dir = Path(config.get("data_dir", "/app/data"))
        self.poll_interval = int(config.get("poll_interval", 60))  # seconds

        self._stop_event = stop_event or threading.Event()
        self._running = False
        self._push_socket = None
        self._strategy = None
        self._paper: "PaperAccount | None" = None
        self._last_timestamps: dict[str, int] = {}  # feed_key -> last seen timestamp_ns

    async def run(self) -> None:
        """Main async entry point."""
        self._setup_zmq()
        await self._push_event("STARTED", {})

        try:
            await self._setup_strategy()
            await self._trading_loop()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            log.exception("Strategy runner fatal error")
            await self._push_event("ERROR", {"message": str(e)})
        finally:
            await self._push_event("STOPPED", {
                "pnl": self._paper.realized_pnl() if self._paper else 0.0,
                "trade_count": self._paper.trade_count() if self._paper else 0,
            })
            self._cleanup_zmq()

    # ------------------------------------------------------------------
    # Setup
    # ------------------------------------------------------------------

    def _setup_zmq(self) -> None:
        import zmq
        ctx = zmq.Context.instance()
        self._push_socket = ctx.socket(zmq.PUSH)
        self._push_socket.connect(self.ipc_endpoint)
        log.info("Connected PUSH socket to %s", self.ipc_endpoint)

    def _cleanup_zmq(self) -> None:
        if self._push_socket:
            self._push_socket.close()

    async def _setup_strategy(self) -> None:
        from dexorder.nautilus.backtest_runner import _load_strategy_class, _setup_custom_indicators
        from dexorder.nautilus.pandas_strategy import PandasStrategyConfig, make_feed_key
        from dexorder.strategy.paper_account import PaperAccount

        # Register custom indicators
        try:
            _setup_custom_indicators(self.data_dir)
        except Exception as e:
            log.warning("Custom indicator setup failed: %s", e)

        # Load strategy class from worktree impl path
        strategy_class = _load_strategy_class(self.impl_path)
        log.info("Loaded strategy class: %s", strategy_class.__name__)

        feed_keys = tuple(make_feed_key(t, p) for t, p in self.feed_configs)
        config = PandasStrategyConfig(
            strategy_id=f"{strategy_class.__name__}-PAPER",
            feed_keys=feed_keys,
            initial_capital=self.allocation,
        )
        self._strategy = strategy_class(config=config)
        self._paper = PaperAccount(self.allocation, feed_keys[0] if feed_keys else None)

        # Wire paper account into strategy's order methods
        self._wire_paper_account(feed_keys)
        log.info("Strategy '%s' initialized with %d feed(s)", self.strategy_name, len(feed_keys))

    def _wire_paper_account(self, feed_keys: tuple) -> None:
        """Replace strategy's order methods with paper account calls."""
        paper = self._paper
        from dexorder.nautilus.pandas_strategy import make_feed_key

        def paper_buy(quantity, feed_key=None):
            fk = feed_key or (feed_keys[0] if feed_keys else "default")
            # Get current close price from last seen bars
            price = self._current_price(fk)
            if price:
                paper.buy(quantity, price, fk)
                asyncio.create_task(self._push_event("ORDER_FILLED", {
                    "side": "buy", "quantity": quantity,
                    "price": price, "feed_key": fk,
                    "pnl": paper.realized_pnl(),
                }))

        def paper_sell(quantity, feed_key=None):
            fk = feed_key or (feed_keys[0] if feed_keys else "default")
            price = self._current_price(fk)
            if price:
                paper.sell(quantity, price, fk)
                asyncio.create_task(self._push_event("ORDER_FILLED", {
                    "side": "sell", "quantity": quantity,
                    "price": price, "feed_key": fk,
                    "pnl": paper.realized_pnl(),
                }))

        def paper_flatten(feed_key=None):
            if feed_key:
                fk_list = [feed_key]
            else:
                fk_list = list(feed_keys)
            for fk in fk_list:
                price = self._current_price(fk)
                if price:
                    paper.flatten(price, fk)

        self._strategy.buy = paper_buy
        self._strategy.sell = paper_sell
        self._strategy.flatten = paper_flatten

    # ------------------------------------------------------------------
    # Trading loop
    # ------------------------------------------------------------------

    async def _trading_loop(self) -> None:
        """Poll DataAPI for new bars and call strategy.evaluate() on each update."""
        import pandas as pd
        from dexorder.api import get_api
        from dexorder.nautilus.pandas_strategy import make_feed_key

        api = get_api()
        accumulated: dict[str, list[dict]] = {
            make_feed_key(t, p): [] for t, p in self.feed_configs
        }
        self._current_prices: dict[str, float] = {}

        heartbeat_task = asyncio.create_task(self._heartbeat_loop())
        self._running = True

        try:
            while self._running and not self._stop_event.is_set():
                now = int(time.time())
                updated_any = False

                for ticker, period_seconds in self.feed_configs:
                    fk = make_feed_key(ticker, period_seconds)
                    last_ts_ns = self._last_timestamps.get(fk, 0)
                    # Request last N bars to catch up
                    lookback = now - max(last_ts_ns // 1_000_000_000, now - 7 * 24 * 3600)
                    from_time = lookback if last_ts_ns == 0 else (last_ts_ns // 1_000_000_000)

                    try:
                        df = await api.data.historical_ohlc(
                            ticker=ticker,
                            period_seconds=period_seconds,
                            start_time=from_time,
                            end_time=now,
                            extra_columns=["volume", "buy_vol", "sell_vol",
                                           "open_time", "high_time", "low_time", "close_time",
                                           "open_interest"],
                        )
                    except Exception as e:
                        log.warning("OHLC fetch failed for %s: %s", fk, e)
                        continue

                    if df.empty:
                        continue

                    # Find new bars
                    ts_col = "timestamp" if "timestamp" in df.columns else df.columns[0]
                    new_bars = df[df[ts_col] > last_ts_ns] if last_ts_ns else df

                    for _, row in new_bars.iterrows():
                        ts_ns = int(row.get(ts_col, 0))
                        entry = {
                            "timestamp": ts_ns,
                            "open": float(row.get("open", 0)),
                            "high": float(row.get("high", 0)),
                            "low": float(row.get("low", 0)),
                            "close": float(row.get("close", 0)),
                            "volume": float(row.get("volume", 0)),
                            "buy_vol": float(row.get("buy_vol", 0)) if "buy_vol" in row else None,
                            "sell_vol": float(row.get("sell_vol", 0)) if "sell_vol" in row else None,
                            "open_interest": float(row.get("open_interest", 0)) if "open_interest" in row else None,
                        }
                        accumulated[fk].append(entry)
                        self._last_timestamps[fk] = max(self._last_timestamps.get(fk, 0), ts_ns)
                        self._current_prices[fk] = entry["close"]
                        updated_any = True

                if updated_any:
                    # Build DataFrames and call evaluate
                    dfs = {fk: pd.DataFrame(rows) for fk, rows in accumulated.items() if rows}
                    try:
                        self._strategy.evaluate(dfs)
                    except Exception as e:
                        log.error("evaluate() error: %s", e)
                        await self._push_event("ERROR", {"message": f"evaluate() error: {e}"})

                    # Push PnL update
                    rpnl = self._paper.realized_pnl() if self._paper else 0.0
                    upnl = self._paper.unrealized_pnl(self._current_prices) if self._paper else 0.0
                    await self._push_event("PNL_UPDATE", {
                        "realized_pnl": rpnl,
                        "unrealized_pnl": upnl,
                        "total_pnl": rpnl + upnl,
                        "trade_count": self._paper.trade_count() if self._paper else 0,
                    })

                # Sleep in 1s increments so stop_event is checked promptly
                for _ in range(self.poll_interval):
                    if self._stop_event.is_set():
                        self._running = False
                        break
                    await asyncio.sleep(1)
        finally:
            heartbeat_task.cancel()
            try:
                await heartbeat_task
            except asyncio.CancelledError:
                pass

    async def _heartbeat_loop(self) -> None:
        while True:
            await asyncio.sleep(10)
            await self._push_event("HEARTBEAT", {})

    def _current_price(self, feed_key: str) -> float | None:
        return getattr(self, "_current_prices", {}).get(feed_key)

    # ------------------------------------------------------------------
    # Event publishing
    # ------------------------------------------------------------------

    async def _push_event(self, event_type: str, payload: dict) -> None:
        from dexorder.strategy.events import StrategyEvent, StrategyEventType
        type_map = {
            "STARTED": StrategyEventType.STARTED,
            "STOPPED": StrategyEventType.STOPPED,
            "HEARTBEAT": StrategyEventType.HEARTBEAT,
            "ORDER_FILLED": StrategyEventType.ORDER_FILLED,
            "POSITION_UPDATE": StrategyEventType.POSITION_UPDATE,
            "PNL_UPDATE": StrategyEventType.PNL_UPDATE,
            "ERROR": StrategyEventType.ERROR,
            "LOG": StrategyEventType.LOG,
        }
        et = type_map.get(event_type, StrategyEventType.LOG)
        event = StrategyEvent(
            event_type=et,
            strategy_name=self.strategy_name,
            payload=payload,
        )
        try:
            if self._push_socket:
                self._push_socket.send(event.serialize(), flags=1)  # NOBLOCK
        except Exception as e:
            log.debug("Failed to push event %s: %s", event_type, e)


def _init_api() -> None:
    """Initialize thread-local API from environment config. Non-fatal on error."""
    try:
        import yaml
        config_path = os.environ.get("CONFIG_PATH", "/app/config/config.yaml")
        secrets_path = os.environ.get("SECRETS_PATH", "/app/config/secrets.yaml")
        config_data, secrets_data = {}, {}
        if Path(config_path).exists():
            with open(config_path) as f:
                config_data = yaml.safe_load(f) or {}
        if Path(secrets_path).exists():
            with open(secrets_path) as f:
                secrets_data = yaml.safe_load(f) or {}

        data_cfg = config_data.get("data", {})
        iceberg_cfg = data_cfg.get("iceberg", {})
        relay_cfg = data_cfg.get("relay", {})

        from dexorder.api import set_api, API
        from dexorder.impl.charting_api_impl import ChartingAPIImpl
        from dexorder.impl.data_api_impl import DataAPIImpl

        data_api = DataAPIImpl(
            iceberg_catalog_uri=iceberg_cfg.get("catalog_uri", "http://iceberg-catalog:8181"),
            relay_endpoint=relay_cfg.get("endpoint", "tcp://relay:5559"),
            notification_endpoint=relay_cfg.get("notification_endpoint", "tcp://relay:5558"),
            namespace=iceberg_cfg.get("namespace", "trading"),
            s3_endpoint=iceberg_cfg.get("s3_endpoint") or secrets_data.get("s3_endpoint"),
            s3_access_key=iceberg_cfg.get("s3_access_key") or secrets_data.get("s3_access_key"),
            s3_secret_key=iceberg_cfg.get("s3_secret_key") or secrets_data.get("s3_secret_key"),
        )
        set_api(API(charting=ChartingAPIImpl(), data=data_api))
    except Exception as e:
        log.warning("API initialization failed: %s", e)


def run_thread(config: dict, stop_event: threading.Event) -> None:
    """
    Entry point for running a strategy in a daemon thread.

    Initializes a thread-local API, creates a StrategyRunner with the given
    stop_event, and runs the async trading loop until stop_event is set.
    """
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )
    _init_api()

    runner = StrategyRunner(config, stop_event=stop_event)

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        loop.run_until_complete(runner.run())
    finally:
        loop.close()


def main():
    """Subprocess entry point (backward compatibility)."""
    import signal

    parser = argparse.ArgumentParser(description="Dexorder strategy subprocess runner")
    parser.add_argument("--config", required=True, help="Path to JSON config file")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )

    with open(args.config) as f:
        config = json.load(f)

    stop_event = threading.Event()

    def _shutdown(signum, frame):
        log.info("Received signal %d, stopping runner", signum)
        stop_event.set()

    signal.signal(signal.SIGTERM, _shutdown)
    signal.signal(signal.SIGINT, _shutdown)

    run_thread(config, stop_event)


if __name__ == "__main__":
    main()
