#!/usr/bin/env python3

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

team_name = "HanyangFloorFunction"

# ==================== 설정 ====================
ORDER_SIZE = 2
MAX_POS = 30

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

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

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

    # 🔥 포지션 직접 관리
    positions = {}
    for sym in hello["symbols"]:
        positions[sym["symbol"]] = sym["position"]

    # 🔥 order id
    order_id = 0
    def next_id():
        nonlocal order_id
        order_id += 1
        return order_id

    # 🔥 active order 관리 (핵심)
    active_orders = {}

    # ==================== MARKET MAKING ====================
    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 pos > MAX_POS:
            return
        if pos < -MAX_POS:
            return

        # 🔥 핵심: 스프레드 내부에서만 주문
        buy_price = bid + 1
        sell_price = ask - 1

        # 🔥 포지션 방향 제어
        if pos > 0:
            buy_size = 1
            sell_size = ORDER_SIZE + 2
        elif pos < 0:
            buy_size = ORDER_SIZE + 2
            sell_size = 1
        else:
            buy_size = sell_size = ORDER_SIZE

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

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

    # ==================== BOND 안전 수익 ====================
    def bond_arb():
        bid = state.bid_prices.get("BOND")
        ask = state.ask_prices.get("BOND")

        if ask and ask < 1000:
            exchange.send_add_message_ioc(next_id(), "BOND", Dir.BUY, ask, 5)

        if bid and bid > 1000:
            exchange.send_add_message_ioc(next_id(), "BOND", Dir.SELL, bid, 5)

    # ==================== 리스크 관리 ====================
    def risk():
        for sym in ["GS", "MS", "WFC"]:
            pos = positions.get(sym, 0)

            if abs(pos) > MAX_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)

            # 🔥 기존 주문 전부 취소 (핵심)
            for oid in list(active_orders.keys()):
                exchange.send_cancel_message(oid)
                del active_orders[oid]

            # 전략 실행
            bond_arb()

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

            risk()

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

            if msg["dir"] == Dir.BUY:
                positions[sym] = positions.get(sym, 0) + qty
            else:
                positions[sym] = positions.get(sym, 0) - qty

# ====================
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()
    parser.add_argument("--test", default="prod-like")
    args = parser.parse_args()

    args.exchange_hostname = "test-exch-" + team_name
    args.port = 22000
    return args

if __name__ == "__main__":
    main()