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

team_name = "HanyangFloorFunction"

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

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", default="prod-like")
    args = parser.parse_args()

    s = socket.socket()
    s.connect(("test-exch-"+team_name, 22000))
    r = s.makefile("r",1)

    def send(m):
        s.send((json.dumps(m)+"\n").encode())

    send({"type":"hello","team":team_name.upper()})

    st = StateManager()

    hello = json.loads(r.readline())

    pos = {}
    for sym in hello["symbols"]:
        pos[sym["symbol"]] = sym["position"]

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

    while True:
        m = json.loads(r.readline())

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

        if m["type"] == "book":
            sym = m["symbol"]

            if not m["buy"] or not m["sell"]:
                continue

            bid = m["buy"][0][0]
            ask = m["sell"][0][0]

            # 🔥 spread 충분할 때만
            if ask - bid >= 3:
                send({"type":"add","order_id":nid(),"symbol":sym,
                      "dir":"BUY","price":bid+1,"size":1})
                send({"type":"add","order_id":nid(),"symbol":sym,
                      "dir":"SELL","price":ask-1,"size":1})

        if m["type"] == "fill":
            sym = m["symbol"]
            q = m["size"]

            pos[sym] = pos.get(sym,0)

            if m["dir"] == "BUY":
                pos[sym] += q
            else:
                pos[sym] -= q

            print("POS:", pos)

if __name__ == "__main__":
    main()