#!/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 = 30
KILL_POS = 40
REFRESH = 0.4
ARB_THRESHOLD = 15

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

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

    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, 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

# ======================
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 = {}
    last_refresh = 0

    # ===== Market Making =====
    def market_make(sym):
        bid = st.bid_prices.get(sym)
        ask = st.ask_prices.get(sym)

        if bid is None or ask is None:
            return

        if ask - bid < 2:
            return

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

        buy_px = bid + 1
        sell_px = ask - 1

        if buy_px >= sell_px:
            return

        buy_sz = ORDER_SIZE
        sell_sz = ORDER_SIZE

        # 포지션 조절
        if p > 0:
            sell_sz += 2
        elif p < 0:
            buy_sz += 2

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

    # ===== XLF Arb =====
    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
        cost = bond_ask*3 + gs_ask*2 + ms_ask*3 + wfc_ask*2
        profit1 = xlf_bid*10 - cost

        if profit1 > ARB_THRESHOLD:
            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
        revenue = bond_bid*3 + gs_bid*2 + ms_bid*3 + wfc_bid*2
        profit2 = revenue - xlf_ask*10

        if profit2 > ARB_THRESHOLD:
            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)

    # ===== Risk =====
    def risk():
        for sym, p in pos.items():
            if abs(p) > KILL_POS:
                bid = st.bid_prices.get(sym)
                ask = st.ask_prices.get(sym)

                if p > 0 and bid:
                    ex.ioc(nid(), sym, Dir.SELL, bid, abs(p))
                elif p < 0 and ask:
                    ex.ioc(nid(), sym, Dir.BUY, ask, abs(p))

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

        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

            st.update_bid_ask_price(sym, bid, ask)

            now = time.time()
            if now - last_refresh > REFRESH:
                last_refresh = now

                # 모든 주문 취소
                for o in list(active.keys()):
                    ex.cancel(o)
                    active.pop(o, None)

                for s in ["GS","MS","WFC"]:
                    market_make(s)

            xlf_arb()
            risk()

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

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

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

# ======================
if __name__ == "__main__":
    main()