#!/usr/bin/env python3

import argparse
from enum import Enum
import socket
import json
from state import StateManager

team_name = "HanyangFloorFunction"

# ==================== 설정 ====================
ORDER_SIZE = 6
MAX_POS = 40
KILL_POS = 50
ARB_THRESHOLD = 10

# ====================
class Dir(str, Enum):
    BUY = "BUY"
    SELL = "SELL"

# ====================
def main():
    args = parse_arguments()
    exchange = ExchangeConnection(args)

    state = StateManager()
    hello = exchange.read_message()

    positions = {s["symbol"]: s["position"] for s in hello["symbols"]}

    order_id = 0
    def next_id():
        nonlocal order_id
        order_id += 1
        return order_id

    active_orders = {}

    # ==================== 공격형 MM ====================
    def market_make(sym):
        bid = state.bid_prices.get(sym)
        ask = state.ask_prices.get(sym)
        if bid is None or ask is None:
            return

        pos = positions.get(sym, 0)

        if abs(pos) > MAX_POS:
            return

        buy_price = bid + 2
        sell_price = ask - 2

        if buy_price >= sell_price:
            return

        buy_size = ORDER_SIZE
        sell_size = ORDER_SIZE

        if pos > 0:
            sell_size += 2
        elif pos < 0:
            buy_size += 2

        oid = next_id()
        exchange.send_add_message(oid, sym, Dir.BUY, buy_price, buy_size)
        active_orders[oid] = sym

        oid = next_id()
        exchange.send_add_message(oid, sym, Dir.SELL, sell_price, sell_size)
        active_orders[oid] = sym

    # ==================== Fill 기반 재주문 ====================
    def refill(sym):
        market_make(sym)

    # ==================== XLF Arb (공격형) ====================
    def xlf_arb():
        bond = state.bid_prices.get("BOND")
        gs = state.bid_prices.get("GS")
        ms = state.bid_prices.get("MS")
        wfc = state.bid_prices.get("WFC")
        xlf = state.ask_prices.get("XLF")

        if None in [bond, gs, ms, wfc, xlf]:
            return

        basket = bond*3 + gs*2 + ms*3 + wfc*2
        profit = basket - xlf*10

        if profit > ARB_THRESHOLD:
            exchange.send_add_message_ioc(next_id(), "XLF", Dir.BUY, xlf, 5)

    # ==================== 리스크 ====================
    def risk():
        for sym in positions:
            pos = positions[sym]

            if abs(pos) > KILL_POS:
                if pos > 0:
                    exchange.send_add_message_ioc(
                        next_id(), sym, Dir.SELL,
                        state.bid_prices.get(sym, 1), abs(pos)
                    )
                else:
                    exchange.send_add_message_ioc(
                        next_id(), sym, Dir.BUY,
                        state.ask_prices.get(sym, 99999), abs(pos)
                    )

    # ==================== LOOP ====================
    while True:
        msg = exchange.read_message()

        if msg["type"] == "close":
            break

        elif msg["type"] == "book":
            sym = msg["symbol"]

            bid = msg["buy"][0][0] if msg["buy"] else None
            ask = msg["sell"][0][0] if msg["sell"] else None

            state.update_bid_ask_price(sym, bid, ask)

            # 기존 주문 일부만 cancel (속도 유지)
            for oid in list(active_orders.keys())[:4]:
                exchange.send_cancel_message(oid)
                del active_orders[oid]

            # 전략 실행
            if sym in ["GS", "MS", "WFC"]:
                market_make(sym)

            xlf_arb()
            risk()

        elif msg["type"] == "fill":
            sym = msg["symbol"]
            qty = msg["size"]

            if msg["dir"] == Dir.BUY:
                positions[sym] += qty
            else:
                positions[sym] -= qty

            # 🔥 체결 즉시 재주문
            if sym in ["GS", "MS", "WFC"]:
                refill(sym)

# ====================
class ExchangeConnection:
    def __init__(self, args):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((args.exchange_hostname, args.port))
        self.reader = s.makefile("r", 1)
        self.writer = s
        self._write({"type": "hello", "team": team_name.upper()})

    def read_message(self):
        msg = json.loads(self.reader.readline())
        if "dir" in msg:
            msg["dir"] = Dir(msg["dir"])
        return msg

    def send_add_message(self, oid, sym, dir, price, size):
        self._write({"type":"add","order_id":oid,"symbol":sym,"dir":dir,"price":price,"size":size,"tif":"DAY"})

    def send_add_message_ioc(self, oid, sym, dir, price, size):
        self._write({"type":"add","order_id":oid,"symbol":sym,"dir":dir,"price":price,"size":size,"tif":"IOC"})

    def send_cancel_message(self, oid):
        self._write({"type": "cancel", "order_id": oid})

    def _write(self, msg):
        self.writer.send((json.dumps(msg)+"\n").encode())

# ====================
def parse_arguments():
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--production", action="store_true")
    group.add_argument("--test", type=str, default="prod-like")

    args = parser.parse_args()

    if args.production:
        args.exchange_hostname = "production"
        args.port = 25000
    else:
        args.exchange_hostname = "test-exch-" + team_name
        args.port = 22000

    return args

if __name__ == "__main__":
    main()