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

team_name = "HanyangFloorFunction"

ORDER_SIZE = 2
MAX_POS = 20
KILL_POS = 30
REFRESH = 0.8   # 🔥 주문 오래 유지

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):
        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 cancel(self, oid):
        self._write({"type":"cancel","order_id":oid})

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

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

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="production"; a.port=25000
    else:
        a.exchange_hostname="test-exch-"+team_name; a.port=22000
    return a

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

    def market_make(sym):
        b=st.bid_prices.get(sym)
        a=st.ask_prices.get(sym)

        if b is None or a is None:
            return

        spread=a-b

        # 🔥 핵심: 넓을 때만 거래
        if spread < 3:
            return

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

        buy_px=b+1
        sell_px=a-1

        if buy_px>=sell_px:
            return

        buy_sz=ORDER_SIZE
        sell_sz=ORDER_SIZE

        # 포지션 줄이기
        if p>0:
            sell_sz+=1
        elif p<0:
            buy_sz+=1

        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

    def risk():
        for s,p in pos.items():
            if abs(p)>KILL_POS:
                b=st.bid_prices.get(s)
                a=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))

    while True:
        m=ex.read()

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

        elif m["type"]=="book":
            s=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(s,b,a)

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

                # cancel 줄임
                for o in list(active.keys()):
                    ex.cancel(o)
                    active.pop(o,None)

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

            risk()

        elif m["type"]=="fill":
            s=m["symbol"]
            q=m["size"]

            if m["dir"]==Dir.BUY:
                pos[s]+=q
            else:
                pos[s]-=q

if __name__=="__main__":
    main()