#!/usr/bin/env python3
import argparse, socket, json, time
from enum import Enum
from state import StateManager

team_name = "HanyangFloorFunction"

# ===== 튜닝 파라미터 =====
ORDER_SIZE = 3
MAX_POS = 40
KILL_POS = 55
REFRESH = 0.25          # 주문 리프레시 주기
ARB_THRESHOLD = 12      # XLF 차익 임계
SKEW_K = 0.2            # 포지션 스큐 강도

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

# ===== Exchange =====
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(self):
        m = json.loads(self.reader.readline())
        if "dir" in m: m["dir"] = Dir(m["dir"])
        return m

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

    def ioc(self, oid, sym, d, px, sz):
        self.add(oid, sym, d, px, sz, "IOC")

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

    def convert(self, oid, sym, d, sz):
        self._write({"type":"convert","order_id":oid,"symbol":sym,"dir":d,"size":sz})

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

# ===== args =====
def parse_arguments():
    p = argparse.ArgumentParser()
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--production", action="store_true")
    g.add_argument("--test", type=str, default="prod-like")
    a = p.parse_args()
    if a.production:
        a.exchange_hostname, a.port = "production", 25000
    else:
        a.exchange_hostname, a.port = "test-exch-"+team_name, 22000
    return a

# ===== main =====
def main():
    args = parse_arguments()
    ex = ExchangeConnection(args)
    st = StateManager()

    hello = ex.read()
    pos = {s["symbol"]: s["position"] for s in hello["symbols"]}

    oid = 0
    def nid():
        nonlocal oid; oid += 1; return oid

    active = {}  # oid -> sym
    last_refresh = 0.0

    # ---- 유틸 ----
    def mid(sym):
        b, a = st.bid_prices.get(sym), st.ask_prices.get(sym)
        return None if (b is None or a is None) else (b+a)//2

    def spread(sym):
        b, a = st.bid_prices.get(sym), st.ask_prices.get(sym)
        return None if (b is None or a is None) else (a-b)

    # ---- MM (스큐 포함) ----
    def place_mm(sym):
        b, a = st.bid_prices.get(sym), st.ask_prices.get(sym)
        if b is None or a is None: return
        if spread(sym) is None or spread(sym) < 2: return

        p = pos.get(sym, 0)
        if abs(p) >= MAX_POS: return

        # 기본: bid+1 / ask-1
        buy_px  = b + 1
        sell_px = a - 1
        if buy_px >= sell_px: return

        # 포지션 스큐: 롱이면 매도 쪽 강화, 숏이면 매수 쪽 강화
        skew = int(SKEW_K * p)
        buy_sz  = max(1, ORDER_SIZE - max(0, skew))
        sell_sz = max(1, ORDER_SIZE + max(0, skew))
        if p < 0:
            buy_sz  = max(1, ORDER_SIZE + max(0, -skew))
            sell_sz = max(1, ORDER_SIZE - max(0, -skew))

        # 안전 클램프
        buy_sz  = min(buy_sz,  MAX_POS - max(p, 0))
        sell_sz = min(sell_sz, MAX_POS + min(p, 0))
        if buy_sz <= 0 and sell_sz <= 0: return

        if buy_sz > 0:
            o = nid(); ex.add(o, sym, Dir.BUY,  buy_px,  buy_sz);  active[o] = sym
        if sell_sz > 0:
            o = nid(); ex.add(o, sym, Dir.SELL, sell_px, sell_sz); active[o] = sym

    # ---- XLF 차익 (정석) ----
    def xlf_arb():
        bond_ask = st.ask_prices.get("BOND")
        gs_ask   = st.ask_prices.get("GS")
        ms_ask   = st.ask_prices.get("MS")
        wfc_ask  = st.ask_prices.get("WFC")
        xlf_bid  = st.bid_prices.get("XLF")

        bond_bid = st.bid_prices.get("BOND")
        gs_bid   = st.bid_prices.get("GS")
        ms_bid   = st.bid_prices.get("MS")
        wfc_bid  = st.bid_prices.get("WFC")
        xlf_ask  = st.ask_prices.get("XLF")

        if None in [bond_ask, gs_ask, ms_ask, wfc_ask, xlf_bid,
                    bond_bid, gs_bid, ms_bid, wfc_bid, xlf_ask]:
            return

        # 케이스1: 바스켓 매수 → XLF 변환 → XLF 매도
        cost = bond_ask*3 + gs_ask*2 + ms_ask*3 + wfc_ask*2
        p1 = xlf_bid*10 - cost
        if p1 > ARB_THRESHOLD and pos.get("XLF",0) < MAX_POS:
            ex.ioc(nid(),"BOND",Dir.BUY,bond_ask,3)
            ex.ioc(nid(),"GS",  Dir.BUY,gs_ask,  2)
            ex.ioc(nid(),"MS",  Dir.BUY,ms_ask,  3)
            ex.ioc(nid(),"WFC", Dir.BUY,wfc_ask, 2)
            ex.convert(nid(),"XLF",Dir.BUY,10)
            ex.ioc(nid(),"XLF",Dir.SELL,xlf_bid,10)

        # 케이스2: XLF 매수 → 바스켓 변환 → 바스켓 매도
        rev = bond_bid*3 + gs_bid*2 + ms_bid*3 + wfc_bid*2
        p2 = rev - xlf_ask*10
        if p2 > ARB_THRESHOLD and pos.get("XLF",0) > -MAX_POS:
            ex.ioc(nid(),"XLF",Dir.BUY,xlf_ask,10)
            ex.convert(nid(),"XLF",Dir.SELL,10)
            ex.ioc(nid(),"BOND",Dir.SELL,bond_bid,3)
            ex.ioc(nid(),"GS",  Dir.SELL,gs_bid,  2)
            ex.ioc(nid(),"MS",  Dir.SELL,ms_bid,  3)
            ex.ioc(nid(),"WFC", Dir.SELL,wfc_bid, 2)

    # ---- 리스크 컷 ----
    def risk():
        for s, p in pos.items():
            if abs(p) >= KILL_POS:
                b, a = st.bid_prices.get(s), st.ask_prices.get(s)
                if p > 0 and b:
                    ex.ioc(nid(), s, Dir.SELL, b, abs(p))
                elif p < 0 and a:
                    ex.ioc(nid(), s, Dir.BUY,  a, abs(p))

    # ---- 주문 리프레시 (전부 취소 후 재호가) ----
    def refresh_all():
        # 전부 취소
        for o in list(active.keys()):
            ex.cancel(o)
            active.pop(o, None)
        # 재호가
        for s in ["GS","MS","WFC"]:
            place_mm(s)

    # ===== LOOP =====
    while True:
        m = ex.read()

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

        elif m["type"] == "book":
            sym = m["symbol"]
            b = m["buy"][0][0]  if m["buy"]  else None
            a = m["sell"][0][0] if m["sell"] else None
            st.update_bid_ask_price(sym, b, a)

            # 주기적 리프레시
            now = time.time()
            if now - last_refresh > REFRESH:
                last_refresh = now
                refresh_all()

            # XLF 차익 + 리스크
            if sym in ["BOND","GS","MS","WFC","XLF"]:
                xlf_arb()
            risk()

        elif m["type"] == "fill":
            s, q, d = m["symbol"], m["size"], m["dir"]
            pos[s] = pos.get(s,0) + (q if d == Dir.BUY else -q)

            # 체결 즉시 해당 심볼만 재호가 (속도 ↑)
            if s in ["GS","MS","WFC"]:
                place_mm(s)

if __name__ == "__main__":
    main()