#!/usr/bin/env python3
"""Example market-making bot for EventTrader Exchange.

A simple two-sided market maker that places bid and ask orders around the
mid price, tracks fills, and manages position/P&L.

Setup:  pip install event-trader
"""

from __future__ import annotations

import argparse, asyncio, logging, os, signal, sys
from decimal import Decimal, ROUND_DOWN
from event_trader import EventTrader

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
log = logging.getLogger("mm-bot")


# ── Position Tracker ──────────────────────────────────────────────────

class PositionTracker:
    """Tracks inventory, realized P&L, and unrealized P&L."""

    def __init__(self, symbol: str):
        self.symbol = symbol.upper()
        self.inventory = Decimal("0")
        self.realized_pnl = Decimal("0")
        self.avg_entry = Decimal("0")
        self.trade_count = 0
        self._seen_trade_ids: set[str] = set()

    def process_fills(self, trades: list[dict]) -> int:
        """Process new fills and update position. Returns count of new fills."""
        new_fills = 0
        for t in trades:
            tid = t.get("trade_id") or t.get("id") or ""
            if tid in self._seen_trade_ids:
                continue
            self._seen_trade_ids.add(tid)
            new_fills += 1
            self.trade_count += 1
            qty = Decimal(str(t.get("quantity", "0")))
            price = Decimal(str(t.get("price", "0")))
            side = t.get("side", "").lower()
            if side == "buy":
                cost = qty * price
                if self.inventory >= 0:
                    total_cost = self.avg_entry * self.inventory + cost
                    self.inventory += qty
                    self.avg_entry = (total_cost / self.inventory) if self.inventory else Decimal("0")
                else:
                    self.realized_pnl += qty * (self.avg_entry - price)
                    self.inventory += qty
            elif side == "sell":
                if self.inventory > 0:
                    self.realized_pnl += qty * (price - self.avg_entry)
                    self.inventory -= qty
                else:
                    total_cost = abs(self.avg_entry * self.inventory) + qty * price
                    self.inventory -= qty
                    self.avg_entry = (total_cost / abs(self.inventory)) if self.inventory else Decimal("0")
        return new_fills

    def unrealized_pnl(self, mid_price: Decimal) -> Decimal:
        if self.inventory == 0 or mid_price == 0:
            return Decimal("0")
        if self.inventory > 0:
            return self.inventory * (mid_price - self.avg_entry)
        return abs(self.inventory) * (self.avg_entry - mid_price)

    def summary(self, mid_price: Decimal) -> str:
        upnl = self.unrealized_pnl(mid_price)
        total = self.realized_pnl + upnl
        return (
            f"Position: {self.inventory:+} {self.symbol} | "
            f"Realized: {self.realized_pnl:+.4f} USDC | "
            f"Unrealized: {upnl:+.4f} USDC | "
            f"Total P&L: {total:+.4f} USDC | Trades: {self.trade_count}"
        )


# ── Helpers ───────────────────────────────────────────────────────────

def round_to_tick(price, tick_size):
    if tick_size <= 0: return price
    return (price / tick_size).to_integral_value(rounding=ROUND_DOWN) * tick_size

def round_to_lot(qty, lot_size):
    if lot_size <= 0: return qty
    return (qty / lot_size).to_integral_value(rounding=ROUND_DOWN) * lot_size


# ── Market Maker ──────────────────────────────────────────────────────

class MarketMaker:
    def __init__(self, client, symbol, spread_bps, order_size, refresh_interval, mode, drift_threshold_bps):
        self.client = client
        self.symbol = symbol
        self.spread_bps = spread_bps
        self.order_size = order_size
        self.refresh_interval = refresh_interval
        self.mode = mode
        self.drift_threshold_bps = drift_threshold_bps
        self.tracker = PositionTracker(symbol)
        self.running = True
        self.tick_size = Decimal("0.000001")
        self.lot_size = Decimal("1")
        self.min_order_size = Decimal("1")
        self._open_order_ids: list[str] = []

    async def initialize(self):
        """Fetch pair info (tick size, lot size) and check balance."""
        try:
            resp = await self.client.clob.list_pairs()
            pairs = resp if isinstance(resp, list) else resp.get("pairs", [])
            for p in pairs:
                base = (p.get("base_symbol") or "").lower()
                if base == self.symbol.lower():
                    self.tick_size = Decimal(str(p.get("tick_size", self.tick_size)))
                    self.lot_size = Decimal(str(p.get("lot_size", self.lot_size)))
                    self.min_order_size = Decimal(str(p.get("min_order_size", self.min_order_size)))
                    log.info("Pair: %s | tick=%s | lot=%s | min=%s", base, self.tick_size, self.lot_size, self.min_order_size)
                    break
        except Exception as e:
            log.warning("Could not fetch pair info: %s", e)
        try:
            bal = await self.client.clob.balance(self.symbol)
            log.info("Balance: %s", bal)
        except Exception as e:
            log.warning("Could not fetch balance: %s", e)

    async def cancel_all(self):
        """Cancel all open orders."""
        for oid in self._open_order_ids:
            try:
                await self.client.clob.cancel_order(self.symbol, oid, mode=self.mode)
                log.info("Cancelled %s", oid)
            except Exception as e:
                log.warning("Failed to cancel %s: %s", oid, e)
        self._open_order_ids.clear()

    async def get_mid_price(self) -> Decimal | None:
        """Get mid price from best bid/offer."""
        try:
            bbo = await self.client.clob.bbo(self.symbol)
            bid = Decimal(str(bbo.get("best_bid") or bbo.get("bid") or 0))
            ask = Decimal(str(bbo.get("best_ask") or bbo.get("ask") or 0))
            if bid > 0 and ask > 0:
                return (bid + ask) / 2
        except Exception as e:
            log.warning("Could not get mid price: %s", e)
        return None

    async def run_cycle(self, last_mid):
        """Run one quoting cycle. Returns the current mid price."""
        mid = await self.get_mid_price()
        if mid is None or mid == 0:
            log.warning("No mid price — skipping cycle")
            return last_mid

        # Check for new fills
        try:
            result = await self.client.clob.my_trades(self.symbol, limit=20)
            trades = result if isinstance(result, list) else result.get("trades", [])
            new = self.tracker.process_fills(trades)
            if new:
                log.info("*** %d new fill(s) ***", new)
        except Exception:
            pass

        # Decide whether to re-quote
        should_requote = not self._open_order_ids
        if last_mid and last_mid > 0 and not should_requote:
            drift = abs(mid - last_mid) / last_mid * Decimal("10000")
            if drift > self.drift_threshold_bps:
                log.info("Mid drifted %.1f bps — re-quoting", drift)
                should_requote = True

        if should_requote:
            await self.cancel_all()

            spread = mid * Decimal(str(self.spread_bps)) / Decimal("10000")
            half = spread / 2
            bid_price = round_to_tick(mid - half, self.tick_size)
            ask_price = round_to_tick(mid + half + self.tick_size, self.tick_size)
            qty = round_to_lot(self.order_size, self.lot_size)

            if qty < self.min_order_size:
                log.warning("Order size below minimum — skipping")
                return mid
            if bid_price >= ask_price:
                ask_price = bid_price + self.tick_size

            log.info("Quoting: BID %.8f x %s | ASK %.8f x %s", bid_price, qty, ask_price, qty)

            # Place bid
            try:
                r = await self.client.clob.place_order(
                    self.symbol, "buy", str(qty),
                    price=str(bid_price), order_type="post_only", mode=self.mode,
                )
                oid = r.get("order_id", r.get("id", ""))
                if oid: self._open_order_ids.append(oid)
                log.info("  BID placed: %s", oid)
            except Exception as e:
                log.warning("  BID failed: %s", e)

            # Place ask
            try:
                r = await self.client.clob.place_order(
                    self.symbol, "sell", str(qty),
                    price=str(ask_price), order_type="post_only", mode=self.mode,
                )
                oid = r.get("order_id", r.get("id", ""))
                if oid: self._open_order_ids.append(oid)
                log.info("  ASK placed: %s", oid)
            except Exception as e:
                log.warning("  ASK failed: %s", e)
        else:
            # Refresh tracked orders (some may have filled)
            try:
                result = await self.client.clob.open_orders(self.symbol, mode=self.mode)
                orders = result if isinstance(result, list) else result.get("orders", [])
                self._open_order_ids = [o.get("order_id", o.get("id", "")) for o in orders]
            except Exception:
                pass

        log.info("Mid: %.8f | Orders: %d | %s", mid, len(self._open_order_ids), self.tracker.summary(mid))
        return mid

    async def run(self):
        """Main loop."""
        log.info("Starting %s/USDC | spread=%dbps | size=%s | mode=%s",
                 self.symbol.upper(), self.spread_bps, self.order_size, self.mode)
        await self.initialize()
        last_mid = None
        while self.running:
            try:
                last_mid = await self.run_cycle(last_mid)
            except Exception as e:
                log.error("Cycle error: %s", e)
            for _ in range(int(self.refresh_interval * 10)):
                if not self.running: break
                await asyncio.sleep(0.1)

        # Shutdown — cancel all open orders
        log.info("Shutting down — cancelling orders...")
        await self.cancel_all()
        log.info("Final: %s", self.tracker.summary(last_mid or Decimal("0")))


# ── Clone Bot Manager ─────────────────────────────────────────────────

class CloneBotManager:
    """Thin async wrapper around the /api/v1/cloned-bots/ REST API."""
    BASE = "/api/v1/cloned-bots"

    def __init__(self, http):
        self._http = http

    async def create_clone(self, source_type, source_id, name=None, is_paper=None):
        return await self._http.post(f"{self.BASE}/clone", json={
            "source_type": source_type, "source_id": source_id,
            "custom_name": name, "is_paper": is_paper,
        })

    async def get_profile(self, cid):
        return await self._http.get(f"{self.BASE}/{cid}/profile")

    async def get_balance(self, cid):
        return await self._http.get(f"{self.BASE}/{cid}/balance")

    async def fund(self, cid, amount):
        return await self._http.post(f"{self.BASE}/{cid}/fund", json={"amount": amount})

    async def withdraw(self, cid, amount):
        return await self._http.post(f"{self.BASE}/{cid}/withdraw", json={"amount": amount})

    async def toggle_trading(self, cid, enabled):
        return await self._http.post(
            f"{self.BASE}/{cid}/toggle-trading", params={"enabled": str(enabled).lower()})

    async def get_skills(self, cid):
        return await self._http.get(f"{self.BASE}/{cid}/skills")

    async def equip_skill(self, cid, skill_id, slot_position=None):
        body = {"skill_id": skill_id}
        if slot_position is not None: body["slot_position"] = slot_position
        return await self._http.post(f"{self.BASE}/{cid}/skills/equip", json=body)

    async def unequip_skill(self, cid, skill_id):
        return await self._http.delete(f"{self.BASE}/{cid}/skills/{skill_id}")


# ── CLI ───────────────────────────────────────────────────────────────

async def main():
    p = argparse.ArgumentParser(description="EventTrader market-making bot")
    p.add_argument("--pair", default="vaix")
    p.add_argument("--api-key", default=os.environ.get("ET_API_KEY"))
    p.add_argument("--email", default=os.environ.get("ET_EMAIL"))
    p.add_argument("--password", default=os.environ.get("ET_PASSWORD"))
    p.add_argument("--base-url", default="https://cymetica.com")
    p.add_argument("--spread-bps", type=int, default=50)
    p.add_argument("--size", type=Decimal, default=Decimal("100"))
    p.add_argument("--refresh", type=float, default=10.0)
    p.add_argument("--drift-bps", type=int, default=25)
    p.add_argument("--live", action="store_true")
    # Clone mode
    p.add_argument("--clone-id", help="Attach to existing clone")
    p.add_argument("--clone-from", help="Create clone from species slug")
    p.add_argument("--clone-name", help="Custom name for new clone")
    p.add_argument("--fund-amount", type=float, default=0)
    p.add_argument("--max-loss", type=Decimal, default=None)
    p.add_argument("--max-position", type=Decimal, default=None)
    p.add_argument("--withdraw-on-stop", action="store_true")
    p.add_argument("--equip-skill", action="append", default=[])
    args = p.parse_args()

    # Authenticate
    if args.api_key:
        client = EventTrader(api_key=args.api_key, base_url=args.base_url)
    elif args.email and args.password:
        client = await EventTrader.from_credentials(args.email, args.password, base_url=args.base_url)
    else:
        log.error("Provide --api-key or --email/--password")
        sys.exit(1)

    clone_mode = bool(args.clone_id or args.clone_from)
    mode = "live" if (clone_mode or args.live) else "paper"
    clone_id = clone_manager = None

    if clone_mode:
        clone_manager = CloneBotManager(client._http)
        if args.clone_from:
            resp = await clone_manager.create_clone("wta_species", args.clone_from,
                                                   args.clone_name, mode != "live")
            clone_id = resp["id"]
        else:
            clone_id = args.clone_id

        # Equip skills, fund, enable trading
        for sid in args.equip_skill:
            await clone_manager.equip_skill(clone_id, sid)
        if args.fund_amount > 0:
            await clone_manager.fund(clone_id, args.fund_amount)
        await clone_manager.toggle_trading(clone_id, True)

    bot = MarketMaker(client, args.pair, args.spread_bps, args.size, args.refresh, mode, args.drift_bps)

    # Graceful shutdown on Ctrl+C
    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, lambda: setattr(bot, "running", False))

    async with client:
        await bot.run()
        if clone_manager and clone_id:
            await clone_manager.toggle_trading(clone_id, False)
            if args.withdraw_on_stop:
                bal = await clone_manager.get_balance(clone_id)
                avail = bal.get("available") or bal.get("balance", 0)
                if avail > 0:
                    await clone_manager.withdraw(clone_id, avail)


if __name__ == "__main__":
    asyncio.run(main())